"""Stage 1 — extract game titles + edition cues from shelf photos.

Each photo is sent to Claude vision once and the raw result is cached under
data/extract_raw/<photo>.json (making re-runs free and --only targeted).
titles.json is rebuilt from all raw files on every run, deduping identical
titles across photos unless their edition cues conflict — conflicting cues
mean two different editions on the shelf, which stay separate entries.
"""

from __future__ import annotations

import base64
import io
import json
import re
from collections.abc import Callable
from pathlib import Path

import typer

from bggpipe.config import Config
from bggpipe.fsio import atomic_write_text
from bggpipe.normalize import normalize_title

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".heic"}
MAX_LONG_EDGE = 1568  # Anthropic vision sweet spot (spec)
JPEG_QUALITY = 85
_CONFIDENCE_ORDER = {"low": 0, "medium": 1, "high": 2}
_CODE_FENCE = re.compile(r"^```(?:json)?\s*|\s*```$", re.MULTILINE)

# vision(image_b64, media_type) -> model's raw text response
VisionFn = Callable[[str, str], str]

VISION_PROMPT = """\
You are cataloging a photo of board game shelves.

List every board game or card game title visible (spines and face-out boxes),
transcribed exactly as printed. Exclude books, card sleeves, storage boxes,
and anything that is not a board/card game.

If the SAME title appears more than once with visibly different boxes, report
each as a separate entry — the owner has multiple editions of some games.

Respond with ONLY a JSON object, no prose:
{
  "titles": [
    {
      "title_raw": "title as printed on the box",
      "confidence": "high" | "medium" | "low",
      "publisher_hint": "publisher name or logo if legible, else null",
      "edition_hint": "edition wording if visible ('2nd Edition', 'Deluxe',
                       'Big Box', anniversary marks), else null",
      "year_hint": copyright or publication year as an integer, ONLY if
                   printed as publishing info (copyright line, edition
                   year). NEVER use a year that is part of the game's
                   title, theme, or subject matter — a wargame about 1942
                   is not published in 1942. When unsure whether a year is
                   thematic, use null,
      "language_hint": "language of the box text if determinable, else null",
      "art_notes": "distinctive box-art notes (colorway, artwork style)
                    that could identify the edition, else null"
    }
  ],
  "unidentified": [
    {
      "location": "where a person standing at this shelf would find the
                   box: shelf row, position, and the identified games on
                   either side of it",
      "partial_text": "any letters or word fragments you can make out,
                       else null",
      "art_notes": "color, artwork, box size/shape — anything that would
                    help identify it"
    }
  ]
}

"unidentified" is for boxes that appear to be games but whose title you
CANNOT confidently transcribe — too blurry, obscured, at too sharp an
angle, or cut off at the frame edge. It is ALWAYS better to report an
unidentifiable box here than to silently omit it, and better here than
guessing a title into "titles". Never list the same box in both arrays.
Use an empty array when everything is identified.
"""


def prepare_image(path: Path) -> tuple[str, str]:
    """Load a photo (JPEG/PNG/HEIC), downscale to <=1568px long edge, return
    (base64 JPEG, media_type). HEIC converts transparently via pillow-heif."""
    from PIL import Image
    from pillow_heif import register_heif_opener

    register_heif_opener()
    with Image.open(path) as img:
        img = img.convert("RGB")
        long_edge = max(img.size)
        if long_edge > MAX_LONG_EDGE:
            scale = MAX_LONG_EDGE / long_edge
            img = img.resize(
                (round(img.width * scale), round(img.height * scale)),
                Image.LANCZOS,
            )
        buffer = io.BytesIO()
        img.save(buffer, format="JPEG", quality=JPEG_QUALITY)
    return base64.standard_b64encode(buffer.getvalue()).decode(), "image/jpeg"


