"""ExtractBench (Contextual AI, arXiv 2602.12247, MIT) as a deg EVAL set. Human gold:
evaluation only, never training data.

``prepare`` maps each of the 5 schemas through :func:`deg.normalize.normalize_schema`
(no leaf cap; ``evaluation_config`` and other non-deg keys are dropped, fields deg
cannot express are dropped), reports how many schemas/leaves compile, and writes one record
per document ``{"uid", "split": "eval", "domain", "context", "schema", "value",
"eval_config"}``: ``context`` is the PDF text (``<name>.txt`` next to the PDF, made with
PyMuPDF), ``value`` the gold projected onto the normalized schema. Image-only PDFs (no text
layer) are skipped and counted.

    uv run --package deg python -m deg.extractbench \\
        --root /data/deg/external/extractbench/dataset \\
        --out /data/deg/extractbench/records.jsonl

PDF -> text first (one-time; PyMuPDF is not a deg dependency, AGPL, used as a tool):

    uv run --package deg --with pymupdf python -m deg.extractbench --pdf-text

Scoring (:func:`score_record`) compares extract fields only (decide/generate fields and
array counts are reported separately): ``exact`` is value equality, ``normalized`` is
equality after case/whitespace/punctuation folding (numbers compare as floats).
"""

from __future__ import annotations

import argparse
import json
import re
import unicodedata
from pathlib import Path
from typing import Any

from deg.normalize import count_leaves, normalize_schema
from deg.schema import SchemaError, compile_schema

ROOT = Path("/data/deg/external/extractbench/dataset")
RECORDS = Path("/data/deg/extractbench/records.jsonl")


def load_schema(path: Path) -> dict[str, Any]:
    doc = json.loads(path.read_text())
    return doc["schema_definition"] if "schema_definition" in doc else doc


def raw_leaves(schema: Any, root: Any = None, ptr: str = "",
               seen: tuple[str, ...] = ()) -> dict[str, str | None]:
    """Abstract pointer -> ExtractBench ``evaluation_config`` (or None) of every leaf of the
    raw schema (local ``$ref`` followed, unions walked at the same pointer)."""
    root = schema if root is None else root
    if not isinstance(schema, dict):
        return {}
    ref = schema.get("$ref")
    if isinstance(ref, str) and ref.startswith("#/") and ref not in seen:
        target: Any = root
        for part in ref[2:].split("/"):
            target = target.get(part, {}) if isinstance(target, dict) else {}
        return raw_leaves({**target, **{k: v for k, v in schema.items() if k != "$ref"}},
                          root, ptr, (*seen, ref))
    out: dict[str, str | None] = {}
    branches = [b for k in ("anyOf", "oneOf", "allOf") for b in schema.get(k) or []]
    for name, sub in (schema.get("properties") or {}).items():
        out |= raw_leaves(sub, root, f"{ptr}/{name}", seen)
    if isinstance(schema.get("items"), dict):
        out |= raw_leaves(schema["items"], root, f"{ptr}/*", seen)
    for branch in branches:
        out |= raw_leaves(branch, root, ptr, seen)
    if not out and ptr:
        config = schema.get("evaluation_config")
        out[ptr] = config if isinstance(config, str) else None
    return out


def project(value: Any, schema: dict[str, Any]) -> Any:
    """Gold restricted to the fields the normalized schema kept."""
    if schema.get("type") == "object":
        if not isinstance(value, dict):
            return value
        return {k: project(v, schema["properties"][k]) for k, v in value.items()
                if k in schema["properties"]}
    if schema.get("type") == "array" and isinstance(value, list):
        return [project(v, schema["items"]) for v in value]
    return value


def pdf_text(root: Path) -> int:
    """Write ``<name>.txt`` (PyMuPDF text layer, pages joined by newlines) per PDF."""
    import pymupdf as fitz

    n = 0
    for pdf in sorted(root.glob("*/*/pdf+gold/*.pdf")):
        with fitz.open(pdf) as doc:
            pdf.with_suffix(".txt").write_text("\n".join(page.get_text() for page in doc))
        n += 1
    return n


