"""Shaktimaan runtime: route a task with the 0.5B router, then call the chosen agent(s) via OpenAI-compatible APIs.

Usage:
  SHAKTIMAAN_ENDPOINTS='{"glm-5.3-flash": "http://host:8000/v1", ...}' \
      python orchestrate.py [--image chart.png] [--code] "task"
"""
import argparse, asyncio, base64, json, mimetypes, os, sys
import httpx
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from common import (CODE_SYSTEM, POOL, ROUTER_SYSTEM, SOLVER_SYSTEM, STRONGEST, capable, endpoints, est_tokens,
                    max_out, router_view, strip_think)

ROUTER = os.environ.get("SHAKTIMAAN_ROUTER", "podstack/Shaktimaan")
ENDPOINTS = endpoints()


def route(task, has_image):
    tok = AutoTokenizer.from_pretrained(ROUTER)
    model = AutoModelForCausalLM.from_pretrained(ROUTER, dtype=torch.bfloat16)
    enc = tok.apply_chat_template([{"role": "system", "content": ROUTER_SYSTEM},
                                   {"role": "user", "content": router_view(task, has_image)}],
                                  add_generation_prompt=True, return_tensors="pt", return_dict=True)
    out = model.generate(**enc, max_new_tokens=96, do_sample=False)
    try:
        plan = json.loads(tok.decode(out[0, enc["input_ids"].shape[1]:], skip_special_tokens=True))
    except Exception:  # noqa: BLE001
        plan = {"strategy": "single", "agents": []}
    # keep only agents that are reachable and can actually take this input
    ok = [a for a in plan.get("agents", []) if a in ENDPOINTS and capable(a, task, has_image)]
    plan["agents"] = ok or [a for a in [STRONGEST, "kimi-vl"] if a in ENDPOINTS and capable(a, task, has_image)][:1]
    return plan


async def ask(client, agent, messages, n_tokens):
    r = await client.post(f"{ENDPOINTS[agent]}/chat/completions",
                          json={"model": POOL[agent]["model"], "messages": messages,
                                "max_tokens": max_out(agent, n_tokens)})
    r.raise_for_status()
    return strip_think(r.json()["choices"][0]["message"]["content"])


async def run(task, image=None, code=False):
    has_image = image is not None
    plan = route(task, has_image)
    print("route:", json.dumps(plan), file=sys.stderr)
    system = CODE_SYSTEM if code else SOLVER_SYSTEM
    user = [{"type": "image_url", "image_url": {"url": image}}, {"type": "text", "text": task}] if has_image else task
    msgs = [{"role": "system", "content": system}, {"role": "user", "content": user}]
    n = est_tokens(task) + 1500 * has_image
    async with httpx.AsyncClient(timeout=1800) as client:
        if plan.get("strategy") != "ensemble" or len(plan["agents"]) == 1:
            return await ask(client, plan["agents"][0], msgs, n)
        drafts = await asyncio.gather(*[ask(client, a, msgs, n) for a in plan["agents"]])
        joined = "\n\n".join(f"### Candidate {i + 1}\n{d}" for i, d in enumerate(drafts))
        agg = f"{task}\n\nIndependent candidate solutions:\n{joined}\n\n" \
              "Check them, resolve disagreements, and give the best final answer."
        return await ask(client, plan["agents"][-1], [{"role": "system", "content": system},
                                                      {"role": "user", "content": agg}], est_tokens(agg))


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("task", nargs="+")
    p.add_argument("--image", help="path to an image to attach")
    p.add_argument("--code", action="store_true", help="ask agents for a Python code solution")
    args = p.parse_args()
    uri = None
    if args.image:
        mime = mimetypes.guess_type(args.image)[0] or "image/png"
        uri = f"data:{mime};base64," + base64.b64encode(open(args.image, "rb").read()).decode()
    print(asyncio.run(run(" ".join(args.task), uri, args.code)))