"""Measure four ways of returning an answer over one set of questions.

Every arm answers from the SAME retrieved context: retrieval runs once per
question and its chunks are handed to all four, so a difference in the results
is a difference in the return shape and not in what was retrieved.

The arms:

  baseline  the playground's current agent — prose plus cited ids
  fixed     one envelope, payload chosen from the closed union in `shapes`
  planned   a planner declares fields per question; they are compiled and applied
  declared  the caller states the shape; nothing is inferred

Scoring is deterministic. Gold values were copied out of printed table cells,
so a returned figure is compared to the gold string directly and, separately,
as a number — no judge is asked whether two figures agree. The baseline has no
value field to compare, so it is scored on whether the gold figure appears
anywhere in its prose: that measures whether the information was found, which
is the point the comparison has to be fair about.

Usage:
    python -m quber.playground.benchmark.run_shapes --gold <path>
    python -m quber.playground.benchmark.run_shapes --gold <path> --limit 4
"""

from __future__ import annotations

import argparse
import asyncio
import json
import re
import time
from decimal import Decimal
from pathlib import Path
from typing import Any, Dict, List, Optional

from loguru import logger
from pydantic import BaseModel, Field

from quber.playground import db
from quber.playground.agent import answer as baseline_answer
from quber.playground.answers import declared, fixed, planned
from quber.playground.answers.figures import as_number, norm_printed
from quber.playground.retrieval import RetrievedChunk, retrieve

OUTPUT = Path("output/benchmark/answer_shapes.json")
CONCURRENCY = 4
K = 10

ARMS = ("baseline", "fixed", "planned", "declared")


class ProseAnswer(BaseModel):
    """The shape a caller declares when it wants an explanation."""

    text: str = Field(description="The explanation, in prose.")
    cited_ids: List[str] = Field(default_factory=list)


# ---------------------------------------------------------------- normalizing

# `norm_printed` and `as_number` live in `answers.figures`, which is the only
# place a printed figure is read as a number. Scoring a run and exporting a
# batch both compare figures, and each keeping its own reading is how the same
# cell comes out with two different signs.


def values_agree(got: Optional[str], gold: str) -> Dict[str, bool]:
    """Exact-as-printed and numeric agreement between a returned and gold figure."""
    if got is None:
        return {"exact": False, "numeric": False}
    exact = norm_printed(got) == norm_printed(gold)
    a, b = as_number(got), as_number(gold)
    return {"exact": exact, "numeric": bool(a is not None and b is not None and a == b)}


def looks_like_prose(s: Optional[str]) -> bool:
    """True when a value field holds a sentence rather than a figure.

    A printed figure is short and at most a few tokens ('$1.2 million',
    '4.6x'). Four or more tokens, or a long string, means the value came back
    wrapped in words — the failure this whole exercise is about.
    """
    if not s:
        return False
    t = norm_printed(s)
    return len(t) > 40 or len(t.split()) >= 4


def gold_in_prose(prose: str, gold: str) -> bool:
    """Whether the gold figure appears in a prose answer, as printed or as a number.

    Whitespace is ignored on the literal comparison: a cell prints "$ (10,550)"
    where prose writes "$(10,550)", and that is the same figure.

    Scale words are deliberately not part of a token. Both sides drop the scale
    when they parse — "$241.4M" and "$241.4 million" each read as 241.4 — so
    absorbing a trailing "thousand" into the token achieves nothing and breaks
    any figure whose closing parenthesis then sits mid-token.
    """
    if not prose:
        return False
    if re.sub(r"\s+", "", gold) in re.sub(r"\s+", "", prose):
        return True
    g = as_number(gold)
    if g is None:
        return False
    # A decimal part must be a real decimal: `\.?\d*` would also match the full
    # stop ending a sentence, leaving a token that cannot parse. The optional
    # parenthesis is allowed on either side of the currency sign because prose
    # writes "$(10,550)" and cells print "(10,550)".
    for tok in re.findall(r"\(?-?\$?\(?[\d][\d,]*(?:\.\d+)?\)?%?[xX]?", prose):
        if as_number(tok) == g:
            return True
    return False


# ------------------------------------------------------------------ grounding


