"""Static export of the playground document viewer.

Experiment for URL deep links to a value's highlight (the QUE-326 capability)
without any server: the output directory is a complete static site that can be
served by any dumb file host (S3 behind CloudFront, `python -m http.server`).

What it writes, from the live playground store:

  site/
    index.html, viewer.js, viewer.css    copied from web/
    data/library.json                    the document list with derived labels
    data/{doc_key}/document.pdf          the staged PDF, byte-for-byte
    data/{doc_key}/groundings.json       every grounding row for the document,
                                         keyed by ref_id for O(1) deep-link
                                         resolution client-side

The public document identifier is the storage key (first 16 hex chars of the
content hash). The playground app deliberately keeps that key server-side, but
here it is the right public slug: opaque, stable across metadata edits, and
derived from the bytes rather than a database serial.

Deep links are hash routes the viewer resolves client-side, so the host needs
no routing rules:

  #/doc/{doc_key}                  the document, page 1
  #/doc/{doc_key}/page/{n}         a page
  #/doc/{doc_key}/ref/{ref_id}     a grounding: its page, with the highlight
                                   box drawn and focused

Run:
  uv run python experiments/static_viewer/export.py [--out DIR] [--doc KEY ...]
"""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
from pathlib import Path
from typing import Any

from quber.core.extractors.base import CELL_STATUS_REFERENCE
from quber.playground import db, metadata
from quber.settings import get_settings

WEB = Path(__file__).with_name("web")
INSPECT_CODES = frozenset(s.code for s in CELL_STATUS_REFERENCE if s.inspect)


def grounding_record(row: tuple[Any, ...]) -> dict[str, Any]:
    """One groundings row -> the JSON the viewer resolves a deep link with.

    Pages are 1-based to match the viewer. Null fields are dropped; the
    figure-value position fields (chart, label, series) are lifted to the top
    level so the sidebar can render a value without knowing the position
    shape.
    """
    ref_type, page, bbox, status, note, cell_text, position = row[1:]
    if isinstance(position, str):
        position = json.loads(position)
    rec: dict[str, Any] = {"type": ref_type, "page": (page or 0) + 1}
    if isinstance(bbox, str):
        bbox = json.loads(bbox)
    if bbox:
        rec["bbox"] = bbox
    if status:
        rec["status"] = status
        if status in INSPECT_CODES:
            rec["flagged"] = True
    if note:
        rec["note"] = note
    if cell_text:
        rec["text"] = cell_text
    for src, dst in (("chart", "chart"), ("label", "segment"), ("series", "series")):
        if position and position.get(src):
            rec[dst] = position[src]
    return rec


def _match_run(words: list, bbox: dict, target: str) -> dict | None:
    """The printed extent of `target` inside (or near) the structural box.

    Candidate words are those centered in the box expanded by one box-height
    vertically — camelot row bands drift off the glyph line on sparse layouts,
    so the true glyphs can sit partly outside the stored band. A contiguous
    run of candidates whose concatenated text equals the cell text (whitespace
    removed) is the cell's print; of several matches the one nearest the box
    center wins. No run matching means no tightening.
    """
    my = bbox["bottom"] - bbox["top"]
    lo_x, hi_x = bbox["left"] - 0.005, bbox["right"] + 0.005
    lo_y, hi_y = bbox["top"] - my, bbox["bottom"] + my
    cand = [
        w
        for w in words
        if lo_x <= (w[0] + w[2]) / 2 <= hi_x and lo_y <= (w[1] + w[3]) / 2 <= hi_y
    ]
    cand.sort(key=lambda w: (round(w[1] * 200), w[0]))
    runs = []
    for i in range(len(cand)):
        text = ""
        for j in range(i, len(cand)):
            text += cand[j][4]
            if len(text) >= len(target):
                if text == target:
                    runs.append(cand[i : j + 1])
                break
    if not runs:
        return None
    cx, cy = (bbox["left"] + bbox["right"]) / 2, (bbox["top"] + bbox["bottom"]) / 2
    def distance(run: list) -> float:
        rx = (min(w[0] for w in run) + max(w[2] for w in run)) / 2
        ry = (min(w[1] for w in run) + max(w[3] for w in run)) / 2
        return (rx - cx) ** 2 + (ry - cy) ** 2
    run = min(runs, key=distance)
    return {
        "left": min(w[0] for w in run),
        "top": min(w[1] for w in run),
        "right": max(w[2] for w in run),
        "bottom": max(w[3] for w in run),
    }


