"""Rebuild a store catalog from packshots, refusing unsafe enrollment.

The LOPIA catalog shipped with views cut from the test shelf itself: 71-119 px
wide, one 18 px sliver, duplicate crops, and five SKUs with no view at all. That
is why every dark package collapsed onto one SKU and why the head was trained on
its own exam.

This rebuilds a catalog directory from named packshots and refuses, per view:

  * a view shorter than MIN_HEIGHT or narrower than MIN_WIDTH,
  * a view whose capture session appears in --eval-manifest,
  * a packshot whose catalog name is not supported by its OCR text,
  * a view that is a near-duplicate of another view of the same SKU,
  * a SKU left with fewer than --min-views distinct views.

Every enrolled image must have a CSV row; a SKU with an image and no row refuses
the build. CSV rows with no image are kept and reported as text-only coverage
gaps, because the text lane names and vetoes those products from the CSV alone.
Capture session ids come from the source manifest (``manifest.json`` beside the
images by default), whose image entries use ``file`` plus
``capture_session_id`` or ``session_id``. Pixel equality is not a leak test:
re-cut crops pass it, so leakage is refused by capture session instead.

    python code/rebuild_catalog.py --source /tmp/lopia-clean \
        --output stock-lopia-clean --csv stock-lopia/stock.csv
    python code/rebuild_catalog.py --source /tmp/lopia-clean \
        --output stock-lopia-clean --csv stock-lopia/stock.csv --replace

Views are named <sku>.jpg, <sku>_2.jpg, ... The CSV is carried over whole, so
text-only products stay in the catalog.
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import shutil
import sys
import tempfile
import unicodedata
from typing import Any, Dict, List, Optional, Sequence, Set, Tuple

current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

import cv2  # noqa: E402
import numpy as np  # noqa: E402

from ocr import OcrUnavailable, build_ocr_engine  # noqa: E402

# A view narrower than this cannot carry the logo and text geometry that tells
# two fine-grained products apart. The shipped LOPIA views were 71-119 px.
MIN_WIDTH = 120
# The enrollment protocol's floor: a reference view under 160 px tall carries no
# text geometry (facings under 60 px wide failed every measured lane).
MIN_HEIGHT = 160
# Correct packshots measured 0.79-1.00 of the catalog name's bigrams in their
# OCR text; the wrong-size Kewpie 450g packshot tops out at 0.64 against every
# name in the catalog. 0.70 sits between them with room for OCR noise.
MIN_NAME_OCR = 0.70
# Vision drops small print at native scale: the l11 packshot reads ハーフ only
# when upscaled (0.57 -> 0.79). Reading at 2x is what makes the check about the
# package rather than about Vision's recall.
OCR_SCALE = 2.0
# Two views of one SKU this similar are the same photograph counted twice.
DUPLICATE_COSINE = 0.98
# One view cannot be checked against anything; two is the floor for a vote.
# Official pack-shot pages publish a single front view per JAN, so a catalog
# built from them runs at 1 and says so; the floor exists to catch accidental
# single-view enrollment from shelf crops, where several views are available.
MIN_VIEWS = 2
IMAGE_EXTENSIONS = (".jpg", ".jpeg", ".png")
SESSION_KEYS = ("capture_session_id", "session_id")
SESSION_LIST_KEYS = ("sessions", "session_ids", "evaluation_sessions")


class RebuildError(Exception):
    """Refused input. The destination, if it existed, is unchanged."""


def sku_of(filename: str) -> str:
    """`sku_l01_2.jpg` -> `sku_l01`; `sku_l01.jpg` -> `sku_l01`."""
    stem = os.path.splitext(filename)[0]
    parts = stem.split("_")
    if len(parts) >= 2 and parts[0] == "sku":
        return "_".join(parts[:2])
    return stem


def view_rank(filename: str) -> int:
    """Order views so the primary view sorts first: `sku_l01.jpg` before `_2`."""
    stem = os.path.splitext(filename)[0]
    parts = stem.split("_")
    return int(parts[2]) if len(parts) > 2 and parts[2].isdigit() else 1


def cosine(left: np.ndarray, right: np.ndarray) -> float:
    a = left.astype(np.float32).ravel()
    b = right.astype(np.float32).ravel()
    denom = float(np.linalg.norm(a) * np.linalg.norm(b))
    return float((a * b).sum() / denom) if denom else 0.0


def upscale(image: np.ndarray) -> np.ndarray:
    """Enlarge for OCR only; the stored view keeps its own pixels."""
    return cv2.resize(image, None, fx=OCR_SCALE, fy=OCR_SCALE, interpolation=cv2.INTER_CUBIC)


def normalize_text(text: str) -> str:
    """NFKC, uppercase, small kana folded, alnum and `%` kept."""
    value = unicodedata.normalize("NFKC", text).upper()
    # Vision varies between small and full-size kana, and Kewpie prints
    # キユーピー with a small ユ; folding both keeps the comparison about the
    # product name rather than the typesetter.
    value = value.translate(str.maketrans("ッャュョァィゥェォヮヵヶ", "ツヤユヨアイウエオワカケ"))
    value = value.replace("キユーピー", "キユピ").replace("キューピー", "キユピ")
    return "".join(char for char in value if char.isalnum() or char == "%")


def bigrams(text: str) -> Set[str]:
    value = normalize_text(text)
    return {value[index:index + 2] for index in range(len(value) - 1)}


def name_ocr_similarity(name: str, ocr_text: str) -> float:
    """Share of the catalog name's bigrams that the OCR text explains."""
    name_bigrams = bigrams(name)
    return len(name_bigrams & bigrams(ocr_text)) / max(1, len(name_bigrams))