def cell_text(doc_key: str, cell_id: str) -> Optional[str]:
    """The text printed in one tagged table cell of one document, if it exists."""
    with db.connect() as conn:
        row = conn.execute(
            """SELECT c.content FROM ade_playground.chunks c
               JOIN ade_playground.documents d ON d.id = c.document_id
               WHERE d.doc_key = %s AND c.content LIKE %s
               LIMIT 1""",
            (doc_key, f'%id="{cell_id}"%'),
        ).fetchone()
    if not row:
        return None
    m = re.search(rf'<td id="{re.escape(cell_id)}"[^>]*>(.*?)</td>', row[0], re.S)
    return norm_printed(re.sub(r"<[^>]+>", "", m.group(1))) if m else None


def retrieval_hit(chunks: List[RetrievedChunk], cell_id: str) -> bool:
    return any(f'id="{cell_id}"' in c.content for c in chunks)


# ----------------------------------------------------------------------- arms


def scalar_field(obj: Any) -> Optional[str]:
    """The single figure a compiled or declared model returned, if it has one.

    A planner-compiled model names its field whatever suited the question, so
    the figure is found by taking the model's first string field that is not
    bookkeeping.
    """
    if obj is None:
        return None
    data = obj.model_dump() if isinstance(obj, BaseModel) else dict(obj)
    for key, val in data.items():
        if key in ("cited_ids", "not_found"):
            continue
        if isinstance(val, str) and val.strip():
            return val
        if isinstance(val, (int, float, Decimal)):
            return str(val)
    return None


async def run_arm(arm: str, q: Dict[str, Any], chunks: List[RetrievedChunk]) -> Dict[str, Any]:
    """One arm's attempt at one question, with what it returned and how long it took."""
    started = time.perf_counter()
    out: Dict[str, Any] = {"arm": arm}
    is_value = q["kind"] == "value"
    try:
        if arm == "baseline":
            res = await baseline_answer(q["question"], chunks)
            out.update(
                shape="prose",
                value=None,
                prose=res.answer,
                cited_ids=res.cited_ids,
            )

        elif arm == "fixed":
            res = await fixed.answer(q["question"], chunks)
            payload = res.payload
            out.update(shape=payload.kind, cited_ids=res.cited_ids)
            if payload.kind == "scalar":
                # The whole scalar, not just the printed form. Anything reading
                # this file downstream needs `number` to have a numeric column
                # without parsing `value` back into one, and a second parser of
                # printed figures is how a negative silently becomes positive.
                out.update(
                    value=payload.value,
                    number=payload.number,
                    unit=payload.unit,
                    period=payload.period,
                    label=payload.label,
                    source_id=payload.source_id,
                )
            elif payload.kind == "prose":
                out.update(value=None, prose=payload.text)
            else:
                out.update(value=None, payload=payload.model_dump(mode="json"))

        elif arm == "planned":
            plan = await planned.plan_for(q["question"])
            out["plan"] = plan.model_dump()
            if plan.wants_value and plan.fields:
                model = planned.compile_model(plan)
                res = await declared.answer(q["question"], chunks, model)
                data = res.model_dump()
                out.update(
                    shape="value",
                    value=scalar_field(res),
                    cited_ids=data.get("cited_ids", []),
                    fields=data,
                )
            else:
                res = await declared.answer(q["question"], chunks, ProseAnswer)
                out.update(shape="prose", value=None, prose=res.text, cited_ids=res.cited_ids)

        elif arm == "declared":
            # The caller knows which it wants, so the shape is supplied, not inferred.
            if is_value:
                res = await declared.answer(q["question"], chunks, declared.DeclaredScalar)
                out.update(
                    shape="value",
                    value=res.value,
                    unit=res.unit,
                    period=res.period,
                    cited_ids=res.cited_ids,
                    not_found=res.not_found,
                )
            else:
                res = await declared.answer(q["question"], chunks, ProseAnswer)
                out.update(shape="prose", value=None, prose=res.text, cited_ids=res.cited_ids)

    except Exception as exc:
        out.update(error=f"{type(exc).__name__}: {exc}", shape=None, value=None, cited_ids=[])

    out["seconds"] = round(time.perf_counter() - started, 2)
    return score(out, q)


def score(out: Dict[str, Any], q: Dict[str, Any]) -> Dict[str, Any]:
    """Add the measurements this comparison turns on."""
    is_value = q["kind"] == "value"
    shape = out.get("shape")
    out["shape_correct"] = shape in ("scalar", "value") if is_value else shape in ("prose", "unanswerable")
    out["cites_any"] = bool(out.get("cited_ids"))

    if is_value:
        gold = q["gold_value"]
        if out.get("value") is not None:
            out.update(values_agree(out["value"], gold))
            out["prose_leak"] = looks_like_prose(out["value"])
            out["found_gold"] = out["exact"] or out["numeric"]
        else:
            # No value field: the arm either answered in prose or failed.
            prose = out.get("prose") or ""
            out.update(exact=False, numeric=False, prose_leak=bool(prose))
            out["found_gold"] = gold_in_prose(prose, gold)
    return out


