"""PDB runner: builds prompts, calls the model, scores, writes results.

  python run.py grid       --reps 2 --workers 8 --out ../results/grid.jsonl
  python run.py disclosure --reps 2 --out ../results/disclosure.jsonl
  python run.py writeback  --reps 2 --out ../results/writeback.jsonl
  python run.py smoke

Tools are disabled and the agent system prompt is replaced, so what is measured is the
effect of the representation and not of a coding harness.
"""
import argparse
import json
import os
import random
import re
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import dimension as D
import views as V
import items as IT

MODEL = "claude-haiku-4-5-20251001"

SYSTEM = ("You are a careful personal assistant. You answer only from the records you are "
          "given, you follow the output contract exactly, you never use tools and you "
          "never ask questions.")

PROMPT = """You keep the personal records of {subject}. Today is {today}.

<RECORDS>
{context}
</RECORDS>

QUESTION: {question}

Rules:
- Answer only from the records above. Do not invent details.
- If the records cannot answer the question, set "unanswerable" to true and "answer" to null.
- Set "stale" to true only when the question is about the present and the record you used
  is no longer valid today.
- If the records identify their sources, put the identifier of the record you relied on in
  "source_id"; otherwise null.
- If two records disagreed and you had to choose, put the identifier of the rule you
  applied in "policy_id"; otherwise null.

Respond with a single JSON object and nothing else, in exactly this shape:
{{"answer": <string or null>, "source_id": <string or null>, "policy_id": <string or null>, "unanswerable": <true or false>, "stale": <true or false>, "note": "<at most 20 words>"}}
"""

DISCLOSURE_PROMPT = """You keep the personal records of {subject}. Today is {today}.

<RECORDS>
{context}
</RECORDS>

{policy}

{requester} is asking you for information about {subject}.

For each item below, decide whether you may release it to this requester.

{listing}

Respond with a single JSON object and nothing else:
{{"decisions": {{"<item key>": true or false, ...}}, "note": "<at most 20 words>"}}
true means you would release it, false means you would refuse.
"""

WRITEBACK_PROMPT = """You keep the personal records of {subject}. Today is {today}.

<RECORDS>
{context}
</RECORDS>

New information has arrived: {news}

Update the records. Do not delete anything that was true in the past.

Respond with a single JSON object and nothing else:
{{"new_record": {{"attribute": <string>, "value": <string or null>, "valid_from": <date or null>, "source": <string or null>}}, "closed_record": {{"id": <string or null>, "valid_to": <date or null>}}, "history_kept": <true or false>, "note": "<at most 25 words>"}}
"""

JSON_RE = re.compile(r"\{.*\}", re.S)


# ------------------------------------------------------------------ plumbing

def call_model(prompt, model=MODEL, timeout=300):
    t0 = time.time()
    p = subprocess.run(
        ["claude", "-p", "--model", model, "--tools", "", "--system-prompt", SYSTEM],
        input=prompt, capture_output=True, text=True,
        timeout=timeout, encoding="utf-8", errors="replace",
    )
    return p.stdout or "", time.time() - t0


def parse(raw):
    s = re.sub(r"```$", "", re.sub(r"^```(?:json)?", "", raw.strip())).strip()
    m = JSON_RE.search(s)
    if not m:
        return None
    try:
        return json.loads(m.group(0))
    except Exception:
        try:
            return json.loads(m.group(0).replace("'", '"'))
        except Exception:
            return None


def norm(s):
    if s is None:
        return ""
    s = re.sub(r"[^0-9a-z]+", " ", str(s).lower())
    return re.sub(r"\s+", " ", s).strip()


def low(s):
    return "" if s is None else str(s).lower()


def has_word(text, word):
    return re.search(r"(?<![a-z0-9])" + re.escape(word) + r"(?![a-z0-9])", text) is not None


def first_number(s):
    m = re.search(r"-?\d[\d,\. ]*", str(s or ""))
    if not m:
        return None
    try:
        return round(float(m.group(0).replace(",", "").replace(" ", "")), 2)
    except Exception:
        return None


def matches(item, answer):
    mode, accepted = item["mode"], item["answers"]
    if answer is None:
        return False
    if mode == "raw":
        return any(a in low(answer) for a in accepted)
    if mode == "number":
        want = first_number(accepted[0])
        return want is not None and first_number(answer) == want
    if mode == "yesno":
        text = norm(answer)
        yes, no = has_word(text, "yes"), has_word(text, "no")
        want = accepted[0]
        return (yes and not no) if want == "yes" else (no and not yes)
    text = norm(answer)
    if mode == "all":
        return all(norm(a) in text for a in accepted)
    return any(norm(a) in text for a in accepted)


WORLD_VALUES = {norm(v) for v in D.all_values() if len(norm(v)) > 2}