def prepare(root: Path, out: Path) -> dict[str, Any]:
    report: dict[str, Any] = {"schemas": {}, "documents": 0, "skipped": {}}
    rows = []
    for schema_path in sorted(root.glob("*/*/*-schema.json")):
        raw = load_schema(schema_path)
        name = schema_path.name.removesuffix("-schema.json")
        try:
            compile_schema(raw)
            as_is = True
        except SchemaError:
            as_is = False
        schema, why = normalize_schema(raw, max_leaves=10_000, max_depth=8)
        info: dict[str, Any] = {"compiles_as_is": as_is, "normalized": why}
        report["schemas"][name] = info
        if schema is None:
            continue
        plan = compile_schema(schema)
        info["leaves_kept"] = count_leaves(schema)
        kept = {f.pointer for f in plan.fields}
        leaves = raw_leaves(raw)
        info["leaves_raw"] = len(leaves)
        info["dropped"] = sorted(p for p in leaves if p not in kept)
        info["modes"] = {m: sum(f.mode == m for f in plan.fields)
                         for m in sorted({f.mode for f in plan.fields})}
        configs = {p: c for p, c in leaves.items() if c is not None}
        domain = "/".join(schema_path.parent.relative_to(root).parts)
        for pdf in sorted(schema_path.parent.glob("pdf+gold/*.pdf")):
            gold_path = pdf.parent / (pdf.stem + ".gold.json")
            txt = pdf.with_suffix(".txt")
            if not gold_path.exists():
                report["skipped"]["no_gold"] = report["skipped"].get("no_gold", 0) + 1
                continue
            if not txt.exists():
                report["skipped"]["no_text"] = report["skipped"].get("no_text", 0) + 1
                continue
            context = txt.read_text()
            if len(context.split()) < 20:
                report["skipped"]["image_only"] = report["skipped"].get("image_only", 0) + 1
                continue
            rows.append({"uid": f"extractbench:{domain}/{pdf.stem}", "split": "eval",
                         "domain": f"extractbench:{domain}", "context": context,
                         "schema": schema,
                         "value": project(json.loads(gold_path.read_text()), schema),
                         "eval_config": configs})
    out.parent.mkdir(parents=True, exist_ok=True)
    with out.open("w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")
    report["documents"] = len(rows)
    report["words"] = sorted(len(r["context"].split()) for r in rows)
    return report


_PUNCT = re.compile(r"[\W_]+", re.UNICODE)


def normalize_value(v: Any) -> Any:
    if isinstance(v, bool) or v is None:
        return v
    if isinstance(v, (int, float)):
        return float(v)
    s = unicodedata.normalize("NFKC", str(v)).casefold()
    s = _PUNCT.sub(" ", s).strip()
    try:
        return float(s.replace(" ", "")) if re.fullmatch(r"[\d ]+(\.\d+)?", s) else s
    except ValueError:
        return s


def _flatten(value: Any, prefix: str = "") -> dict[str, Any]:
    if isinstance(value, dict):
        out: dict[str, Any] = {}
        for k, v in value.items():
            out |= _flatten(v, f"{prefix}/{k.replace('~', '~0').replace('/', '~1')}")
        return out
    if isinstance(value, list):
        out = {f"{prefix}#len": len(value)}
        for i, v in enumerate(value):
            out |= _flatten(v, f"{prefix}/{i}")
        return out
    return {prefix: value}


def _abstract(ptr: str) -> str:
    return "/".join("*" if p.isdigit() else p for p in ptr.split("/"))


def score_record(schema: dict[str, Any], gold: Any, pred: Any, align: bool = False
                 ) -> dict[str, int]:
    """Counts for one document: extract fields (gold non-null or predicted), exact and
    normalized matches; array count matches. ``align``: array items are first matched to
    their best gold item (order-insensitive, normalized equality) instead of by index."""
    if align:
        from deg.evaluate import align_arrays

        pred = align_arrays(gold, pred, lambda g, p: normalize_value(g) == normalize_value(p))
    modes = {f.pointer: f.mode for f in compile_schema(schema).fields}
    g, p = _flatten(gold), _flatten(pred)
    c = {"extract_n": 0, "exact": 0, "normalized": 0, "count_n": 0, "count_ok": 0}
    for ptr in sorted(set(g) | set(p)):
        if ptr.endswith("#len"):
            if ptr in g:
                c["count_n"] += 1
                c["count_ok"] += p.get(ptr) == g[ptr]
            continue
        if not modes.get(_abstract(ptr), "").startswith("extract."):
            continue
        gv, pv = g.get(ptr), p.get(ptr)
        if gv is None and pv is None:
            continue
        c["extract_n"] += 1
        c["exact"] += gv == pv
        c["normalized"] += normalize_value(gv) == normalize_value(pv)
    return c


