"""Workstream-B PILOT: DeepSeek structured summaries for ~5k judgments (all rig gold docs + random
fill), to be embedded as issues/holding/facts vectors and ablated on the frozen rig BEFORE any
corpus-wide spend. Output: summaries JSONL (resumable; safe to re-run).

Summary contract: retrieval keys ONLY — never shown as evidence. Anchored on HELD where present.
"""
import os, sys, json, time, random
import concurrent.futures as cf
import requests

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.environ.get("THEMIS_DATA", ".")
OUT = os.environ.get("OUT", os.path.join(HERE, "summaries_pilot.jsonl"))
N = int(os.environ.get("N", "5000"))

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"}

SYS = ('You summarise an Indian Supreme Court judgment into structured retrieval keys. Output ONLY JSON: '
  '{"issues": [2-4 short phrases, the distinct legal issues decided], '
  '"holding": 2-3 sentences — the ratio decidendi, what this case DECIDES (anchor on the HELD headnote if provided; '
  'plain modern legal English), '
  '"facts": 2-3 sentences — the fact pattern in plain words (who did what; the dispute), '
  '"statutes": [provisions central to the decision, e.g. "IPC 302", "Article 21"], '
  '"outcome": one word/phrase (allowed/dismissed/quashed/remanded/reference answered)}. '
  'Never invent; if the text is unclear on a field, keep it minimal.')

print("loading corpus ...", flush=True)
meta = {}
for l in open(os.path.join(DATA, "escr_meta.jsonl"), encoding="utf-8"):
    m = json.loads(l); meta[m["doc_id"]] = m
doc_chunks = {}
texts = []
with open(os.path.join(DATA, "escr_chunks.jsonl"), encoding="utf-8") as f:
    for i, l in enumerate(f):
        c = json.loads(l); texts.append(c["text"]); doc_chunks.setdefault(c["doc_id"], []).append(i)

# selection: every gold doc the rig knows + random fill to N
want = set()
for fn in ("qrels.tsv", "authority_qrels.tsv", "recall_recovery_qrels.tsv"):
    p = os.path.join(HERE, fn)
    if os.path.exists(p):
        for l in open(p):
            want.add(l.split("\t")[1])
want = {d for d in want if d in meta}
rng = random.Random(7)
rest = [d for d in meta if d not in want]
rng.shuffle(rest)
docs = list(want) + rest[:max(0, N - len(want))]
print(f"selected {len(docs)} docs ({len(want)} rig-gold + fill)", flush=True)

done = set()
if os.path.exists(OUT):
    for l in open(OUT):
        try: done.add(json.loads(l)["doc_id"])
        except Exception: pass
todo = [d for d in docs if d not in done]
print(f"{len(todo)} to summarise ({len(done)} already done)", flush=True)

def doc_text(d, cap=48000):
    cis = doc_chunks.get(d, [])
    full = "\n".join(texts[i] for i in cis)
    if len(full) <= cap: return full
    return full[:int(cap * 0.75)] + "\n[...]\n" + full[-int(cap * 0.2):]

def one(d):
    m = meta[d]
    held = (m.get("held") or "")[:5000]
    body = doc_text(d)
    user = (f"CASE: {m.get('case_name')} ({m.get('year') or m.get('date')})\n"
            + (f"HELD (reporter headnote): {held}\n\n" if held.strip() else "")
            + f"JUDGMENT TEXT:\n{body}\n\nJSON:")
    for attempt in range(3):
        try:
            r = requests.post("https://api.deepseek.com/chat/completions", headers=HDR, timeout=90,
                json={"model": "deepseek-chat", "temperature": 0, "max_tokens": 500,
                      "messages": [{"role": "system", "content": SYS}, {"role": "user", "content": user}]})
            if r.status_code == 200:
                t = r.json()["choices"][0]["message"]["content"]
                j = json.loads(t[t.find("{"):t.rfind("}") + 1])
                return {"doc_id": d, **{k: j.get(k) for k in ("issues", "holding", "facts", "statutes", "outcome")}}
        except Exception:
            time.sleep(2 * (attempt + 1))
    return None

t0 = time.time(); n_ok = 0
with cf.ThreadPoolExecutor(max_workers=24) as ex, open(OUT, "a", encoding="utf-8") as fh:
    for res in ex.map(one, todo):
        if res:
            fh.write(json.dumps(res, ensure_ascii=False) + "\n"); n_ok += 1
            if n_ok % 200 == 0:
                fh.flush(); rate = n_ok / (time.time() - t0)
                print(f"  {n_ok}/{len(todo)} ({rate:.1f}/s, eta {int((len(todo)-n_ok)/max(rate,0.1)/60)}min)", flush=True)
print(f"DONE {n_ok}/{len(todo)} in {(time.time()-t0)/60:.0f}min -> {OUT}", flush=True)