def parse_vision_response(text: str) -> tuple[list[dict], list[dict], int]:
    """Parse the model's JSON defensively: strip code fences, locate the
    payload amid any prose. Returns (title entries, unidentified sightings,
    dropped-malformed-entry count).
    Accepts either response shape: a bare JSON array (all titles) or an
    object with titles/unidentified keys."""
    cleaned = _CODE_FENCE.sub("", text).strip()
    if cleaned[:1] in ("[", "{"):
        # trim trailing prose after a leading JSON payload ("{...}\nNote:")
        end = cleaned.rfind("]" if cleaned[0] == "[" else "}")
        if end != -1:
            cleaned = cleaned[: end + 1]
    if cleaned[:1] not in ("[", "{"):
        starts = [i for i in (cleaned.find("["), cleaned.find("{")) if i != -1]
        if not starts:
            raise ValueError(f"no JSON in vision response: {text[:200]!r}")
        start = min(starts)
        end = cleaned.rfind("]" if cleaned[start] == "[" else "}")
        if end == -1:
            raise ValueError(f"unterminated JSON in vision response: {text[:200]!r}")
        cleaned = cleaned[start : end + 1]
    data = json.loads(cleaned)
    if isinstance(data, list):
        titles_raw, unidentified_raw = data, []
    elif isinstance(data, dict):
        titles_raw = data.get("titles") or []
        unidentified_raw = data.get("unidentified") or []
    else:
        raise ValueError("vision response is neither a JSON object nor array")
    titles = [e for e in titles_raw if isinstance(e, dict) and e.get("title_raw")]
    unidentified = [
        u
        for u in unidentified_raw
        if isinstance(u, dict)
        and (u.get("location") or u.get("partial_text") or u.get("art_notes"))
    ]
    dropped = len(titles_raw) - len(titles)
    return titles, unidentified, dropped


def default_vision(model: str) -> VisionFn:
    """The real Anthropic-API-backed vision callable (needs ANTHROPIC_API_KEY)."""
    import anthropic

    client = anthropic.Anthropic()

    def vision(image_b64: str, media_type: str) -> str:
        response = client.messages.create(
            model=model,
            max_tokens=4000,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": media_type,
                                "data": image_b64,
                            },
                        },
                        {"type": "text", "text": VISION_PROMPT},
                    ],
                }
            ],
        )
        return next(b.text for b in response.content if b.type == "text")

    return vision


def extract_photo(photo: Path, vision: VisionFn) -> dict:
    """One photo -> {"titles": [...], "unidentified": [...]} (the raw-cache
    file format)."""
    image_b64, media_type = prepare_image(photo)
    raw_entries, raw_unidentified, dropped = parse_vision_response(
        vision(image_b64, media_type)
    )
    unidentified = [
        {
            "location": str(u.get("location") or "").strip(),
            "partial_text": str(u.get("partial_text") or "").strip(),
            "art_notes": str(u.get("art_notes") or "").strip(),
        }
        for u in raw_unidentified
    ]
    entries = []
    for raw in raw_entries:
        year = raw.get("year_hint")
        entries.append(
            {
                "title_raw": str(raw["title_raw"]).strip(),
                "confidence": raw.get("confidence") or "medium",
                "publisher_hint": raw.get("publisher_hint") or "",
                "edition_hint": raw.get("edition_hint") or "",
                "year_hint": int(year)
                if isinstance(year, int | str) and str(year).isdigit()
                else None,
                "language_hint": raw.get("language_hint") or "",
                "art_notes": raw.get("art_notes") or "",
                "source_photos": [photo.name],
            }
        )
    return {"titles": entries, "unidentified": unidentified, "dropped": dropped}


def cues_conflict(a: dict, b: dict) -> bool:
    """Two sightings conflict if any edition cue is set on both and differs —
    that means visibly different boxes, i.e. separate editions (spec)."""
    for key in ("publisher_hint", "edition_hint", "language_hint"):
        va, vb = a.get(key) or "", b.get(key) or ""
        if va and vb and normalize_title(va) != normalize_title(vb):
            return True
    ya, yb = a.get("year_hint"), b.get("year_hint")
    return bool(ya and yb and ya != yb)


def _merge(a: dict, b: dict) -> dict:
    merged = dict(a)
    merged["source_photos"] = sorted(set(a["source_photos"]) | set(b["source_photos"]))
    if _CONFIDENCE_ORDER.get(b["confidence"], 1) > _CONFIDENCE_ORDER.get(
        a["confidence"], 1
    ):
        merged["confidence"] = b["confidence"]
    for key in (
        "publisher_hint",
        "edition_hint",
        "year_hint",
        "language_hint",
        "art_notes",
    ):
        merged[key] = merged.get(key) or b.get(key)
    return merged


# the fields a human correction may override on an extracted entry
EDIT_FIELDS = (
    "title_raw",
    "confidence",
    "publisher_hint",
    "edition_hint",
    "year_hint",
    "language_hint",
    "art_notes",
)