def collect(source: str) -> Dict[str, List[str]]:
    by_sku: Dict[str, List[str]] = {}
    for name in sorted(os.listdir(source)):
        if os.path.splitext(name)[1].lower() not in IMAGE_EXTENSIONS:
            continue
        by_sku.setdefault(sku_of(name), []).append(name)
    for views in by_sku.values():
        views.sort(key=view_rank)
    return by_sku


def read_catalog(csv_path: str) -> Tuple[List[str], List[Dict[str, str]], Dict[str, str]]:
    """Columns, rows, and SKU -> product_name. Refuses a malformed catalog."""
    with open(csv_path, newline="", encoding="utf-8-sig") as handle:
        reader = csv.DictReader(handle)
        fieldnames = list(reader.fieldnames or [])
        rows = list(reader)
    if "SKU_ID" not in fieldnames or "product_name" not in fieldnames:
        raise RebuildError(f"{csv_path}: needs SKU_ID and product_name columns")
    names: Dict[str, str] = {}
    duplicates: Set[str] = set()
    for row in rows:
        sku = (row.get("SKU_ID") or "").strip()
        if not sku:
            raise RebuildError(f"{csv_path}: a row has an empty SKU_ID")
        if sku in names:
            duplicates.add(sku)
        names[sku] = (row.get("product_name") or "").strip()
    if duplicates:
        raise RebuildError("duplicate CSV row(s): " + ", ".join(sorted(duplicates)))
    return fieldnames, rows, names


def _session_values(payload: Any) -> Set[str]:
    """Every session id anywhere in a manifest, whatever its shape."""
    found: Set[str] = set()
    if isinstance(payload, dict):
        for key, value in payload.items():
            if key in SESSION_KEYS and isinstance(value, (str, int)):
                found.add(str(value))
            elif key in SESSION_LIST_KEYS and isinstance(value, list):
                for item in value:
                    found.update([str(item)] if isinstance(item, (str, int)) else _session_values(item))
            elif isinstance(value, (dict, list)):
                found.update(_session_values(value))
    elif isinstance(payload, list):
        for item in payload:
            found.update(_session_values(item))
    return found


def load_eval_sessions(path: str) -> Set[str]:
    """Sessions an evaluation used. A manifest naming none is refused, not trusted."""
    with open(path, encoding="utf-8") as handle:
        sessions = _session_values(json.load(handle))
    if not sessions:
        raise RebuildError(f"{path}: evaluation manifest names no capture session")
    return sessions


def load_capture_sessions(path: Optional[str]) -> Dict[str, str]:
    """Image basename -> capture session, from an enrollment manifest."""
    if not path:
        return {}
    with open(path, encoding="utf-8") as handle:
        payload = json.load(handle)
    if isinstance(payload, dict):
        entries = payload.get("images", [])
        default = payload.get("capture_session_id") or payload.get("session_id")
    else:
        entries, default = payload, None
    if not isinstance(entries, list):
        raise RebuildError(f"{path}: 'images' must be a list")
    sessions: Dict[str, str] = {}
    for entry in entries:
        if not isinstance(entry, dict):
            continue
        filename = entry.get("file") or entry.get("filename") or entry.get("image")
        session = entry.get("capture_session_id") or entry.get("session_id") or default
        if filename and session is not None:
            sessions[os.path.basename(str(filename))] = str(session)
    return sessions


