"""Run the agentic controller over a query file -> run.tsv (final ranked docs), for nDCG scoring.
Two phases: (1) batch-plan all queries in parallel (DeepSeek, cached per qid so ablation runs reuse
the SAME plans), (2) assemble each query (parallel fetch + rerank + rank) sequentially.

Env: THEMIS_DATA, THEMIS_STATUTE, THEMIS_QFILE, OUT, PLANCACHE, THEMIS_ENABLED=all|vector,authority,...
"""
import os, sys, json, time
import concurrent.futures as cf
import requests
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "scripts"))
from tools import Corpus
import agent as A

HERE = os.path.dirname(os.path.abspath(__file__))
def _load_env(p):
    for l in open(p):
        l = l.strip()
        if l and not l.startswith("#") and "=" in l:
            k, v = l.split("=", 1); os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))
_load_env(os.path.join(HERE, "..", "scripts", ".env"))
HDR = {"Authorization": f"Bearer {os.environ['DEEPSEEK_API_KEY']}", "Content-Type": "application/json"}

def llm_fn(msgs):
    for _ in range(3):
        try:
            r = requests.post("https://api.deepseek.com/chat/completions", headers=HDR, timeout=60,
                json={"model": "deepseek-chat", "temperature": 0, "max_tokens": 400, "messages": msgs})
            if r.status_code == 200: return r.json()["choices"][0]["message"]["content"]
        except Exception: time.sleep(2)
    return "{}"

QFILE = os.environ.get("THEMIS_QFILE", "authority_queries.tsv")
OUT = os.environ.get("OUT", "agent_run.tsv")
PLANCACHE = os.environ.get("PLANCACHE", "plan_" + os.path.basename(QFILE).split(".")[0] + ".json")
ENABLED = A.ALL_TOOLS if os.environ.get("THEMIS_ENABLED", "all") == "all" else set(os.environ["THEMIS_ENABLED"].split(","))
ALPHA = float(os.environ.get("ALPHA", "0.3"))

rows = []
for l in open(QFILE, encoding="utf-8"):
    qid, it, t = l.rstrip("\n").split("\t", 2); rows.append((qid, t))

# Phase 1: batch-plan (parallel, cached)
plans = json.load(open(PLANCACHE)) if os.path.exists(PLANCACHE) else {}
todo = [(qid, t) for qid, t in rows if qid not in plans]
if todo:
    print(f"planning {len(todo)} queries ...", flush=True); t0 = time.time()
    with cf.ThreadPoolExecutor(max_workers=24) as ex:
        futs = {ex.submit(A.plan, t, llm_fn): qid for qid, t in todo}
        done = 0
        for f in cf.as_completed(futs):
            plans[futs[f]] = f.result(); done += 1
            if done % 50 == 0: json.dump(plans, open(PLANCACHE, "w")); print(f"  {done}/{len(todo)} {time.time()-t0:.0f}s", flush=True)
    json.dump(plans, open(PLANCACHE, "w"))
    print(f"plans done {time.time()-t0:.0f}s", flush=True)

# Phase 2: assemble (sequential; CE is CPU-bound)
C = Corpus(os.environ.get("THEMIS_DATA", "."), os.environ.get("THEMIS_STATUTE", "."), device=os.environ.get("THEMIS_DEVICE", "cpu"))
print(f"assembling {len(rows)} queries (enabled={sorted(ENABLED)}) ...", flush=True)
t0 = time.time()
with open(OUT, "w", encoding="utf-8") as f:
    for i, (qid, t) in enumerate(rows):
        ranked, info = A.assemble(C, t, plans[qid], enabled=ENABLED, alpha=ALPHA)
        for rank, d in enumerate(ranked, 1): f.write(f"{qid}\t{rank}\t{d}\n")
        if (i + 1) % 30 == 0: print(f"  {i+1}/{len(rows)} {time.time()-t0:.0f}s", flush=True)
print(f"DONE {len(rows)} -> {OUT} in {time.time()-t0:.0f}s", flush=True)