"""CMB runner.

  python run.py grid   --reps 2 --workers 8 --out ../results/grid.jsonl
  python run.py scope  --reps 2 --out ../results/scope.jsonl
  python run.py restate --reps 2 --out ../results/restate.jsonl
  python run.py smoke

Tools disabled and the agent system prompt replaced, so what is measured is the shape of
the record the agent reads and not 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 org as O
import views as V
import items as IT

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

SYSTEM = ("You are a careful assistant working from an organisation's shared memory. 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 read the shared memory of {org}, which six teams write into: sales,
finance, support, legal, engineering and people. Today is {today}.

<MEMORY>
{context}
</MEMORY>

QUESTION: {question}

Rules:
- Answer only from the memory above. Do not invent details.
- If the memory cannot answer the question, set "unanswerable" to true and "answer" to null.
- Put the team accountable for the concept in "owner_team" if the memory identifies one,
  otherwise null.
- Put the identifier of the record or document you relied on in "source_id" if the memory
  carries identifiers, otherwise null.
- If two teams' records disagreed and you had to choose, put the identifier of the rule you
  applied in "policy_id", otherwise null.
- Set "stale" to true when the record the question asks about has been superseded.

Respond with a single JSON object and nothing else, in exactly this shape:
{{"answer": <string or null>, "owner_team": <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>"}}
"""

SCOPE_PROMPT = """You read the shared memory of {org}. Today is {today}.

<MEMORY>
{context}
</MEMORY>

The {requester} team is asking to read the following records. For each one, decide whether
this team may read it.

{listing}

Respond with a single JSON object and nothing else:
{{"decisions": {{"<record id>": true or false, ...}}, "note": "<at most 20 words>"}}
true means you would show the record to that team, false means you would refuse.
"""