def _load_store(path: Path) -> list:
    """Curation stores hold irreplaceable human decisions and are committed
    (merge conflicts are a realistic corruption vector) — so a broken file
    must stop the pipeline loudly, never quietly reset it."""
    if not path.exists():
        return []
    try:
        return json.loads(path.read_text())
    except json.JSONDecodeError as err:
        raise ValueError(
            f"{path} is corrupt ({err}) — fix or delete it; it holds human "
            "review decisions, so check git history before deleting"
        ) from err


def load_title_edits(path: Path) -> list[dict]:
    """Human corrections to raw reads (fixed misspellings, cues the owner
    knows offhand). Each record: {"match": <title as displayed when the fix
    was made>, "photos": [...] to target one copy (optional), <EDIT_FIELDS
    to override>}. Applied on every rebuild, before dedupe."""
    return _load_store(path)


def record_title_edit(path: Path, record: dict) -> None:
    existing = load_title_edits(path)
    if existing and existing[-1] == record:
        return  # a retried request must not double-record
    atomic_write_text(
        path, json.dumps([*existing, record], indent=2, ensure_ascii=False) + "\n"
    )


def apply_title_edits(entries: list[dict], edits: list[dict]) -> list[dict]:
    """Apply stored corrections in order. Matching is by the title an entry
    CURRENTLY carries, so a later edit made against an earlier edit's result
    chains naturally; each record applies at most once per entry."""
    if not edits:
        return entries
    for entry in entries:
        applied: set[int] = set()
        while True:
            norm = normalize_title(entry["title_raw"])
            photos = set(entry.get("source_photos") or [])
            hit = next(
                (
                    i
                    for i, e in enumerate(edits)
                    if i not in applied
                    and normalize_title(e["match"]) == norm
                    and (not e.get("photos") or set(e["photos"]) & photos)
                ),
                None,
            )
            if hit is None:
                break
            applied.add(hit)
            for key in EDIT_FIELDS:
                if key in edits[hit]:
                    entry[key] = edits[hit][key]
    return entries


def _scoped_records(path: Path) -> list[dict]:
    """Parse a store of photo-scoped title decisions. Each stored record is
    {"title": ..., "photos": [...]} scoping the decision to the sightings
    that were on the line (a bare string is legacy: unscoped). Returns
    [{"norm": <normalized title>, "photos": set | None}]."""
    records = []
    for item in _load_store(path):
        if isinstance(item, str):
            records.append({"norm": normalize_title(item), "photos": None})
        else:
            photos = item.get("photos")
            records.append(
                {
                    "norm": normalize_title(item["title"]),
                    "photos": set(photos) if photos else None,
                }
            )
    return records


def load_title_splits(path: Path) -> list[dict]:
    """The human's split-into-copies decisions — durable: they must survive
    extract rebuilds and resolve --force."""
    return _scoped_records(path)


def load_title_removals(path: Path) -> list[dict]:
    """Titles the human removed from the catalog (not a game, a misread of
    box art, out of scope) — their sightings are filtered out of every
    rebuild, so re-extraction cannot resurrect them. Undo by deleting the
    record from the store file."""
    return _scoped_records(path)


def is_split(norm: str, photos, splits: list[dict]) -> bool:
    """Does a split decision cover this (normalized title, photo set)?
    Photo-scoped records only bind sightings from the photos that were on
    the split line — a same-named different edition keeps deduping."""
    photo_set = set(photos or [])
    return any(
        r["norm"] == norm and (r["photos"] is None or r["photos"] & photo_set)
        for r in splits
    )


def _record_scoped(path: Path, title: str, photos: list[str] | None) -> None:
    if is_split(normalize_title(title), photos or [], _scoped_records(path)):
        return  # an existing record already covers this line
    existing = _load_store(path)
    record = {"title": title, "photos": sorted(photos) if photos else None}
    atomic_write_text(
        path, json.dumps([*existing, record], indent=2, ensure_ascii=False) + "\n"
    )


def record_title_split(path: Path, title: str, photos: list[str] | None = None) -> None:
    _record_scoped(path, title, photos)


def record_title_removal(
    path: Path, title: str, photos: list[str] | None = None
) -> None:
    _record_scoped(path, title, photos)


