"""Rule-based vs LLM extractors (multi-provider) on the fixtures, same agent, same scorer.
  python compare_extractors.py --providers flint --limit 5 --sleep 4
  python compare_extractors.py --mock                               # plumbing check only (not a result)
Options: --dev-only | --heldout-only, --only ID[,ID...] (e.g. --only s18 or --only real10,real11),
--limit N (new scenarios per provider this run), --sleep S, --fresh, --allow-heldout-rerun
(voids the "scored once" held-out lock — see README).
Partial results go in results/extractor_compare_state.json, and reruns skip finished (provider, scenario) pairs.
A provider without a key, or one marked unsupported, is SKIPPED with a reason. No score is ever filled in.
Per-scenario LLM errors fall back to the rule extractor inside LLMExtractor, so llm_errors are counted and shown."""
import argparse, hashlib, json, os, sys, time
from agent import Config, run
from eval_harness import bounds_for, load_scenarios, score, summary
from bounds import load_bounds
from interfaces import FixedPlanner, LLMExtractor, LLMPlanner, MockLLMClient, OpenAICompatClient, RegexGuardExtractor, RuleExtractor
MAX_FALLBACK = 0.10  # above this share of chunks falling back to rules, a row is NOT an LLM result
from tools import MockTools

STATE = "results/extractor_compare_state.json"
LOCK = "results/heldout_scored.json"
CODE_FILES = ("interfaces.py", "tools.py", "agent.py", "bounds.py", "packing.py", "providers.json")


def code_hash():
    h = hashlib.sha256()
    for f in CODE_FILES:
        with open(f, "rb") as fh:
            h.update(fh.read())
    return h.hexdigest()[:16]


def run_one(sc, extractor, planner, base):
    ms = sc.get("max_steps", 8)
    rep = run(sc["prompt"], MockTools(sc["tools"]), bounds_for(sc, base), ms, Config(extractor=extractor, planner=planner))
    r = score(rep, sc["expected"], ms)
    r.update(id=sc["id"], held_out=sc["held_out"], llm_errors=sum("llm error" in n for n in rep["notes"]),
             llm_chunks=sum(("llm error" in n) or ("llm chunk ok" in n) for n in rep["notes"]))
    return r


