import argparse
import json
import re
import statistics
from math import comb

import jev
from harness import run
from rerank import rerank
from retrieve import retrieve


def page_rank(chunks, gold, section=None):
    for i, chunk in enumerate(chunks):
        if chunk["url"] in gold and section in (None, chunk["section"]):
            return i + 1
    return None


def mcnemar_exact(b, c):
    n = b + c
    if n == 0:
        return 1.0
    tail = sum(comb(n, k) for k in range(min(b, c) + 1)) / 2 ** n
    return min(1.0, 2 * tail)


def retrieval_metrics(ranks):
    hits = [r is not None and r <= 5 for r in ranks]
    top1 = sum(r == 1 for r in ranks) / len(ranks)
    mrr = statistics.mean(1 / r if r else 0 for r in ranks)
    return hits, {"recall@1": top1, "recall@5": sum(hits) / len(hits), "mrr@20": mrr, "n": len(hits)}


def mcnemar(a_ranks, b_ranks, k):
    a_hits = [r is not None and r <= k for r in a_ranks]
    b_hits = [r is not None and r <= k for r in b_ranks]
    b = sum(x and not y for x, y in zip(a_hits, b_hits))
    c = sum(y and not x for x, y in zip(a_hits, b_hits))
    return {"only_first": b, "only_second": c, "p_exact": mcnemar_exact(b, c)}


def cited_urls(answer, chunks):
    ids = {int(i) for i in re.findall(r"\[c(\d+)\]", answer)}
    return {chunks[i - 1]["url"] for i in ids if 0 < i <= len(chunks)}


def p50(values):
    return statistics.median(values) if values else None




NUM = re.compile(r"(?<![A-Za-z_\d.])\d[\d,]*(?:\.\d+)?")
CITE = re.compile(r"\[c\d+\]")


def numbers(text):
    return {n.replace(",", "") for n in NUM.findall(CITE.sub("", text))}


def number_check(q, record):
    asked = numbers(q["question"])
    said = numbers(record["answer"]) - asked
    cited = [record["chunks"][int(i) - 1] for i in re.findall(r"\[c(\d+)\]", record["answer"])
             if 0 < int(i) <= len(record["chunks"])]
    source = set().union(*(numbers(c["text"]) for c in cited)) if cited else set()
    key = [n.replace(",", "") for n in NUM.findall(q["answer"]) if n.replace(",", "") not in asked][:1]
    return {"numbers_in_cited": None if not said else said <= source,
            "gold_number_in_answer": None if not key or q["kind"] not in ("number", "table") else key[0] in numbers(record["answer"])}