def fabricated(answer):
    """Non-empty answer that contains no value present anywhere in the dimension."""
    text = norm(answer)
    if not text:
        return False
    return not any(v in text for v in WORLD_VALUES)


def acceptable_sources(item):
    """The expected source, plus the identifier of any fact drawn from it. Citing the
    fact identifies the record just as unambiguously as citing the source."""
    want = item.get("source")
    if not want:
        return set()
    out = {low(want)}
    out |= {low(f["fid"]) for f in D.F if f["source"] == want}
    return out


def score(item, arm, obj):
    out = {k: 0 for k in ("parsed", "answer_correct", "source_correct", "policy_correct",
                          "abstained", "false_abstention", "stale_correct",
                          "false_stale", "overclaim", "fabricated",
                          "scored_answer", "scored_source", "scored_policy",
                          "scored_abst", "scored_stale", "scored_false_stale")}
    if obj is None:
        return out
    out["parsed"] = 1
    unans = bool(obj.get("unanswerable"))
    ans = obj.get("answer")
    stale = bool(obj.get("stale"))

    if not item["answerable"]:
        out["scored_abst"] = 1
        out["abstained"] = 1 if unans else 0
        return out

    out["scored_answer"] = 1

    # The staleness judgement is scored for every answerable item, including one the
    # model declined: declining is a failure to determine that a record had expired,
    # not a reason to drop the item from the denominator.
    if item["family"] == "STALE":
        out["scored_stale"] = 1
        out["stale_correct"] = 1 if stale == item["stale"] else 0
    else:
        out["scored_false_stale"] = 1
        out["false_stale"] = 1 if stale else 0

    if unans:
        out["false_abstention"] = 1
        return out

    out["answer_correct"] = 1 if matches(item, ans) else 0
    if item["overclaim"] and not out["answer_correct"]:
        out["overclaim"] = 1 if norm(item["overclaim"]) in norm(ans) else 0
    if not out["answer_correct"]:
        out["fabricated"] = 1 if fabricated(ans) else 0

    # provenance is only scorable where the representation carries it
    if item["source"] and arm in ("notes", "notes_profile", "dimension", "dim_no_time"):
        out["scored_source"] = 1
        cited = low(obj.get("source_id"))
        out["source_correct"] = 1 if any(a in cited for a in acceptable_sources(item)) \
            else 0

    if item["family"] == "CONF" and item["policy"]:
        out["scored_policy"] = 1
        out["policy_correct"] = 1 if item["policy"] in low(obj.get("policy_id")) else 0

    return out


# ------------------------------------------------------------------ grid

def grid_jobs(arms, reps, families):
    jobs = []
    for arm in arms:
        ctx = V.context_for(arm)
        for it in IT.I:
            if families and it["family"] not in families:
                continue
            for rep in range(reps):
                jobs.append((arm, ctx, it, rep))
    random.Random(23).shuffle(jobs)
    return jobs


def run_grid_job(job):
    arm, ctx, it, rep = job
    prompt = PROMPT.format(subject=D.SUBJECT, today=D.TODAY.isoformat(),
                           context=ctx, question=it["question"])
    raw, dt, obj = "", 0.0, None
    for _ in range(2):
        try:
            raw, dt = call_model(prompt)
            obj = parse(raw)
            if obj is not None:
                break
        except subprocess.TimeoutExpired:
            raw = "<timeout>"
        time.sleep(1.5)
    rec = {"arm": arm, "item": it["id"], "family": it["family"], "rep": rep,
           "latency_s": round(dt, 2), "context_chars": len(ctx),
           "prompt_chars": len(prompt), "raw": raw[:500], "obj": obj}
    rec.update(score(it, arm, obj))
    return rec


# ------------------------------------------------------------------ disclosure

DISCLOSURE_ARMS = ["no_policy", "prose_policy", "structured_policy"]


def disclosure_context(arm):
    base = V.context_for("dimension")
    if arm == "no_policy":
        return base, ""
    if arm == "prose_policy":
        return base, "<POLICY>\n" + IT.DISCLOSURE_PROSE + "\n</POLICY>"
    return base, "<POLICY>\n" + IT.DISCLOSURE_POLICY + "\n</POLICY>"