def evaluate_records(model: Any, records: Path, n: int | None = None,
                     predictions: Path | None = None) -> dict[str, Any]:
    """Run :func:`deg.infer.fill` on prepared ExtractBench records; extract-field exact
    and normalized match per domain and overall (human gold; eval only). ``predictions``:
    also write one ``{uid, gold, pred, error}`` JSON line per document."""
    import gc
    import time

    import torch

    from deg.infer import fill
    from deg.request import FillOptions
    from deg.response import FillError

    rows = [json.loads(line) for line in records.read_text().splitlines() if line.strip()]
    rows = rows[:n] if n else rows
    per: dict[str, dict[str, Any]] = {}
    docs: list[dict[str, Any]] = []
    pf = predictions.open("w", encoding="utf-8") if predictions else None
    cuda = torch.cuda.is_available() and model.device.type == "cuda"
    for row in rows:
        m = per.setdefault(row["domain"], {"docs": 0, "valid": 0, "errors": {},
                                           "extract_n": 0, "exact": 0, "normalized": 0,
                                           "count_n": 0, "count_ok": 0})
        m["docs"] += 1
        doc: dict[str, Any] = {"uid": row["uid"], "words": len(row["context"].split()),
                               "tokens": len(model.renderer.context(row["context"]).ids)}
        docs.append(doc)
        if cuda:
            torch.cuda.synchronize()
            torch.cuda.reset_peak_memory_stats()
            base = torch.cuda.memory_allocated()
        t0 = time.perf_counter()
        try:
            resp = fill(model, row["context"], row["schema"], FillOptions())
            err = None
        except (FillError, torch.OutOfMemoryError) as exc:
            # An OOM is reported, never worked around by truncating the context.
            resp, err = None, (f"{type(exc).__name__}: {exc}" if isinstance(
                exc, torch.OutOfMemoryError) else str(exc))[:80]
        if cuda:
            torch.cuda.synchronize()
            doc["peak_gpu_mb"] = round(torch.cuda.max_memory_allocated() / 2**20)
            doc["peak_over_weights_mb"] = round(
                (torch.cuda.max_memory_allocated() - base) / 2**20)
        doc["seconds"] = round(time.perf_counter() - t0, 3)
        if pf:
            pf.write(json.dumps({"uid": row["uid"], "gold": row["value"],
                                 "pred": resp["value"] if resp else None, "error": err},
                                ensure_ascii=False) + "\n")
            pf.flush()
        if resp is None:
            doc["error"] = err
            m["errors"][err] = m["errors"].get(err, 0) + 1
            if cuda:  # the failed request's tensors sit in exception/frame cycles
                gc.collect()
                torch.cuda.empty_cache()
            print(json.dumps(doc), flush=True)
            continue
        m["valid"] += 1
        for k, v in score_record(row["schema"], row["value"], resp["value"]).items():
            m[k] += v
            doc[k] = v
        for k, v in score_record(row["schema"], row["value"], resp["value"], align=True).items():
            if k in ("exact", "normalized"):
                m[f"{k}_aligned"] = m.get(f"{k}_aligned", 0) + v
                doc[f"{k}_aligned"] = v
        print(json.dumps(doc), flush=True)
    if pf:
        pf.close()
    total = {k: sum(m[k] for m in per.values())
             for k in ("docs", "valid", "extract_n", "exact", "normalized", "count_n",
                       "count_ok")}
    for k in ("exact_aligned", "normalized_aligned"):
        total[k] = sum(m.get(k, 0) for m in per.values())
    for m in [*per.values(), total]:
        m["exact_rate"] = round(m["exact"] / max(1, m["extract_n"]), 4)
        m["normalized_rate"] = round(m["normalized"] / max(1, m["extract_n"]), 4)
        m["normalized_aligned_rate"] = round(m.get("normalized_aligned", 0)
                                             / max(1, m["extract_n"]), 4)
    return {"by_domain": per, "all": total, "docs": docs}


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser(description="ExtractBench -> deg eval records")
    ap.add_argument("--root", type=Path, default=ROOT)
    ap.add_argument("--out", type=Path, default=RECORDS)
    ap.add_argument("--pdf-text", action="store_true", help="first convert PDFs to .txt")
    args = ap.parse_args(argv)
    if args.pdf_text:
        print(json.dumps({"pdfs_converted": pdf_text(args.root)}))
    print(json.dumps(prepare(args.root, args.out), indent=1))


if __name__ == "__main__":
    main()