# ------------------------------------------------------------------ the sweep


async def run_question(q: Dict[str, Any], sem: asyncio.Semaphore) -> Dict[str, Any]:
    async with sem:
        chunks = await retrieve(q["slug"], q["question"], k=K)
    row: Dict[str, Any] = {
        **q,
        "retrieved_pages": sorted({c.page + 1 for c in chunks}),
        "retrieval_hit": retrieval_hit(chunks, q["gold_cell_id"]) if q["kind"] == "value" else None,
    }
    results = []
    for arm in ARMS:
        async with sem:
            results.append(await run_arm(arm, q, chunks))
    row["arms"] = results
    return row


def summarize(rows: List[Dict[str, Any]]) -> Dict[str, Any]:
    """Per-arm totals over the value questions and the prose controls."""
    value_rows = [r for r in rows if r["kind"] == "value"]
    hit_rows = [r for r in value_rows if r["retrieval_hit"]]
    prose_rows = [r for r in rows if r["kind"] == "prose"]
    out: Dict[str, Any] = {
        "value_questions": len(value_rows),
        "value_questions_with_retrieval_hit": len(hit_rows),
        "prose_controls": len(prose_rows),
        "arms": {},
    }

    def arm_of(row: Dict[str, Any], arm: str) -> Dict[str, Any]:
        return next(a for a in row["arms"] if a["arm"] == arm)

    for arm in ARMS:
        vs = [arm_of(r, arm) for r in hit_rows]
        ps = [arm_of(r, arm) for r in prose_rows]
        n = len(vs) or 1
        out["arms"][arm] = {
            "value_exact_as_printed": sum(a.get("exact", False) for a in vs),
            "value_numeric_match": sum(a.get("numeric", False) for a in vs),
            "found_gold_anywhere": sum(a.get("found_gold", False) for a in vs),
            "returned_a_value_field": sum(a.get("value") is not None for a in vs),
            "prose_leak": sum(a.get("prose_leak", False) for a in vs),
            "shape_correct_on_values": sum(a["shape_correct"] for a in vs),
            "shape_correct_on_prose": sum(a["shape_correct"] for a in ps),
            "cites_any": sum(a["cites_any"] for a in vs),
            "errors": sum("error" in a for a in vs + ps),
            "mean_seconds": round(sum(a["seconds"] for a in vs) / n, 2),
        }
    return out


async def run(gold_path: Path, limit: Optional[int]) -> None:
    gold = json.loads(gold_path.read_text())
    questions = gold[:limit] if limit else gold
    logger.info("{n} questions x {a} arms", n=len(questions), a=len(ARMS))
    sem = asyncio.Semaphore(CONCURRENCY)
    rows = await asyncio.gather(*(run_question(q, sem) for q in questions))
    summary = summarize(list(rows))

    OUTPUT.parent.mkdir(parents=True, exist_ok=True)
    OUTPUT.write_text(json.dumps({"summary": summary, "rows": rows}, indent=2, default=str))
    logger.success("written to {p}", p=OUTPUT)
    print(json.dumps(summary, indent=2))


def rescore(path: Path) -> None:
    """Recompute the measurements over a finished run's stored answers.

    Every arm's raw output is kept in the results file, so a correction to the
    scoring is applied by re-reading it — no model is called again, and the
    answers being scored are byte-identical to the ones originally returned.
    """
    data = json.loads(path.read_text())
    for row in data["rows"]:
        row["arms"] = [score(arm, row) for arm in row["arms"]]
    data["summary"] = summarize(data["rows"])
    path.write_text(json.dumps(data, indent=2, default=str))
    logger.success("rescored {p}", p=path)
    print(json.dumps(data["summary"], indent=2))


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--gold", type=Path, help="gold question set JSON")
    p.add_argument("--limit", type=int, default=None)
    p.add_argument("--rescore", type=Path, help="recompute scores over a finished run")
    args = p.parse_args()
    if args.rescore:
        rescore(args.rescore)
        return
    if not args.gold:
        p.error("--gold is required unless --rescore is given")
    asyncio.run(run(args.gold, args.limit))


if __name__ == "__main__":
    main()