"""Round-0 intent classifier (the 'router' as a state-setter). Labels each query AUTHORITY
(wants the leading/landmark case on a doctrine -> apply the authority prior) vs SPECIFIC (a
particular case / fact-pattern / narrow holding -> no prior). DeepSeek, parallel, cached to JSON.

Usage: THEMIS_QFILE=authority_queries.tsv OUT=intent_authority.json python classify_intent.py
Reads .env for DEEPSEEK_API_KEY. Cached by qid so re-runs are free; delete the OUT file to refresh.
"""
import os, json, sys, time, concurrent.futures as cf
import requests

HERE = os.path.dirname(os.path.abspath(__file__))
def _load_env(p):
    if os.path.exists(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"))
KEY = os.environ["DEEPSEEK_API_KEY"]
HDR = {"Authorization": f"Bearer {KEY}", "Content-Type": "application/json"}
QFILE = os.environ.get("THEMIS_QFILE", "authority_queries.tsv")
OUT = os.environ.get("OUT", "intent.json")

SYS = ("You route a legal search query for an Indian Supreme Court case-law engine. "
       "Decide what the user is after:\n"
       "AUTHORITY = they want the leading / landmark / controlling case(s) on a legal PRINCIPLE, "
       "doctrine, right, or test (e.g. 'is privacy a fundamental right', 'test for sedition', "
       "'doctrine of basic structure').\n"
       "SPECIFIC = they want a particular named case, a narrow fact-pattern match, a specific "
       "statutory provision's application, or a procedural/factual lookup where the single most "
       "authoritative landmark is NOT necessarily the right answer.\n"
       "Reply with EXACTLY one word: AUTHORITY or SPECIFIC.")

def classify(text):
    for attempt in range(3):
        try:
            r = requests.post("https://api.deepseek.com/chat/completions", headers=HDR, timeout=40,
                json={"model": "deepseek-chat", "temperature": 0, "max_tokens": 4,
                      "messages": [{"role": "system", "content": SYS}, {"role": "user", "content": text}]})
            if r.status_code == 200:
                t = r.json()["choices"][0]["message"]["content"].strip().upper()
                return "AUTHORITY" if "AUTHORITY" in t else "SPECIFIC"
        except Exception:
            time.sleep(2 * (attempt + 1))
    return "SPECIFIC"  # fail-safe: no prior

def main():
    rows = []
    for l in open(QFILE, encoding="utf-8"):
        qid, intent, text = l.rstrip("\n").split("\t", 2); rows.append((qid, text))
    cache = json.load(open(OUT)) if os.path.exists(OUT) else {}
    todo = [(qid, text) for qid, text in rows if qid not in cache]
    print(f"{len(rows)} queries, {len(todo)} to classify ({len(cache)} cached)", flush=True)
    t0 = time.time()
    with cf.ThreadPoolExecutor(max_workers=24) as ex:
        futs = {ex.submit(classify, text): qid for qid, text in todo}
        done = 0
        for f in cf.as_completed(futs):
            cache[futs[f]] = f.result(); done += 1
            if done % 50 == 0:
                json.dump(cache, open(OUT, "w")); print(f"  {done}/{len(todo)} {time.time()-t0:.0f}s", flush=True)
    json.dump(cache, open(OUT, "w"))
    n_auth = sum(1 for v in cache.values() if v == "AUTHORITY")
    print(f"done {len(cache)} -> {OUT} | AUTHORITY={n_auth} ({100*n_auth/len(cache):.0f}%) SPECIFIC={len(cache)-n_auth}", flush=True)

if __name__ == "__main__":
    main()