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

Usage:
  python run.py --model claude-haiku-4-5-20251001 --reps 3 --workers 6 --out results/haiku.jsonl
  python run.py --smoke
"""

import argparse
import json
import random
import os
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 world as W
import layers as L
import tasks as T

PROMPT = """You are a data agent. Answer the question using only the tables and the context below.

<DATA>
{data}
</DATA>

<CONTEXT>
{context}
</CONTEXT>

QUESTION (as of {as_of}): {question}

Rules:
- Compute the answer from the data above. Counts are integers, money amounts are whole units.
- If the question cannot be answered from the data under any definition available to you,
  set "unanswerable" to true and "answer" to null.
- State which definition you applied. If no definition catalogue is available, use null.

Respond with a single JSON object and nothing else, in exactly this shape:
{{"answer": <number or null>, "definition_id": <string or null>, "definition_version": <integer or null>, "unanswerable": <true or false>, "note": "<at most 15 words>"}}
"""

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


def build_prompt(data, arm, task):
    return PROMPT.format(
        data=data,
        context=L.context_for(arm),
        as_of=task["as_of"].isoformat(),
        question=task["question"],
    )


SYSTEM = ("You are a precise data analyst. You compute answers from tabular data and "
          "follow the output contract exactly. You never use tools and never ask questions.")


def call_model(prompt, model, timeout=300):
    """Plain model call: tools disabled and the agent system prompt replaced, so that
    what is measured is the effect of the context layer and not of a coding harness."""
    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 = raw.strip()
    s = re.sub(r"^```(?:json)?", "", s).strip()
    s = re.sub(r"```$", "", s).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_num(v):
    if v is None:
        return None
    if isinstance(v, bool):
        return None
    if isinstance(v, (int, float)):
        return round(float(v), 2)
    s = str(v).replace(",", "").replace(" ", "").replace("$", "").replace("EUR", "")
    try:
        return round(float(s), 2)
    except Exception:
        return None


def norm_def(v):
    if v is None:
        return None
    return str(v).strip().strip(".").lower()


def score(task, obj):
    """Returns a dict of per-run outcome flags."""
    out = {
        "parsed": obj is not None,
        "answer_correct": 0,
        "def_correct": 0,
        "abstained": 0,
        "used_alt_definition": 0,
        "false_abstention": 0,
        "scored_answer": 0,
        "scored_def": 0,
        "scored_abst": 0,
    }
    if obj is None:
        return out

    unans = bool(obj.get("unanswerable"))
    ans = norm_num(obj.get("answer"))
    did = norm_def(obj.get("definition_id"))
    dver = obj.get("definition_version")
    try:
        dver = int(dver) if dver is not None else None
    except Exception:
        dver = None

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

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

    gt = norm_num(task["gt"])
    out["answer_correct"] = 1 if (ans is not None and gt is not None and ans == gt) else 0

    alt = norm_num(task.get("decisive_alt"))
    if alt is not None and gt is not None and alt != gt and ans == alt:
        out["used_alt_definition"] = 1

    want_id = norm_def(task["def_id"])
    out["def_correct"] = 1 if (did == want_id and dver == task["def_version"]) else 0
    return out


def one_run(job):
    data, model, arm, task, rep = job
    prompt = build_prompt(data, arm, task)
    raw, dt = "", 0.0
    obj = None
    for attempt in range(2):
        try:
            raw, dt = call_model(prompt, model)
            obj = parse(raw)
            if obj is not None:
                break
        except subprocess.TimeoutExpired:
            raw = "<timeout>"
        time.sleep(1.5)
    rec = {
        "model": model, "arm": arm, "task": task["id"], "family": task["family"],
        "rep": rep, "latency_s": round(dt, 2), "prompt_chars": len(prompt),
        "raw": raw[:600], "obj": obj,
    }
    rec.update(score(task, obj))
    return rec


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="claude-haiku-4-5-20251001")
    ap.add_argument("--reps", type=int, default=3)
    ap.add_argument("--workers", type=int, default=6)
    ap.add_argument("--out", default="results/run.jsonl")
    ap.add_argument("--smoke", action="store_true")
    ap.add_argument("--arms", default=",".join(L.ARMS))
    ap.add_argument("--families", default="", help="comma list, e.g. COMP,TEMP,XORG")
    args = ap.parse_args()

    w = W.build_world()
    data = W.render_data(w)
    tl = T.build_tasks(w)
    arms = args.arms.split(",")
    if args.families:
        keep = set(args.families.split(","))
        tl = [t for t in tl if t["family"] in keep]

    if args.smoke:
        tl = [t for t in tl if t["id"] in ("COMP-01", "TEMP-01", "XORG-01", "ABST-01")]
        args.reps = 1

    jobs = [(data, args.model, arm, t, r)
            for arm in arms for t in tl for r in range(args.reps)]
    # randomise submission order so time-dependent serving effects are not confounded
    # with arm (arms were previously submitted arm-by-arm)
    random.Random(11).shuffle(jobs)

    print("world: %d customers, %d orders, %d shipments, %d invoices, %d returns"
          % (len(w["customers"]), len(w["orders"]), len(w["shipments"]),
             len(w["invoices"]), len(w["returns"])))
    print("data block: %d chars" % len(data))
    print("tasks: %d | arms: %s | reps: %d | total calls: %d"
          % (len(tl), arms, args.reps, len(jobs)))

    outpath = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", args.out)
    outpath = os.path.normpath(outpath)
    os.makedirs(os.path.dirname(outpath), exist_ok=True)

    done = 0
    t0 = time.time()
    with open(outpath, "w", encoding="utf-8") as fh, \
            ThreadPoolExecutor(max_workers=args.workers) as ex:
        futs = [ex.submit(one_run, j) for j in jobs]
        for f in as_completed(futs):
            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("  %d/%d  (%.0fs elapsed)" % (done, len(jobs), time.time() - t0))
    print("written: %s" % outpath)


if __name__ == "__main__":
    main()