"""LLM-based dual-arm task allocation (the 'T' in TAMP: decide WHICH arm does WHAT).

Given the scene (object world positions + the two YAM arm bases) and a goal, ask an LLM
to output a per-arm plan: which arm grasps which object, in what order, and whether a
hand-off is needed. Reachability heuristic + collision avoidance are described to the model.
The motion for each assigned sub-task is then executed by the existing PRM + diff-IK layer.

Uses HuggingFace router (OpenAI-compatible). Key via env HF_TOKEN. No key is printed.
"""
import os, json, urllib.request

HF = os.environ["HF_TOKEN"]
MODEL = os.environ.get("VLM_MODEL", "deepseek-ai/DeepSeek-V4-Flash")

# ---- real scene (env-local world frame; z up; table top ~0.45) ----
SCENE = {
    "arms": {
        "left_robot":  {"base": [-0.2,  0.2, 0.45], "reach_m": 0.55},
        "right_robot": {"base": [-0.2, -0.2, 0.45], "reach_m": 0.55},
    },
    "objects": {
        "apple": {"pos": [0.00,  0.00, 0.47], "graspable": "sphere ~4cm"},
        "grape": {"pos": [0.00,  0.10, 0.44], "graspable": "cluster ~15x8cm"},
        "can":   {"pos": [0.00, -0.12, 0.50], "graspable": "cylinder d4.4 h11cm"},
    },
    "placement_targets": {"cart": [-0.5, 0.0, 0.45]},
    "notes": "Both arms are fixed-base. An arm reaches an object if |object - base| < reach_m. "
             "Prefer assigning each object to the CLOSER arm; the two arms must not be sent to "
             "the same object at the same time (collision). Grasp is top-down.",
}
GOAL = os.environ.get("GOAL", "Pick up all three objects and place them into the cart, using both arms in parallel where possible.")

SYS = (
    "You are the task-allocation planner for a BIMANUAL robot (two 6-DoF arms, left_robot and "
    "right_robot, each with a parallel-jaw gripper). Decide which arm performs which sub-task. "
    "Respect reachability (|obj-base|<reach) and avoid sending both arms to the same object at once. "
    "Return STRICT JSON only, no prose, with schema:\n"
    "{\"assignments\":[{\"arm\":\"left_robot|right_robot\",\"object\":\"...\",\"action\":\"pick_and_place\","
    "\"target\":\"cart\",\"order\":1,\"reason\":\"...\"}],\"parallelizable\":[[...],[...]],"
    "\"handoffs\":[],\"notes\":\"...\"}"
)
USER = f"SCENE:\n{json.dumps(SCENE, indent=2)}\n\nGOAL: {GOAL}\n\nReturn the JSON plan."

body = json.dumps({
    "model": MODEL,
    "messages": [{"role": "system", "content": SYS}, {"role": "user", "content": USER}],
    "temperature": 0.2, "max_tokens": 900,
}).encode()

req = urllib.request.Request("https://router.huggingface.co/v1/chat/completions", data=body,
    headers={"Authorization": f"Bearer {HF}", "Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=90) as r:
    out = json.load(r)
msg = out["choices"][0]["message"]["content"]
print("=== RAW LLM OUTPUT ===")
print(msg)
# try to extract JSON
try:
    s = msg[msg.index("{"): msg.rindex("}")+1]
    plan = json.loads(s)
    print("\n=== PARSED PLAN ===")
    for a in plan.get("assignments", []):
        print(f"  [{a.get('order')}] {a.get('arm')}: {a.get('action')} {a.get('object')} -> {a.get('target')}  ({a.get('reason','')[:70]})")
    print("  parallelizable:", plan.get("parallelizable"))
    json.dump(plan, open("outputs/dualarm_allocation.json", "w"), indent=2)
    print("\nsaved -> outputs/dualarm_allocation.json")
except Exception as e:
    print("\n[parse warn]", e)