"""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


def dedupe_entries(entries: list[dict]) -> list[dict]:
    """Collapse same-normalized-title sightings unless their cues conflict."""
    result: list[dict] = []
    for entry in entries:
        entry = {**entry, "title_normalized": normalize_title(entry["title_raw"])}
        for existing in result:
            if existing["title_normalized"] == entry[
                "title_normalized"
            ] and not cues_conflict(existing, entry):
                existing.update(_merge(existing, entry))
                break
        else:
            result.append(entry)
    return result


def rebuild_artifacts(
    raw_dir: Path, titles_path: Path, unidentified_path: Path
) -> 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")):
        data = json.loads(raw_file.read_text())
        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"]
    deduped = dedupe_entries(entries)
    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 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] = []
    for photo in photos:
        raw_path = raw_dir / f"{photo.name}.json"
        if raw_path.exists() and not only and not force:
            typer.echo(f"  {photo.name}: already extracted, skipping")
            continue
        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")
            continue
        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
    )
    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 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