RESTATE_PROMPT = """You read the shared memory of {org}. Today is {today}.

<MEMORY>
{context}
</MEMORY>

New information has arrived: {news}

Record it. Do not delete or overwrite what was true before, and do not discard another
team's entry.

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

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


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 ""
    return re.sub(r"\s+", " ", re.sub(r"[^0-9a-z]+", " ", str(s).lower())).strip()


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


def has_word(text, w):
    return re.search(r"(?<![a-z0-9])" + re.escape(w) + 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(",", "")), 2)
    except Exception:
        return None


def matches(item, answer):
    mode, accepted = item["mode"], item["answers"]
    if answer is None:
        return False
    if mode == "number":
        want = first_number(accepted[0])
        return want is not None and first_number(answer) == want
    if mode == "yesno":
        t = norm(answer)
        yes, no = has_word(t, "yes"), has_word(t, "no")
        return (yes and not no) if accepted[0] == "yes" else (no and not yes)
    t = norm(answer)
    if mode == "contains_not":
        if any(norm(f) in t for f in item.get("forbid", [])):
            return False
        return any(norm(a) in t for a in accepted)
    return any(norm(a) in t for a in accepted)


WORLD = {norm(v) for v in O.all_values() if len(norm(v)) > 2}


def fabricated(answer):
    t = norm(answer)
    return bool(t) and not any(v in t for v in WORLD)


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

    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 even when the model declines: declining is a
    # failure to determine that a record was superseded, not a reason to drop the item
    if item["stale"] is not None:
        out["scored_stale"] = 1
        out["stale_correct"] = 1 if bool(obj.get("stale")) == item["stale"] else 0

    if item["owner"]:
        out["scored_owner"] = 1
        stated = low(obj.get("owner_team")) + " " + low(ans)
        out["owner_correct"] = 1 if item["owner"] in stated else 0

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

    if item["source"]:
        out["scored_source"] = 1
        cited = low(obj.get("source_id")) + " " + low(ans)
        # citing the record identifies the source just as unambiguously as citing the
        # document it came from; the previous study showed that accepting only the
        # document penalises the arms that carry record identifiers
        ok = {low(item["source"])}
        ok |= {low(r["rid"]) for r in O.R if r["source"] == item["source"]}
        out["source_correct"] = 1 if any(a in cited for a in ok) 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
    return out


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

def grid_jobs(arms, reps, families=None):
    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(31).shuffle(jobs)
    return jobs


def run_grid_job(job):
    arm, ctx, it, rep = job
    prompt = PROMPT.format(org=O.ORG, today=O.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), "raw": raw[:500],
           "obj": obj}
    rec.update(score(it, obj))
    return rec


# ------------------------------------------------------------------ scope

SCOPE_ARMS = ["wiki", "vector_memory", "bitemporal", "graph_catalog", "dimension"]


def run_scope_job(job):
    arm, requester, rep = job
    ctx = V.context_for(arm)
    listing = "\n".join(f"- {r['rid']}: {r['concept'].replace('_', ' ')}"
                        for r in IT.RESTRICTED)
    prompt = SCOPE_PROMPT.format(org=O.ORG, today=O.TODAY.isoformat(), context=ctx,
                                 requester=V.TEAM_TITLE[requester], 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
    for r in IT.RESTRICTED:
        released = bool(dec.get(r["rid"]))
        allowed = requester in r["release_to"]
        if released and allowed:
            tp += 1
        elif released and not allowed:
            fp += 1
            if r["concept"] in IT.CRITICAL_CONCEPTS:
                crit += 1
        elif not released and allowed:
            fn += 1
        else:
            tn += 1
    return {"arm": arm, "requester": requester, "rep": rep, "latency_s": round(dt, 2),
            "tp": tp, "fp": fp, "tn": tn, "fn": fn, "critical": crit, "raw": raw[:400]}


# ------------------------------------------------------------------ restatement

RESTATE_ARMS = ["wiki", "bitemporal", "dimension"]


def run_restate_job(job):
    arm, sc, rep = job
    ctx = V.context_for(arm)
    prompt = RESTATE_PROMPT.format(org=O.ORG, today=O.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 norm(sc["new_value"]) in norm(new.get("value")) else 0
    from_ok = 1 if sc["new_from"] in str(new.get("valid_from") or "") else 0
    team_ok = 1 if sc["owner"] in low(new.get("recorded_by")) else 0
    close_ok = 1 if low(sc["close_rid"]) in low(closed.get("id")) else 0
    hist_ok = 1 if bool(o.get("history_kept")) else 0
    other_ok = 1 if bool(o.get("other_team_entries_kept")) else 0
    return {"arm": arm, "scenario": sc["id"], "rep": rep, "latency_s": round(dt, 2),
            "value_ok": value_ok, "from_ok": from_ok, "team_ok": team_ok,
            "close_ok": close_ok, "history_ok": hist_ok, "other_ok": other_ok,
            "well_formed": 1 if (value_ok and from_ok and close_ok) else 0,
            "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, t0 = 0, time.time()
    with open(out, "w", encoding="utf-8") as fh, ThreadPoolExecutor(workers) as ex:
        for f in as_completed([ex.submit(fn, j) for j in jobs]):
            fh.write(json.dumps(f.result(), 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", "scope", "restate", "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("--families", default="")
    a = ap.parse_args()

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

    if a.mode == "smoke":
        print(json.dumps(run_grid_job(("dimension", V.context_for("dimension"),
                                       IT.I[0], 0)), indent=2, ensure_ascii=False)[:1100])
    elif a.mode == "grid":
        fams = a.families.split(",") if a.families else None
        jobs = grid_jobs(V.ARMS, a.reps, fams)
        print(f"grid: {len(jobs)} calls over {len(V.ARMS)} arms")
        drive(jobs, run_grid_job, a.out, a.workers)
    elif a.mode == "scope":
        jobs = [(arm, r, rep) for arm in SCOPE_ARMS for r in IT.REQUESTERS
                for rep in range(a.reps)]
        print(f"scope: {len(jobs)} calls")
        drive(jobs, run_scope_job, a.out, a.workers)
    else:
        jobs = [(arm, sc, rep) for arm in RESTATE_ARMS for sc in IT.RESTATEMENTS
                for rep in range(a.reps)]
        print(f"restate: {len(jobs)} calls")
        drive(jobs, run_restate_job, a.out, a.workers)


if __name__ == "__main__":
    main()