def _provider_cfg(cfg, name):
    p = dict(cfg.get(name) or {})
    loc_path = "fixtures/private_unredacted/providers.local.json"
    if os.path.exists(loc_path):
        loc = (json.load(open(loc_path)).get(name) or {})
        p.update({k: v for k, v in loc.items() if not str(k).startswith("_")})
    return p


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--providers", default=""); ap.add_argument("--config", default="providers.json")
    ap.add_argument("--mock", action="store_true"); ap.add_argument("--limit", type=int)
    ap.add_argument("--sleep", type=float); ap.add_argument("--fresh", action="store_true")
    ap.add_argument("--only", default="", help="comma-separated scenario ids to keep (e.g. s18 or real10,real11)")
    ap.add_argument("--allow-heldout-rerun", action="store_true",
                    help="voids the held-out scored-once lock (results no longer count as scored once)")
    ap.add_argument("--with-synthetic", action="store_true", help="include synthetic scenarios alongside real")
    g = ap.add_mutually_exclusive_group(); g.add_argument("--dev-only", action="store_true"); g.add_argument("--heldout-only", action="store_true")
    a = ap.parse_args()
    split = "dev" if a.dev_only else "heldout" if a.heldout_only else "all"
    # --only may name synthetic (s18) or real (real10) ids; without it, keep the real-fixture default.
    scs = load_scenarios(split, real=None if (a.only or a.with_synthetic) else True)
    if a.only:
        want = {x.strip() for x in a.only.split(",") if x.strip()}
        scs = [sc for sc in scs if sc["id"] in want or any(sc["id"].startswith(w) for w in want)]
        if not scs:
            raise SystemExit(f"--only matched no scenarios: {sorted(want)}")
    watch = set(scs[0]["prompt"]["n"]) if scs else None
    cfg = json.load(open(a.config))
    base = load_bounds("bounds.json")
    os.makedirs("results", exist_ok=True)

    ch = code_hash()
    state = {} if a.fresh or not os.path.exists(STATE) else json.load(open(STATE))
    if state.get("_code_hash") != ch:
        if any(k for k in state if not str(k).startswith("_")):
            print("code/prompt/provider config changed: discarding resumed results")
        state = {"_code_hash": ch}
    else:
        state["_code_hash"] = ch
    lock = {} if not os.path.exists(LOCK) else json.load(open(LOCK))

    state["rule"] = {sc["id"]: run_one(sc, RuleExtractor(), FixedPlanner(), base) for sc in scs}
    state["regex_guards"] = {sc["id"]: run_one(sc, RegexGuardExtractor(), FixedPlanner(), base) for sc in scs}
    meta, skipped = {}, {}
    names = (["mock"] if a.mock else []) + [p for p in a.providers.split(",") if p]
    for name in names:
        if name == "mock":
            client = MockLLMClient([{"claims": [], "queries": []}]); meta[name] = {"model": "mock (plumbing only, NOT a result)"}
        else:
            p = _provider_cfg(cfg, name)
            if not p: skipped[name] = "not in config"; continue
            if not p.get("supported", True): skipped[name] = p.get("reason", "unsupported"); continue
            env = p.get("env", "SQ_LLM_API_KEY")
            key = os.environ.get(env, "")
            if not key: skipped[name] = f"${env} not set"; continue
            model = p.get("model") or os.environ.get(p.get("model_env", "SQ_LLM_MODEL"), "")
            base_url = p.get("base_url") or os.environ.get(p.get("base_url_env", "SQ_LLM_BASE_URL"), "")
            if not model or not base_url:
                skipped[name] = "model/base_url missing (set env or private local config)"; continue
            client = OpenAICompatClient(model=model, base_url=base_url, api_key=key,
                                        sleep_s=a.sleep if a.sleep is not None else p.get("sleep_s", 1.0),
                                        timeout=p.get("timeout", 120),
                                        json_mode=p.get("json_mode", True))
            meta[name] = {"model": model, "cost": p.get("cost")}
        done = state.setdefault(name, {}); n_new = 0
        if not isinstance(done, dict):
            done = {}; state[name] = done
        for sc in scs:
            if sc["id"] in done: continue
            if a.limit is not None and n_new >= a.limit: break
            if sc.get("held_out") and name != "mock" and name in lock and lock[name].get("code_hash") != ch and not a.allow_heldout_rerun:
                skipped[name] = "held-out already scored under different code; pass --allow-heldout-rerun to override"
                break
            try:
                done[sc["id"]] = run_one(sc, LLMExtractor(client, watch=watch), LLMPlanner(client), base)
            except Exception as e:  # unexpected failure: stop this provider, keep what we have
                skipped[name] = f"stopped at {sc['id']}: {type(e).__name__}: {e}"[:200]; break
            n_new += 1
            if sc.get("held_out") and name != "mock":
                entry = lock.get(name) or {"code_hash": ch, "arm": name, "held_out_scenarios": []}
                entry["code_hash"] = ch
                entry["arm"] = name
                hs = list(entry.get("held_out_scenarios") or [])
                if sc["id"] not in hs:
                    hs.append(sc["id"])
                entry["held_out_scenarios"] = hs
                lock[name] = entry
                json.dump(lock, open(LOCK, "w"), indent=1)
            json.dump(state, open(STATE, "w"), indent=1)  # checkpoint after every scenario
        meta[name]["usage_this_run"] = getattr(client, "usage", {})
    json.dump(state, open(STATE, "w"), indent=1)

    ids = [sc["id"] for sc in scs]
    lines = [f"# Extractor comparison (local-run), split={split}",
             "", "| system | model | scenarios done | pass | Wilson 95% | false imp | false cand | missed cand | llm fallback chunks | status |",
             "|---|---|---|---|---|---|---|---|---|---|"]
    rows = {}
    name_list = ["rule", "regex_guards"] + [n for n in names if (n in state and n not in skipped) or n in meta]
    for name in name_list:
        if str(name).startswith("_"): continue
        bucket = state.get(name, {})
        if not isinstance(bucket, dict): continue
        R = [bucket[i] for i in ids if i in bucket]
        if not R: continue
        s = summary(R, name)
        fb, chs = sum(r.get("llm_errors", 0) for r in R), sum(r.get("llm_chunks", 0) for r in R)
        is_llm = name not in ("rule", "regex_guards")
        if is_llm and (chs == 0 or fb / chs > MAX_FALLBACK):
            s["status"] = f"INSUFFICIENT: {fb}/{chs} chunks fell back to rules (> {MAX_FALLBACK:.0%}); not an LLM result"
            lines.append(f"| {name} | {meta.get(name, {}).get('model', '?')} | {len(R)}/{len(ids)} | - | - | - | - | - | {fb}/{chs} | {s['status']} |")
        else:
            s["status"] = "ok" if is_llm else "deterministic"
            lines.append(f"| {name} | {meta.get(name, {}).get('model', 'regex' if name == 'regex_guards' else 'rules')} | {len(R)}/{len(ids)} | {s['pass']} | {s['wilson95']} | "
                         f"{s['false_improvements']} | {s['false_candidates']} | {s['missed_candidates']} | {fb}/{chs} | {s['status']} |")
        rows[name] = s
    lines.append("")
    lines += [f"- SKIPPED {k}: {v}" for k, v in skipped.items()]
    lines.append("- Partial rows (scenarios done < total) only cover those scenarios and can't be compared head to head yet.")
    lines.append("- --allow-heldout-rerun voids the held-out 'scored once' guarantee.")
    for name in [k for k in rows if not str(rows[k]["status"]).startswith("INSUFFICIENT")]:
        for i in ids:
            bucket = state.get(name, {})
            r = bucket.get(i) if isinstance(bucket, dict) else None
            if r and not r["pass"]: lines.append(f"- {name} {'*' if r['held_out'] else ''}{i}: {'; '.join(r['fails'])}")
    md = "\n".join(lines); print(md)
    ts = "local"
    open(f"results/extractor_compare_{ts}.md", "w").write(md + "\n")
    per = {k: state[k] for k in rows if not str(k).startswith("_")}
    json.dump({"split": split, "summary": rows, "skipped": skipped, "meta": meta, "code_hash": ch,
               "per_scenario": per}, open(f"results/extractor_compare_{ts}.json", "w"), indent=1)


if __name__ == "__main__":
    main()