"""Replay a recorded walk through LiveSession and count what it actually decided.

The recording format is FrameRecorder's: DIR/<session>/index.jsonl plus the JPEG
bytes the client sent. Optional labels live in labels.json next to that index
(or at --labels). A label is {"frame", "box", "sku_id"}.

Template leakage is refused, not scored. A label whose crop matches a catalog
template (byte-identical, or IoU >= 0.98 against a same-sized template pasted
into the frame) is excluded and counted under leakage. Those crops are how the
store was enrolled; scoring them is training-set accuracy.

    python code/eval_live.py RECORDING_DIR [--stock-dir DIR] [--labels PATH]
                              [--json PATH] [--iou 0.5]

Prints a human count table and writes the same report as JSON (default
<recording>/eval_live.json). Missing input exits 2 with {"error": ...} on
stdout. The same pinned recording, catalog, weights, and labels produce the
same counts; timings are omitted because they are not.

No held-out labels means denominators and fingerprints only. That is not a
walk accuracy number.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

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

LEVELS = ("accepted", "ambiguous", "rejected", "unknown", "unsettled")
TAG_STATES = ("readable", "unreadable", "absent", "detected")


class EvalInputError(ValueError):
    """The recording, labels, or catalog cannot be scored as given."""


def file_sha256(path: str) -> Optional[str]:
    digest = hashlib.sha256()
    try:
        with open(path, "rb") as handle:
            for chunk in iter(lambda: handle.read(1 << 20), b""):
                digest.update(chunk)
    except OSError:
        return None
    return digest.hexdigest()


def content_sha256(payload: bytes) -> str:
    return hashlib.sha256(payload).hexdigest()


def _iou(a: Sequence[float], b: Sequence[float]) -> float:
    if len(a) != 4 or len(b) != 4:
        return 0.0
    ax1, ay1, ax2, ay2 = (float(v) for v in a)
    bx1, by1, bx2, by2 = (float(v) for v in b)
    ix1, iy1 = max(ax1, bx1), max(ay1, by1)
    ix2, iy2 = min(ax2, bx2), min(ay2, by2)
    inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
    if inter <= 0.0:
        return 0.0
    area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1)
    area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1)
    union = area_a + area_b - inter
    return inter / union if union > 0.0 else 0.0


def load_index(recording_dir: str) -> List[Dict[str, Any]]:
    path = os.path.join(recording_dir, "index.jsonl")
    if not os.path.isfile(path):
        raise EvalInputError(f"missing recording index: {path}")
    rows: List[Dict[str, Any]] = []
    with open(path, "r", encoding="utf-8") as handle:
        for number, line in enumerate(handle, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                entry = json.loads(line)
            except json.JSONDecodeError as exc:
                raise EvalInputError(f"{path}:{number}: {exc.msg}") from exc
            if not isinstance(entry, dict) or "frame" not in entry or "path" not in entry:
                raise EvalInputError(f"{path}:{number}: line needs frame and path")
            rows.append(entry)
    if not rows:
        raise EvalInputError(f"empty recording index: {path}")
    return rows


def recording_fingerprint(recording_dir: str, rows: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
    """Hash the JPEGs the client sent, in index order. Missing files are an error."""
    digest = hashlib.sha256()
    missing: List[str] = []
    for row in rows:
        name = str(row["path"])
        path = os.path.join(recording_dir, name)
        payload_hash = file_sha256(path)
        if payload_hash is None:
            missing.append(name)
            continue
        digest.update(f"{row['frame']}:{name}:{payload_hash}\n".encode("utf-8"))
    if missing:
        raise EvalInputError("missing recorded frames: " + ", ".join(missing[:8]))
    return {"sha256": digest.hexdigest(), "frames": len(rows)}


def load_labels(path: Optional[str]) -> List[Dict[str, Any]]:
    if not path:
        return []
    if not os.path.isfile(path):
        raise EvalInputError(f"missing labels: {path}")
    try:
        with open(path, "r", encoding="utf-8") as handle:
            payload = json.loads(handle.read())
    except json.JSONDecodeError as exc:
        raise EvalInputError(f"labels are not JSON: {exc.msg}") from exc
    rows = payload.get("labels") if isinstance(payload, dict) else payload
    if not isinstance(rows, list):
        raise EvalInputError("labels must be a list or {\"labels\": [...]}")
    cleaned: List[Dict[str, Any]] = []
    for index, row in enumerate(rows):
        if not isinstance(row, dict):
            raise EvalInputError(f"labels[{index}] is not an object")
        if "frame" not in row or "sku_id" not in row or "box" not in row:
            raise EvalInputError(f"labels[{index}] needs frame, box, sku_id")
        box = row["box"]
        if not isinstance(box, (list, tuple)) or len(box) != 4:
            raise EvalInputError(f"labels[{index}].box must be [x1, y1, x2, y2]")
        cleaned.append(
            {
                "frame": int(row["frame"]),
                "box": [int(v) for v in box],
                "sku_id": str(row["sku_id"]),
                "source": row.get("source"),
            }
        )
    return cleaned


def catalog_fingerprint(stock_dir: str, stock_csv: str = "stock.csv") -> Dict[str, Any]:
    """Same inputs as RecognitionEngine.catalog_version, plus the full digest and head."""
    digest = hashlib.sha256()
    csv_path = os.path.join(stock_dir, stock_csv)
    csv_hash = file_sha256(csv_path)
    if csv_hash is None:
        digest.update(b"missing-csv")
    else:
        digest.update(bytes.fromhex(csv_hash))
    try:
        entries = sorted(
            (name, os.path.getsize(os.path.join(stock_dir, name)))
            for name in os.listdir(stock_dir)
            if name.lower().endswith((".jpeg", ".jpg", ".png"))
        )
    except OSError as exc:
        raise EvalInputError(f"cannot read catalog {stock_dir}: {exc.strerror}") from exc
    for name, size in entries:
        digest.update(f"{name}:{size}".encode("utf-8"))
    full = digest.hexdigest()
    head_path = os.path.join(stock_dir, "identity_head.pt")
    return {
        "stock_dir": os.path.basename(os.path.abspath(stock_dir)),
        "version": full[:12],
        "sha256": full,
        "templates": len(entries),
        "csv_sha256": csv_hash,
        "identity_head_sha256": file_sha256(head_path),
    }


def model_fingerprint(weights_path: str) -> Dict[str, Any]:
    try:
        import ultralytics

        library = f"ultralytics-{ultralytics.__version__}"
    except Exception:
        library = "ultralytics-unknown"
    return {
        "version": f"{os.path.basename(weights_path)}@{library}",
        "weights_sha256": file_sha256(weights_path) if weights_path else None,
    }


def _decode_image(payload: bytes) -> Any:
    import cv2
    import numpy as np

    image = cv2.imdecode(np.frombuffer(payload, dtype=np.uint8), cv2.IMREAD_COLOR)
    if image is None:
        raise EvalInputError("recorded frame is not a decodable image")
    return image


def template_crops(stock_dir: str) -> List[Tuple[str, Any]]:
    """Catalog photos keyed by SKU stem (sku_l05, sku_l05_1, ...)."""
    import cv2

    crops: List[Tuple[str, Any]] = []
    try:
        names = sorted(os.listdir(stock_dir))
    except OSError:
        return crops
    for name in names:
        stem, extension = os.path.splitext(name)
        if extension.lower() not in (".jpeg", ".jpg", ".png"):
            continue
        image = cv2.imread(os.path.join(stock_dir, name))
        if image is None or image.size == 0:
            continue
        crops.append((stem, image))
    return crops


def _same_picture(crop: Any, template: Any) -> bool:
    """True when two same-sized images are the same photo after JPEG."""
    import numpy as np

    left = crop.astype("float32").reshape(-1)
    right = template.astype("float32").reshape(-1)
    if left.size < 16 or left.size != right.size:
        return False
    left = left - float(left.mean())
    right = right - float(right.mean())
    denom = float(np.sqrt((left * left).sum() * (right * right).sum()))
    if denom <= 1e-6:
        return bool(np.max(np.abs(crop.astype("int16") - template.astype("int16"))) <= 2)
    return float((left * right).sum() / denom) >= 0.98


def leaked_labels(
    recording_dir: str,
    rows: Sequence[Dict[str, Any]],
    labels: Sequence[Dict[str, Any]],
    stock_dir: str,
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
    """Drop labels whose crop is a catalog template.

    Byte identity catches a frame that is the template file. Geometric identity
    catches a template pasted into a larger frame at the labelled box.
    """
    if not labels:
        return [], []
    templates = template_crops(stock_dir)
    template_hashes = {}
    try:
        names = os.listdir(stock_dir)
    except OSError:
        names = []
    for name in names:
        digest = file_sha256(os.path.join(stock_dir, name))
        if digest:
            template_hashes.setdefault(digest, os.path.splitext(name)[0])
    by_frame: Dict[int, List[Dict[str, Any]]] = {}
    for label in labels:
        by_frame.setdefault(int(label["frame"]), []).append(label)
    held: List[Dict[str, Any]] = []
    leaked: List[Dict[str, Any]] = []
    frames = {int(row["frame"]): row for row in rows}
    for frame_id, frame_labels in by_frame.items():
        row = frames.get(frame_id)
        image = None
        payload_hash = None
        if row is not None:
            path = os.path.join(recording_dir, str(row["path"]))
            try:
                with open(path, "rb") as handle:
                    payload = handle.read()
            except OSError:
                payload = b""
            payload_hash = content_sha256(payload) if payload else None
            if payload:
                image = _decode_image(payload)
        for label in frame_labels:
            reason = None
            matched = None
            if payload_hash in template_hashes:
                reason = "frame_is_template"
                matched = template_hashes[payload_hash]
            elif image is not None:
                x1, y1, x2, y2 = label["box"]
                crop = image[max(0, y1) : max(0, y2), max(0, x1) : max(0, x2)]
                if crop.size:
                    import cv2

                    ch, cw = crop.shape[:2]
                    for stem, template in templates:
                        th, tw = template.shape[:2]
                        if abs(th - ch) > 1 or abs(tw - cw) > 1:
                            continue
                        resized = template if (th, tw) == (ch, cw) else cv2.resize(template, (cw, ch))
                        # JPEG re-encoding keeps layout and drops exact pixels.
                        if _same_picture(crop, resized):
                            reason = "crop_is_template"
                            matched = stem
                            break
            if reason:
                leaked.append({**label, "leak": reason, "template": matched})
            else:
                held.append(label)
    return held, leaked


def identity_level(track: Dict[str, Any]) -> str:
    identity = track.get("identity") or {}
    level = identity.get("identity_level")
    if level in LEVELS:
        return str(level)
    state = identity.get("state")
    if state == "recognized":
        return "accepted"
    if state == "needs_confirmation":
        return "ambiguous"
    if state == "unrecognized":
        return "rejected"
    return "unsettled"


def tag_state(track: Dict[str, Any]) -> str:
    tag = track.get("price_tag") or track.get("tag") or {}
    state = tag.get("state")
    if state in TAG_STATES:
        return str(state)
    if tag:
        return "detected"
    return "absent"


def score_tracks(tracks: Iterable[Dict[str, Any]]) -> Dict[str, Any]:
    """Denominators over settled decisions. Unsettled tracks are not in the rate."""
    counts = {level: 0 for level in LEVELS}
    tags = {state: 0 for state in TAG_STATES}
    for track in tracks:
        counts[identity_level(track)] += 1
        tags[tag_state(track)] += 1
    settled = counts["accepted"] + counts["ambiguous"] + counts["rejected"] + counts["unknown"]
    return {
        "tracks": sum(counts.values()),
        "settled": settled,
        "accepted": counts["accepted"],
        "ambiguous": counts["ambiguous"],
        "rejected": counts["rejected"],
        "unknown": counts["unknown"],
        "unsettled": counts["unsettled"],
        "accepted_rate": _rate(counts["accepted"], settled),
        "ambiguous_rate": _rate(counts["ambiguous"], settled),
        "rejected_rate": _rate(counts["rejected"], settled),
        "tags_detected": tags["detected"] + tags["readable"] + tags["unreadable"],
        "tags_readable": tags["readable"],
        "tags_unreadable": tags["unreadable"],
        "tags_absent": tags["absent"],
    }


def _rate(numerator: int, denominator: int) -> Optional[float]:
    if denominator <= 0:
        return None
    return round(numerator / denominator, 4)


def track_churn(frames: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
    """How often one track id changes the SKU it is willing to name.

    A change counts only after the id has already named a SKU. Restarting or
    dropping the label (relabels/restarts) is reported beside that, from the
    counters LiveSession already publishes.
    """
    previous: Dict[int, Optional[str]] = {}
    changes: Dict[int, int] = {}
    seen: Dict[int, int] = {}
    relabels: Dict[int, int] = {}
    restarts: Dict[int, int] = {}
    for frame in frames:
        for track in frame.get("tracks") or []:
            track_id = track.get("track_id")
            if track_id is None:
                continue
            track_id = int(track_id)
            seen[track_id] = seen.get(track_id, 0) + 1
            relabels[track_id] = max(relabels.get(track_id, 0), int(track.get("relabels") or 0))
            restarts[track_id] = max(restarts.get(track_id, 0), int(track.get("restarts") or 0))
            named = None
            if identity_level(track) == "accepted":
                named = (track.get("identity") or {}).get("sku_id")
            prior = previous.get(track_id)
            if prior and named and named != prior:
                changes[track_id] = changes.get(track_id, 0) + 1
            if named:
                previous[track_id] = named
    churned = sum(1 for count in changes.values() if count)
    return {
        "tracks": len(seen),
        "tracks_with_sku_change": churned,
        "sku_changes": sum(changes.values()),
        "relabels": sum(relabels.values()),
        "restarts": sum(restarts.values()),
        "per_track": [
            {
                "track_id": track_id,
                "observations": seen[track_id],
                "sku_changes": changes.get(track_id, 0),
                "relabels": relabels.get(track_id, 0),
                "restarts": restarts.get(track_id, 0),
            }
            for track_id in sorted(seen)
            if changes.get(track_id, 0) or relabels.get(track_id, 0) or restarts.get(track_id, 0)
        ],
    }


def labelled_top1(
    frames: Sequence[Dict[str, Any]],
    labels: Sequence[Dict[str, Any]],
    *,
    iou_min: float,
) -> Dict[str, Any]:
    """Top-1 over held-out labels. Accepted wrong SKU is a miss, not a reject.

    Denominator is labels that overlap a track. A label with no overlapping
    track is unmatched and stays out of the top-1 rate.
    """
    by_frame: Dict[int, List[Dict[str, Any]]] = {}
    for frame in frames:
        by_frame.setdefault(int(frame.get("frame", frame.get("frame_id"))), []).append(frame)
    correct = wrong = rejected = ambiguous = missed = 0
    matched = 0
    for label in labels:
        observations = by_frame.get(int(label["frame"]), [])
        best_iou = 0.0
        best: Optional[Dict[str, Any]] = None
        for observation in observations:
            for track in observation.get("tracks") or []:
                score = _iou(label["box"], track.get("box") or [])
                if score > best_iou:
                    best_iou = score
                    best = track
        if best is None or best_iou < iou_min:
            missed += 1
            continue
        matched += 1
        level = identity_level(best)
        guess = (best.get("identity") or {}).get("sku_id")
        if level == "accepted" and guess == label["sku_id"]:
            correct += 1
        elif level == "accepted":
            wrong += 1
        elif level == "ambiguous":
            ambiguous += 1
        else:
            rejected += 1
    return {
        "labels": len(labels),
        "matched": matched,
        "unmatched": missed,
        "correct": correct,
        "wrong": wrong,
        "rejected": rejected,
        "ambiguous": ambiguous,
        "top1": _rate(correct, matched),
        "iou_min": iou_min,
    }


def score_recording(
    recording_dir: str,
    *,
    stock_dir: str,
    weights_path: str = "",
    labels_path: Optional[str] = None,
    iou_min: float = 0.5,
    live_frames: Optional[Sequence[Dict[str, Any]]] = None,
) -> Dict[str, Any]:
    """Score one session directory. live_frames, when given, are the replayed results.

    Without live_frames the recorded index is what the server already decided.
    That still produces denominators; it is not a fresh model run.
    """
    recording_dir = os.path.abspath(recording_dir)
    rows = load_index(recording_dir)
    sample = recording_fingerprint(recording_dir, rows)
    labels_file = labels_path or os.path.join(recording_dir, "labels.json")
    labels = load_labels(labels_file if os.path.isfile(labels_file) or labels_path else None)
    held, leaked = leaked_labels(recording_dir, rows, labels, stock_dir)
    source = list(live_frames) if live_frames is not None else rows
    tracks = [track for frame in source for track in (frame.get("tracks") or [])]
    decisions = score_tracks(tracks)
    report = {
        "recording": recording_dir,
        "source": "live_replay" if live_frames is not None else "recorded_index",
        "fingerprints": {
            "sample": sample,
            "model": model_fingerprint(weights_path),
            "catalog": catalog_fingerprint(stock_dir),
        },
        "decisions": decisions,
        "tags": {
            "detected": decisions["tags_detected"],
            "readable": decisions["tags_readable"],
            "unreadable": decisions["tags_unreadable"],
            "absent": decisions["tags_absent"],
        },
        "churn": track_churn(source),
        "leakage": {"excluded": len(leaked), "labels": leaked},
        "labelled_top1": labelled_top1(source, held, iou_min=iou_min),
        "evidence": _evidence(labels, held, leaked, live_frames is not None),
    }
    return report


def _evidence(
    labels: Sequence[Dict[str, Any]],
    held: Sequence[Dict[str, Any]],
    leaked: Sequence[Dict[str, Any]],
    replayed: bool,
) -> List[str]:
    gaps: List[str] = []
    if not replayed:
        gaps.append("counts are from the recorded index, not a fresh LiveSession replay")
    if not labels:
        gaps.append("no held-out labels; labelled top-1 is not a walk accuracy")
    elif not held:
        gaps.append("every label matched a catalog template and was excluded")
    if leaked:
        gaps.append(f"{len(leaked)} template-derived label(s) excluded")
    return gaps


def replay_live(
    recording_dir: str,
    engine: Any,
    *,
    settings: Optional[Any] = None,
) -> List[Dict[str, Any]]:
    """Feed recorded JPEGs through one LiveSession, in index order.

    Classification is asynchronous. Each frame is scored only after its
    in-flight crops finish, then processed again so the label is on the
    payload. Without that wait the same recording does not replay twice.
    """
    from live import LiveSession

    rows = load_index(recording_dir)
    session = LiveSession("eval", engine, settings=settings)
    session.refresh_catalog()
    if hasattr(session, "enable_tags"):
        session.enable_tags()
    frames: List[Dict[str, Any]] = []
    for row in rows:
        path = os.path.join(recording_dir, str(row["path"]))
        with open(path, "rb") as handle:
            payload = handle.read()
        frame_id = int(row["frame"])
        result = session.process_frame(frame_id, payload)
        # min_hits is 2, so the first pass only detects. Keep passing the same
        # bytes until a pass schedules nothing and everything in flight lands.
        for _ in range(4):
            _wait_until_settled(session)
            result = session.process_frame(frame_id, payload)
            if int((session.snapshot().get("stats") or {}).get("classify_pending") or 0) == 0:
                break
        result["frame"] = frame_id
        frames.append(result)
    engine_shutdown = getattr(engine, "shutdown", None)
    if callable(engine_shutdown):
        engine_shutdown()
    return frames


def _wait_until_settled(session: Any, timeout: float = 30.0) -> None:
    import time

    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        pending = int((session.snapshot().get("stats") or {}).get("classify_pending") or 0)
        if pending == 0:
            return
        time.sleep(0.02)
    raise EvalInputError("classification did not finish before the replay timeout")

def format_report(report: Dict[str, Any]) -> str:
    decisions = report["decisions"]
    tags = report["tags"]
    top1 = report["labelled_top1"]
    churn = report["churn"]
    fingerprints = report["fingerprints"]
    lines = [
        f"source: {report['source']}",
        (
            f"decisions: accepted {decisions['accepted']}  ambiguous {decisions['ambiguous']}  "
            f"rejected {decisions['rejected']}  unknown {decisions['unknown']}  "
            f"unsettled {decisions['unsettled']}  settled {decisions['settled']}"
        ),
        (
            f"rates: accepted {decisions['accepted_rate']}  "
            f"ambiguous {decisions['ambiguous_rate']}  rejected {decisions['rejected_rate']}  "
            f"(denominator settled)"
        ),
        (
            f"tags: detected {tags['detected']}  readable {tags['readable']}  "
            f"unreadable {tags['unreadable']}  absent {tags['absent']}"
        ),
        (
            f"churn: tracks {churn['tracks']}  sku_changes {churn['sku_changes']}  "
            f"relabels {churn['relabels']}  restarts {churn['restarts']}"
        ),
        (
            f"labelled top-1: {top1['correct']}/{top1['matched']} = {top1['top1']}  "
            f"wrong {top1['wrong']}  unmatched {top1['unmatched']}  "
            f"leakage excluded {report['leakage']['excluded']}"
        ),
        (
            f"fingerprints: sample {fingerprints['sample']['sha256'][:12]}  "
            f"model {fingerprints['model']['version']}  "
            f"catalog {fingerprints['catalog']['version']}"
        ),
    ]
    for gap in report["evidence"]:
        lines.append(f"gap: {gap}")
    return "\n".join(lines)


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = argparse.ArgumentParser(description="Score a recorded live walk.")
    parser.add_argument("recording", nargs="?", help="session directory with index.jsonl")
    parser.add_argument("--stock-dir", default="stock")
    parser.add_argument("--weights", default="")
    parser.add_argument("--labels", default=None)
    parser.add_argument("--json", dest="json_path", default=None)
    parser.add_argument("--iou", type=float, default=0.5)
    parser.add_argument(
        "--replay",
        action="store_true",
        help="re-run frames through LiveSession (needs the recognition stack)",
    )
    parser.add_argument("--method", default="hybrid")
    args = parser.parse_args(argv)

    def fail(message: str) -> int:
        print(json.dumps({"error": message}, ensure_ascii=False))
        print(f"error: {message}", file=sys.stderr)
        return 2

    if not args.recording:
        return fail("recording directory is required")
    recording = os.path.abspath(args.recording)
    if not os.path.isdir(recording):
        return fail(f"recording directory not found: {recording}")
    try:
        live_frames = None
        weights = args.weights
        if args.replay:
            from engine import EngineSettings, build_default_engine
            from live import LiveSettings

            settings = EngineSettings(stock_dir=args.stock_dir)
            if weights:
                settings.weights_path = weights
            engine = build_default_engine(settings)
            weights = engine.settings.weights_path
            live_frames = replay_live(
                recording, engine, settings=LiveSettings(method=args.method, tags=True)
            )
        report = score_recording(
            recording,
            stock_dir=args.stock_dir,
            weights_path=weights,
            labels_path=args.labels,
            iou_min=args.iou,
            live_frames=live_frames,
        )
    except EvalInputError as exc:
        return fail(str(exc))
    text = json.dumps(report, ensure_ascii=False, indent=1, sort_keys=True)
    destination = args.json_path or os.path.join(recording, "eval_live.json")
    try:
        with open(destination, "w", encoding="utf-8") as handle:
            handle.write(text + "\n")
    except OSError as exc:
        return fail(f"cannot write {destination}: {exc.strerror}")
    print(format_report(report))
    print(f"json: {destination}")
    return 0


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