def audit(
    source: str,
    min_views: int = MIN_VIEWS,
    names: Optional[Dict[str, str]] = None,
    ocr_engine: Any = None,
    capture_sessions: Optional[Dict[str, str]] = None,
    eval_sessions: Optional[Set[str]] = None,
) -> Tuple[Dict[str, List[str]], List[Dict[str, str]]]:
    """Keep the views that pass, and report every refusal with its reason."""
    kept: Dict[str, List[str]] = {}
    refused: List[Dict[str, str]] = []
    capture_sessions = capture_sessions or {}
    eval_sessions = eval_sessions or set()
    for sku, views in sorted(collect(source).items()):
        accepted: List[Tuple[str, np.ndarray]] = []
        for name in views:
            image = cv2.imread(os.path.join(source, name))
            if image is None:
                refused.append({"sku": sku, "view": name, "reason": "unreadable"})
                continue
            # Leakage first: a view from an evaluation session is refused even if
            # it would fail another check, so the leak is never hidden by one.
            session = capture_sessions.get(name)
            if session and session in eval_sessions:
                refused.append({"sku": sku, "view": name,
                                "reason": f"capture session {session} is in the evaluation manifest"})
                continue
            if image.shape[0] < MIN_HEIGHT:
                refused.append({"sku": sku, "view": name,
                                "reason": f"{image.shape[0]}px tall, below the {MIN_HEIGHT}px floor"})
                continue
            if image.shape[1] < MIN_WIDTH:
                refused.append({"sku": sku, "view": name,
                                "reason": f"{image.shape[1]}px wide, below the {MIN_WIDTH}px floor"})
                continue
            if ocr_engine is not None and names is not None:
                try:
                    text = " ".join(line.text for line in ocr_engine.read(upscale(image)).lines)
                except OcrUnavailable as error:
                    print(f"OCR skipped for {name}: {error}", file=sys.stderr)
                else:
                    similarity = name_ocr_similarity(names.get(sku, ""), text)
                    if similarity < MIN_NAME_OCR:
                        refused.append({"sku": sku, "view": name,
                                        "reason": f"name/OCR {similarity:.2f} below {MIN_NAME_OCR:.2f}: "
                                                  f"{text[:100] or '<no text read>'}"})
                        continue
            duplicate = next((other for other, pixels in accepted
                              if image.shape == pixels.shape and cosine(image, pixels) >= DUPLICATE_COSINE), None)
            if duplicate:
                refused.append({"sku": sku, "view": name,
                                "reason": f"near-duplicate of {duplicate}"})
                continue
            accepted.append((name, image))
        if len(accepted) < min_views:
            for name, _ in accepted:
                refused.append({"sku": sku, "view": name,
                                "reason": f"only {len(accepted)} usable view(s), {min_views} required"})
            continue
        kept[sku] = [name for name, _ in accepted]
    return kept, refused


