"""Check every authored gold value against the cell it claims to come from.

The questions were authored by agents reading printed tables, so each one
names a table cell and quotes what that cell prints. This reads the cell out
of the database and compares. A question whose gold value does not match its
cell is dropped rather than corrected — a benchmark answer that had to be
repaired is not a benchmark answer.

Input is the authoring workflow's journal; output is the gold set the
comparison runs against.

Usage:
    python -m quber.playground.benchmark.verify_gold \\
        --journal <path>/journal.jsonl --out output/benchmark/gold.json
"""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple

from loguru import logger

from quber.playground import db


def norm(s: str) -> str:
    return re.sub(r"\s+", " ", (s or "").strip())


def load_documents(journal: Path) -> List[Dict[str, Any]]:
    docs = []
    for line in journal.read_text().splitlines():
        row = json.loads(line)
        if row.get("type") == "result" and isinstance(row.get("result"), dict):
            docs.append(row["result"])
    return docs


def tagged_cells(content: str) -> List[Tuple[str, str]]:
    """The id-tagged cells of one chunk's markup, as (cell id, printed text)."""
    return [
        (cid, norm(re.sub(r"<[^>]+>", "", inner)))
        for cid, inner in re.findall(r'<td id="([^"]+)"[^>]*>(.*?)</td>', content or "", re.S)
    ]


def cell_index(doc_key: str) -> Dict[Tuple[str, str], str]:
    """Every tagged cell of every chunk in one document, keyed by chunk and cell id."""
    with db.connect() as conn:
        rows = conn.execute(
            """SELECT c.chunk_id, c.content FROM ade_playground.chunks c
               JOIN ade_playground.documents d ON d.id = c.document_id
               WHERE d.doc_key = %s""",
            (doc_key,),
        ).fetchall()
    index: Dict[Tuple[str, str], str] = {}
    for chunk_id, content in rows:
        for cid, text in tagged_cells(content):
            index.setdefault((chunk_id, cid), text)
            index.setdefault(("", cid), text)  # same cell id, whichever chunk carries it
    return index


def occurrences(index: Dict[Tuple[str, str], str], printed: str) -> int:
    """How many distinct cells in the document print this exact figure.

    Counted over cell ids, not over chunks: a table's cells are repeated in
    each of its line-item chunks under the same ids, so counting chunk
    entries would report every figure as duplicated. Reported so a reviewer
    can see which questions rest on a figure that appears more than once,
    where a right-looking answer read off the wrong row would still score.
    """
    target = norm(printed)
    return len({cid for (chunk, cid), text in index.items() if not chunk and text == target})


def verify(
    docs: List[Dict[str, Any]],
    index_for: Callable[[str], Dict[Tuple[str, str], str]] = cell_index,
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
    kept: List[Dict[str, Any]] = []
    dropped: List[Dict[str, Any]] = []
    for doc in docs:
        slug = doc["slug"]
        index = index_for(slug)
        for q in doc.get("value_questions", []):
            cid = q["gold_cell_id"]
            chunk = q.get("chunk_id", "")
            printed: Optional[str] = index.get((chunk, cid)) or index.get(("", cid))
            gold = norm(q["gold_value"])
            if printed is None:
                dropped.append({**q, "slug": slug, "why": f"cell {cid} not found in {slug}"})
                continue
            if printed != gold:
                dropped.append(
                    {**q, "slug": slug, "why": f"cell {cid} prints {printed!r}, gold says {gold!r}"}
                )
                continue
            kept.append(
                {
                    "slug": slug,
                    "kind": "value",
                    "question": q["question"],
                    "gold_value": printed,
                    "gold_cell_id": cid,
                    "chunk_id": chunk,
                    "page": q.get("page"),
                    "unit": q.get("unit"),
                    "period": q.get("period"),
                    "style": q.get("style"),
                    "uniqueness_note": q.get("uniqueness_note"),
                    "occurrences_in_document": occurrences(index, printed),
                }
            )
        for q in doc.get("prose_questions", []):
            kept.append(
                {
                    "slug": slug,
                    "kind": "prose",
                    "question": q["question"],
                    "gold_value": None,
                    "gold_cell_id": None,
                    "why_prose": q.get("why_prose"),
                }
            )
    return kept, dropped


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--journal", type=Path, required=True)
    p.add_argument("--out", type=Path, default=Path("output/benchmark/gold.json"))
    args = p.parse_args()

    docs = load_documents(args.journal)
    kept, dropped = verify(docs)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(kept, indent=2))
    (args.out.parent / "gold_dropped.json").write_text(json.dumps(dropped, indent=2))

    values = [k for k in kept if k["kind"] == "value"]
    repeated = [k for k in values if k["occurrences_in_document"] > 1]
    logger.success(
        "kept {v} value + {p} prose from {d} documents; dropped {x}",
        v=len(values),
        p=len(kept) - len(values),
        d=len(docs),
        x=len(dropped),
    )
    if repeated:
        logger.warning(
            "{n} kept value questions have a gold figure printed in more than one cell",
            n=len(repeated),
        )
    for d in dropped:
        logger.warning("dropped [{s}] {q!r}: {w}", s=d["slug"], q=d["question"][:60], w=d["why"])


if __name__ == "__main__":
    main()