def dedupe_entries(entries: list[dict], splits: list[dict] | None = None) -> list[dict]:
    """Collapse same-normalized-title sightings unless their cues conflict —
    or unless the human declared them split (several physical copies):
    split sightings stay one entry per photo and never absorb new ones."""
    splits = splits or []

    def covered(e: dict) -> bool:
        return is_split(e["title_normalized"], e.get("source_photos"), splits)

    result: list[dict] = []
    for entry in entries:
        entry = {**entry, "title_normalized": normalize_title(entry["title_raw"])}
        if covered(entry):
            result.append(entry)
            continue
        for existing in result:
            if (
                existing["title_normalized"] == entry["title_normalized"]
                and not covered(existing)
                and not cues_conflict(existing, entry)
            ):
                existing.update(_merge(existing, entry))
                break
        else:
            result.append(entry)
    return result


def apply_title_removals(entries: list[dict], removals: list[dict]) -> list[dict]:
    """Filter out sightings the human removed. Runs AFTER edits (records
    key on the title as displayed when removal was clicked) and before
    dedupe (so a removed sighting can't be absorbed into a survivor)."""
    if not removals:
        return entries
    return [
        e
        for e in entries
        if not is_split(
            normalize_title(e["title_raw"]), e.get("source_photos"), removals
        )
    ]


def rebuild_artifacts(
    raw_dir: Path,
    titles_path: Path,
    unidentified_path: Path,
    splits: list[dict] | None = None,
    edits: list[dict] | None = None,
    removals: list[dict] | None = None,
) -> tuple[list[dict], dict[str, list[dict]]]:
    """Regenerate titles.json and unidentified.json from the per-photo raw
    cache. A raw file is either an object with titles/unidentified or a bare
    arrays — still readable."""
    entries: list[dict] = []
    unidentified: dict[str, list[dict]] = {}
    for raw_file in sorted(raw_dir.glob("*.json")):
        try:
            data = json.loads(raw_file.read_text())
        except json.JSONDecodeError as err:
            raise ValueError(
                f"{raw_file} is corrupt ({err}) — delete it (or re-upload "
                "the photo) and run extract again"
            ) from err
        if isinstance(data, list):  # bare-array shape
            entries.extend(data)
            continue
        entries.extend(data.get("titles") or [])
        photo = raw_file.name.removesuffix(".json")
        if data.get("unidentified"):
            unidentified[photo] = data["unidentified"]
    entries = apply_title_removals(
        apply_title_edits(entries, edits or []), removals or []
    )
    deduped = dedupe_entries(entries, splits)
    titles_path.parent.mkdir(parents=True, exist_ok=True)
    atomic_write_text(
        titles_path, json.dumps(deduped, indent=2, ensure_ascii=False) + "\n"
    )
    atomic_write_text(
        unidentified_path, json.dumps(unidentified, indent=2, ensure_ascii=False) + "\n"
    )
    return deduped, unidentified


def replay_titles(cfg: Config) -> None:
    """Re-derive titles.json after a stored split or edit changed the rules.
    Prefers the raw caches (per-photo cue fidelity); without them (trimmed
    or hand-built data) it replays the current titles.json entries through
    the same edit + dedupe path, exploding multi-photo entries first so
    photo-level splits and targeted edits can take hold."""
    splits = load_title_splits(cfg.title_splits_path)
    edits = load_title_edits(cfg.title_edits_path)
    removals = load_title_removals(cfg.title_removals_path)
    raw_dir = cfg.extract_raw_dir
    raw_photos = (
        {f.name.removesuffix(".json") for f in raw_dir.glob("*.json")}
        if raw_dir.is_dir()
        else set()
    )
    known_photos: set[str] = set()
    if cfg.titles_path.exists():
        for entry in json.loads(cfg.titles_path.read_text()):
            known_photos.update(entry.get("source_photos") or [])
    # raw caches are gitignored: a fresh clone (or a killed/partial extract)
    # can have FEWER raw files than titles.json has photos — rebuilding from
    # them would silently truncate the committed catalog
    if raw_photos and raw_photos >= known_photos:
        rebuild_artifacts(
            raw_dir, cfg.titles_path, cfg.unidentified_path, splits, edits, removals
        )
        return
    if not cfg.titles_path.exists():
        return
    exploded: list[dict] = []
    for entry in json.loads(cfg.titles_path.read_text()):
        photos = entry.get("source_photos") or []
        if len(photos) > 1:
            exploded.extend({**entry, "source_photos": [p]} for p in photos)
        else:
            exploded.append(dict(entry))
    deduped = dedupe_entries(
        apply_title_removals(apply_title_edits(exploded, edits), removals), splits
    )
    atomic_write_text(
        cfg.titles_path, json.dumps(deduped, indent=2, ensure_ascii=False) + "\n"
    )