def write_catalog(kept: Dict[str, List[str]], source: str, output: str, csv_path: Optional[str]) -> None:
    """Copy the kept views and the whole CSV into a staging dir, then swap it in."""
    staging = tempfile.mkdtemp(prefix=".catalog-", dir=os.path.dirname(os.path.abspath(output)) or ".")
    try:
        for views in kept.values():
            for name in views:
                shutil.copy2(os.path.join(source, name), os.path.join(staging, name))
        if csv_path:
            fieldnames, rows, _ = read_catalog(csv_path)
            with open(os.path.join(staging, "stock.csv"), "w", newline="", encoding="utf-8") as handle:
                writer = csv.DictWriter(handle, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(rows)
        if os.path.exists(output):
            shutil.rmtree(output)
        os.replace(staging, output)
    except BaseException:
        shutil.rmtree(staging, ignore_errors=True)
        raise


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--source", required=True, help="directory of pack shots")
    parser.add_argument("--output", required=True, help="catalog directory to write")
    parser.add_argument("--csv", help="stock.csv to validate and carry over whole")
    parser.add_argument("--replace", action="store_true", help="overwrite an existing output")
    parser.add_argument("--min-views", type=int, default=MIN_VIEWS,
                        help=f"views required per SKU (default {MIN_VIEWS}; "
                             "official pack-shot pages publish one front view per JAN)")
    parser.add_argument("--ocr", choices=("auto", "vision", "none"), default=None,
                        help="packshot/name OCR backend (default: DORAYAKI_OCR, else vision "
                             "on macOS; an unavailable backend is skipped, not fatal)")
    parser.add_argument("--source-manifest",
                        help="image capture metadata (default: SOURCE/manifest.json when present)")
    parser.add_argument("--eval-manifest",
                        help="manifest of evaluation sessions; images captured in them are refused")
    parser.add_argument("--json", help="write the audit here")
    args = parser.parse_args(argv)

    if os.path.exists(args.output) and not args.replace:
        raise SystemExit(f"{args.output} exists; pass --replace to overwrite")

    try:
        if args.csv:
            _, _, names = read_catalog(args.csv)
            missing_rows = set(collect(args.source)) - set(names)
            if missing_rows:
                raise RebuildError("enrolled image SKU(s) with no CSV row: " + ", ".join(sorted(missing_rows)))
        else:
            names = {}
            print("no --csv: image/CSV match and name/OCR checks skipped", file=sys.stderr)

        source_manifest = args.source_manifest
        default_manifest = os.path.join(args.source, "manifest.json")
        if source_manifest is None and os.path.isfile(default_manifest):
            source_manifest = default_manifest
        capture_sessions = load_capture_sessions(source_manifest)
        eval_sessions = load_eval_sessions(args.eval_manifest) if args.eval_manifest else set()
        if eval_sessions:
            unlabelled = sorted(name for views in collect(args.source).values() for name in views
                                if name not in capture_sessions)
            if unlabelled:
                print("eval-session check cannot cover image(s) with no capture_session_id: "
                      + ", ".join(unlabelled), file=sys.stderr)

        engine = build_ocr_engine(args.ocr)
        ocr_engine = engine if engine.available() else None
        if ocr_engine is None:
            print("name/OCR check skipped: no OCR backend available", file=sys.stderr)
        kept, refused = audit(args.source, min_views=args.min_views, names=names,
                              ocr_engine=ocr_engine, capture_sessions=capture_sessions,
                              eval_sessions=eval_sessions)
        if refused:
            print(f"refused {len(refused)} view(s):", file=sys.stderr)
            for row in refused:
                print(f"  {row['sku']:<10} {row['view']:<20} {row['reason']}", file=sys.stderr)
        if not kept:
            raise RebuildError("no SKU has usable views; nothing written")
        write_catalog(kept, args.source, args.output, args.csv)
    except (OSError, ValueError, json.JSONDecodeError, RebuildError, OcrUnavailable) as error:
        raise SystemExit(str(error)) from error

    # A row with no kept image is a text-only entry: the text lane still names
    # and vetoes it from the CSV alone, so it is a coverage gap, not a refusal.
    text_only = sorted(set(names) - set(kept))

    report = {
        "source": args.source,
        "output": args.output,
        "min_width": MIN_WIDTH,
        "min_height": MIN_HEIGHT,
        "min_views": args.min_views,
        "duplicate_cosine": DUPLICATE_COSINE,
        "min_name_ocr": MIN_NAME_OCR,
        "ocr": engine.name if ocr_engine is not None else "skipped",
        "text_only_skus": text_only,
        "skus": {sku: views for sku, views in sorted(kept.items())},
        "refused": refused,
    }
    if args.json:
        with open(args.json, "w", encoding="utf-8") as handle:
            json.dump(report, handle, indent=1, ensure_ascii=False)

    print(f"{len(kept)} SKU(s), {sum(len(v) for v in kept.values())} view(s) -> {args.output}")
    for sku, views in sorted(kept.items()):
        print(f"  {sku:<10} {len(views)} view(s): {', '.join(views)}")
    if text_only:
        print(f"\ntext-only CSV row(s), no enrolled image ({len(text_only)}): {', '.join(text_only)}")
    if refused:
        print(f"\nrefused {len(refused)} view(s):")
        for row in refused:
            print(f"  {row['sku']:<10} {row['view']:<20} {row['reason']}")
    return 0


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