def share(rows, key):
    vals = [r[key] for r in rows if r.get(key) is not None]
    return {"share": sum(vals) / len(vals) if vals else None, "n": len(vals)}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--no-jev", action="store_true")
    parser.add_argument("--decompose", action="store_true")
    parser.add_argument("--questions", default="eval/questions.jsonl")
    parser.add_argument("--tag", required=True)
    args = parser.parse_args()
    use_jev = not args.no_jev
    questions = [json.loads(line) for line in open(args.questions)]
    rows, ranks = [], {}
    for q in questions:
        record = run(q["question"], use_jev=use_jev, split=args.decompose)
        gold = set(q["gold"])
        row = {"id": q["id"], "kind": q["kind"], "question": q["question"], "answer": record["answer"],
               "subqueries": record["subqueries"], "route": record["route"], "stages": record["stages"],
               "seconds": record["seconds"], "jev_tokens": record["jev_tokens"], "jev_dollars": record["jev_dollars"],
               "chunks": [{"id": c["id"], "url": c["url"], "section": c["section"]} for c in record["chunks"]],
               "verify": record["verify"] and {k: record["verify"][k] for k in ("checked", "passed", "escalate", "claims")},
               "retrieved": bool(record["chunks"])}
        row.update(number_check(q, record))
        if gold:
            plain = retrieve(q["question"])
            orders = {"emb": plain}
            if use_jev:
                by_id = {c["id"]: c for c in record["candidates"]}
                pipeline = [by_id[i] for i in record.get("reranked_ids", [])]
                if args.decompose:
                    orders["rerank"] = rerank(q["question"], plain, None)
                    orders["decomp_emb"] = record["candidates"] or plain
                    orders["decomp_rerank"] = pipeline or orders["rerank"]
                else:
                    orders["rerank"] = pipeline or rerank(q["question"], plain, None)
            for name, order in orders.items():
                row[f"{name}_page"] = page_rank(order, gold)
                ranks.setdefault(("page", name), []).append(row[f"{name}_page"])
                if q["section"]:
                    row[f"{name}_section"] = page_rank(order, gold, q["section"])
                    ranks.setdefault(("section", name), []).append(row[f"{name}_section"])
            row["gold_cited"] = bool(cited_urls(record["answer"], record["chunks"]) & gold)
            row["all_gold_in_top5"] = gold <= {c["url"] for c in record["chunks"]}
        rows.append(row)
        print(f"{q['id']} {record['seconds']:6.1f}s emb={row.get('emb_page')} rr={row.get('rerank_page')} "
              f"cited={row.get('gold_cited')} {record['answer'][:80]!r}", flush=True)

    gold_rows = [r for r in rows if "emb_page" in r]
    metrics = {"tag": args.tag, "questions": args.questions, "n": len(rows), "decompose": args.decompose}
    for level in ("page", "section"):
        block = {name: retrieval_metrics(r)[1] for (lv, name), r in ranks.items() if lv == level}
        if (level, "rerank") in ranks:
            for k in (1, 5):
                block[f"mcnemar_emb_vs_rerank@{k}"] = mcnemar(ranks[(level, "emb")], ranks[(level, "rerank")], k)
        if (level, "decomp_rerank") in ranks:
            for k in (1, 5):
                block[f"mcnemar_emb_vs_decomp_rerank@{k}"] = mcnemar(ranks[(level, "emb")], ranks[(level, "decomp_rerank")], k)
        metrics[f"retrieval_{level}"] = block
    if use_jev:
        pairs = [(r["route"]["intent"], q["intent"]) for r, q in zip(rows, questions)]
        metrics["router"] = {"intent_accuracy": sum(a == e for a, e in pairs) / len(pairs),
                             "retrieve_decision_accuracy": sum(r["retrieved"] == bool(q["gold"]) for r, q in zip(rows, questions)) / len(rows),
                             "misses": [f"{q['id']}: {a} (want {e})" for (a, e), q in zip(pairs, questions) if a != e]}
        verified = [r for r in rows if r["verify"]]
        metrics["verify"] = {"answers_checked": len(verified),
                             "pass_rate": sum(r["verify"]["passed"] for r in verified) / max(1, len(verified)),
                             "escalated": [r["id"] for r in verified if r["verify"]["escalate"]]}
        per_stage = {}
        for c in jev.CALLS:
            s = per_stage.setdefault(c["stage"], {"calls": 0, "tokens": 0, "dollars": 0.0})
            s["calls"] += 1
            s["tokens"] += c["input_tokens"]
            s["dollars"] += c["dollars"]
        metrics["jev_cost"] = {"pipeline_tokens_per_question": statistics.mean(r["jev_tokens"] for r in rows),
                               "pipeline_dollars_per_question": statistics.mean(r["jev_dollars"] for r in rows),
                               "all_calls_dollars_total": sum(c["dollars"] for c in jev.CALLS),
                               "by_stage_all_calls": per_stage, "models": sorted({c["model"] for c in jev.CALLS})}
    metrics["grounded"] = {"gold_page_cited": sum(r["gold_cited"] for r in gold_rows) / len(gold_rows),
                           "all_gold_pages_in_context": sum(r["all_gold_in_top5"] for r in gold_rows) / len(gold_rows),
                           "dont_know": [r["id"] for r in gold_rows if "don't know from my notes" in r["answer"].lower()],
                           "numbers_in_cited_chunks": share(rows, "numbers_in_cited"),
                           "gold_number_in_answer": share(rows, "gold_number_in_answer")}
    stages = sorted({s for r in rows for s in r["stages"]})
    metrics["latency_p50_ms"] = {s: 1000 * p50([r["stages"][s]["seconds"] for r in rows if s in r["stages"]]) for s in stages}
    metrics["latency_p50_ms"]["total"] = 1000 * p50([r["seconds"] for r in rows])
    with open(f"eval/results-{args.tag}.jsonl", "w") as f:
        f.writelines(json.dumps(r) + "\n" for r in rows)
    with open(f"eval/metrics-{args.tag}.json", "w") as f:
        json.dump(metrics, f, indent=2)
    print(json.dumps(metrics, indent=2))


if __name__ == "__main__":
    main()