def tighten_cell_boxes(groundings: dict, pdf_path: Path) -> tuple[int, int]:
    """Shrink each table cell's box to its own printed text.

    The stored cell box is structural geometry — the column's x-extent crossed
    with the row band — which on sparse layouts runs far wider than the value
    and can sit half a line off the glyphs. Cells whose text has no match in
    the page's text layer (image-based tables, OCR divergence) keep the
    structural box. Returns (tightened, kept) counts.
    """
    import pdfplumber

    by_page: dict[int, list] = {}
    for rec in groundings.values():
        if rec["type"] == "tableCell" and rec.get("text") and rec.get("bbox"):
            by_page.setdefault(rec["page"], []).append(rec)
    if not by_page:
        return 0, 0
    tightened = kept = 0
    with pdfplumber.open(pdf_path) as pdf:
        for page1, recs in by_page.items():
            if not 1 <= page1 <= len(pdf.pages):
                kept += len(recs)
                continue
            page = pdf.pages[page1 - 1]
            w, h = page.width, page.height
            words = [
                (wd["x0"] / w, wd["top"] / h, wd["x1"] / w, wd["bottom"] / h, "".join(wd["text"].split()))
                for wd in page.extract_words()
            ]
            for rec in recs:
                target = "".join(rec["text"].split())
                box = _match_run(words, rec["bbox"], target) if target else None
                if box:
                    rec["bbox"] = box
                    tightened += 1
                else:
                    kept += 1
    return tightened, kept


def export(out: Path, only: set[str] | None) -> int:
    data_dir = get_settings().playground.data_dir
    out.mkdir(parents=True, exist_ok=True)
    for asset in WEB.iterdir():
        shutil.copy2(asset, out / asset.name)
    # Stamp the asset references with a content hash so a browser that cached
    # a previous deploy picks up changed files on a plain reload.
    index = out / "index.html"
    html = index.read_text()
    for name in ("viewer.css", "viewer.js"):
        digest = hashlib.sha256((out / name).read_bytes()).hexdigest()[:8]
        html = html.replace(name, f"{name}?v={digest}")
    index.write_text(html)

    with db.connect() as conn:
        docs = conn.execute(
            """SELECT id, doc_key, filename, page_count, ade_version, title,
                      folder, filing_type, year, period, version
               FROM ade_playground.documents
               ORDER BY lower(coalesce(title, filename))"""
        ).fetchall()

        library: list[dict[str, Any]] = []
        missing_pdf: list[str] = []
        for d in docs:
            doc_id, key, filename, page_count, ade_version, title = d[:6]
            folder, filing_type, year, period, version = d[6:]
            if only and key not in only:
                continue
            pdf = data_dir / f"{key}.pdf"
            if not pdf.exists():
                missing_pdf.append(f"{key} ({title or filename})")
                continue

            rows = conn.execute(
                """SELECT ref_id, ref_type, page, bbox, status, note, cell_text, position
                   FROM ade_playground.groundings WHERE document_id = %s""",
                (doc_id,),
            ).fetchall()
            groundings = {r[0]: grounding_record(r) for r in rows}
            # Chunk-level groundings (text, table, line_item, picture) carry no
            # cell_text; pull a display excerpt from the chunk content so the
            # viewer can list them as more than a bare ref id.
            chunk_rows = conn.execute(
                "SELECT chunk_id, content FROM ade_playground.chunks WHERE document_id = %s",
                (doc_id,),
            ).fetchall()
            content_by_id = dict(chunk_rows)
            for ref_id, rec in groundings.items():
                if "text" in rec:
                    continue
                content = content_by_id.get(ref_id)
                if content:
                    rec["text"] = " ".join(content.split())[:160]

            tightened, kept = tighten_cell_boxes(groundings, pdf)

            doc_out = out / "data" / key
            doc_out.mkdir(parents=True, exist_ok=True)
            shutil.copy2(pdf, doc_out / "document.pdf")
            (doc_out / "groundings.json").write_text(
                json.dumps(groundings, separators=(",", ":"))
            )

            values = sum(1 for g in groundings.values() if g["type"] == "figureValue")
            flags = sum(1 for g in groundings.values() if g.get("flagged"))
            library.append(
                {
                    "id": key,
                    "title": title or filename,
                    "filename": filename,
                    "label": metadata.filing_label(filing_type, year, period, version),
                    "folder": folder,
                    "page_count": page_count,
                    "format": "fusion" if ade_version == "quber-fusion" else "ade",
                    "groundings": len(groundings),
                    "values": values,
                    "flags": flags,
                }
            )
            print(
                f"  {key}  {values:>4} values  {flags:>3} flags  {len(groundings):>5} groundings"
                f"  cells tightened {tightened}/{tightened + kept}  {title or filename}"
            )

    (out / "data").mkdir(exist_ok=True)
    (out / "data" / "library.json").write_text(json.dumps(library, separators=(",", ":")))
    print(f"\n{len(library)} documents exported to {out}")
    if missing_pdf:
        print(f"\nEXCLUDED — no staged PDF on disk for {len(missing_pdf)} document(s):")
        for m in missing_pdf:
            print(f"  {m}")
    return 0 if library else 1


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--out", type=Path, default=Path(__file__).with_name("site"))
    parser.add_argument("--doc", action="append", help="doc_key to export (repeatable); default all")
    args = parser.parse_args()
    return export(args.out, set(args.doc) if args.doc else None)


if __name__ == "__main__":
    raise SystemExit(main())