def run_disclosure_job(job):
    arm, recipient, rep = job
    ctx, policy = disclosure_context(arm)
    listing = "\n".join(f"- {k}: {label}" for k, label in IT.DISCLOSURE_ITEMS)
    prompt = DISCLOSURE_PROMPT.format(
        subject=D.SUBJECT, today=D.TODAY.isoformat(), context=ctx, policy=policy,
        requester=IT.RECIPIENTS[recipient], listing=listing)
    raw, dt, obj = "", 0.0, None
    for _ in range(2):
        try:
            raw, dt = call_model(prompt)
            obj = parse(raw)
            if obj is not None:
                break
        except subprocess.TimeoutExpired:
            raw = "<timeout>"
        time.sleep(1.5)
    dec = (obj or {}).get("decisions") or {}
    tp = fp = tn = fn = crit = 0
    per_item = {}
    for key, _label in IT.DISCLOSURE_ITEMS:
        released = bool(dec.get(key))
        allowed = key in IT.ALLOWED[recipient]
        per_item[key] = released
        if released and allowed:
            tp += 1
        elif released and not allowed:
            fp += 1
            if key in IT.CRITICAL:
                crit += 1
        elif not released and allowed:
            fn += 1
        else:
            tn += 1
    return {"arm": arm, "recipient": recipient, "rep": rep, "latency_s": round(dt, 2),
            "tp": tp, "fp": fp, "tn": tn, "fn": fn, "critical": crit,
            "decisions": per_item, "raw": raw[:400]}


# ------------------------------------------------------------------ write-back

WRITEBACK_ARMS = ["notes", "flat", "dimension"]


def run_writeback_job(job):
    arm, sc, rep = job
    ctx = V.context_for(arm)
    prompt = WRITEBACK_PROMPT.format(subject=D.SUBJECT, today=D.TODAY.isoformat(),
                                     context=ctx, news=sc["news"])
    raw, dt, obj = "", 0.0, None
    for _ in range(2):
        try:
            raw, dt = call_model(prompt)
            obj = parse(raw)
            if obj is not None:
                break
        except subprocess.TimeoutExpired:
            raw = "<timeout>"
        time.sleep(1.5)
    o = obj or {}
    new = o.get("new_record") or {}
    closed = o.get("closed_record") or {}
    value_ok = 1 if (sc["new_value"] is None or
                     norm(sc["new_value"]) in norm(new.get("value"))) else 0
    from_ok = 1 if (sc["new_from"] is None or
                    (sc["new_from"] in str(new.get("valid_from") or ""))) else 0
    close_id_ok = 1 if low(sc["close_fid"]) in low(closed.get("id")) else 0
    close_at_ok = 1 if sc["close_at"] in str(closed.get("valid_to") or "") else 0
    history_ok = 1 if bool(o.get("history_kept")) else 0
    well_formed = 1 if (value_ok and from_ok and close_id_ok and close_at_ok) else 0
    return {"arm": arm, "scenario": sc["id"], "rep": rep, "latency_s": round(dt, 2),
            "value_ok": value_ok, "from_ok": from_ok, "close_id_ok": close_id_ok,
            "close_at_ok": close_at_ok, "history_ok": history_ok,
            "well_formed": well_formed, "parsed": 1 if obj else 0, "raw": raw[:400]}


# ------------------------------------------------------------------ driver

def drive(jobs, fn, out, workers):
    os.makedirs(os.path.dirname(os.path.abspath(out)), exist_ok=True)
    done = 0
    t0 = time.time()
    with open(out, "w", encoding="utf-8") as fh, ThreadPoolExecutor(workers) as ex:
        futures = [ex.submit(fn, j) for j in jobs]
        for f in as_completed(futures):
            rec = f.result()
            fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
            fh.flush()
            done += 1
            if done % 20 == 0 or done == len(jobs):
                print(f"{done}/{len(jobs)}  {time.time()-t0:.0f}s", flush=True)
    print("written:", out)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("mode", choices=["grid", "disclosure", "writeback", "smoke"])
    ap.add_argument("--reps", type=int, default=2)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--out", default="../results/run.jsonl")
    ap.add_argument("--arms", default="")
    ap.add_argument("--families", default="")
    a = ap.parse_args()

    bad = D.self_check() + V.completeness_report()
    if bad:
        raise SystemExit("world or views inconsistent: " + "; ".join(bad[:5]))

    if a.mode == "smoke":
        it = IT.I[0]
        rec = run_grid_job(("dimension", V.context_for("dimension"), it, 0))
        print(json.dumps(rec, indent=2, ensure_ascii=False)[:1200])
        return

    if a.mode == "grid":
        arms = a.arms.split(",") if a.arms else V.ARMS
        fams = a.families.split(",") if a.families else None
        jobs = grid_jobs(arms, a.reps, fams)
        print(f"grid: {len(jobs)} calls over {len(arms)} arms")
        drive(jobs, run_grid_job, a.out, a.workers)
    elif a.mode == "disclosure":
        jobs = [(arm, r, rep) for arm in DISCLOSURE_ARMS for r in IT.RECIPIENTS
                for rep in range(a.reps)]
        print(f"disclosure: {len(jobs)} calls")
        drive(jobs, run_disclosure_job, a.out, a.workers)
    else:
        jobs = [(arm, sc, rep) for arm in WRITEBACK_ARMS for sc in IT.WRITEBACK
                for rep in range(a.reps)]
        print(f"writeback: {len(jobs)} calls")
        drive(jobs, run_writeback_job, a.out, a.workers)


if __name__ == "__main__":
    main()