def run_extract(
    cfg: Config,
    *,
    only: str | None = None,
    force: bool = False,
    vision: VisionFn | None = None,
) -> list[dict]:
    photos = (
        sorted(
            p for p in cfg.photos_dir.iterdir() if p.suffix.lower() in IMAGE_EXTENSIONS
        )
        if cfg.photos_dir.is_dir()
        else []
    )
    if only:
        photos = [p for p in photos if p.name == only]
        if not photos:
            raise FileNotFoundError(
                f"--only {only!r}: no such photo in {cfg.photos_dir}"
            )
    if not photos:
        typer.echo(
            f"No photos found in {cfg.photos_dir}/ — add JPEG/PNG/HEIC files first."
        )
        raise typer.Exit(code=1)

    raw_dir = cfg.extract_raw_dir
    raw_dir.mkdir(parents=True, exist_ok=True)
    vision = vision or default_vision(cfg.model)

    failed: list[str] = []
    consecutive: tuple[str, int] = ("", 0)
    attempted = 0
    for photo in photos:
        raw_path = raw_dir / f"{photo.name}.json"
        if raw_path.exists() and not only and not force:
            try:
                json.loads(raw_path.read_text())
                typer.echo(f"  {photo.name}: already extracted, skipping")
                continue
            except json.JSONDecodeError:
                typer.echo(f"  {photo.name}: cache corrupt — re-extracting")
        attempted += 1
        try:
            result = extract_photo(photo, vision)
        except Exception as err:  # one bad photo must not block the rest
            failed.append(photo.name)
            typer.echo(f"  {photo.name}: FAILED ({err}) — continuing")
            kind = type(err).__name__
            consecutive = (
                kind,
                consecutive[1] + 1 if kind == consecutive[0] else 1,
            )
            if consecutive[1] >= 3:
                typer.echo(
                    "  aborting — 3 identical consecutive failures look "
                    "systemic (bad API key? wrong model in config.toml?), "
                    "not per-photo; nothing more will be attempted"
                )
                break
            continue
        consecutive = ("", 0)
        atomic_write_text(
            raw_path, json.dumps(result, indent=2, ensure_ascii=False) + "\n"
        )
        note = (
            f" ({len(result['unidentified'])} unidentified)"
            if result["unidentified"]
            else ""
        )
        if result.get("dropped"):
            # the prompt forbids silently omitting a box; so do we
            note += (
                f" — DROPPED {result['dropped']} malformed entr"
                f"{'y' if result['dropped'] == 1 else 'ies'}; inspect "
                f"{raw_path}"
            )
        typer.echo(f"  {photo.name}: {len(result['titles'])} title(s){note}")

    deduped, unidentified = rebuild_artifacts(
        raw_dir,
        cfg.titles_path,
        cfg.unidentified_path,
        load_title_splits(cfg.title_splits_path),
        load_title_edits(cfg.title_edits_path),
        load_title_removals(cfg.title_removals_path),
    )
    typer.echo(f"Wrote {len(deduped)} unique title(s) to {cfg.titles_path}.")
    if failed:
        typer.echo(
            f"\n{len(failed)} photo(s) failed extraction (re-run to retry): "
            + ", ".join(failed)
        )
        if len(failed) == attempted:
            # every attempted photo failed: this run accomplished nothing
            # and a wrapper (or the web job runner) must not report success
            raise typer.Exit(code=1)

    if unidentified:
        typer.echo(
            "\nSaw but couldn't identify — take a closer photo of each, drop "
            "it in photos/, and run extract again:"
        )
        for photo_name, sightings in unidentified.items():
            for s in sightings:
                detail = "; ".join(
                    part
                    for part in (
                        s["location"],
                        f"text visible: {s['partial_text']!r}"
                        if s["partial_text"]
                        else "",
                        s["art_notes"],
                    )
                    if part
                )
                typer.echo(f"  [{photo_name}] {detail}")

    shaky = [e for e in deduped if e["confidence"] != "high"]
    if shaky:
        typer.echo("\nLow-confidence reads worth double-checking:")
        for e in shaky:
            typer.echo(
                f"  {e['title_raw']!r} ({e['confidence']}) in "
                f"{', '.join(e['source_photos'])}"
            )
    return deduped
