#!/usr/bin/env python3
"""Tolerant, evidence-producing jackpot checks for AOV reroll screenshots.

This intentionally does not use exact pixel equality.  Every screenshot is
normalized to the native ADB canvas, cropped to the two user-selected regions,
brightness-normalized, and compared after a small translation search.  Known
"asset still downloading" placeholders are accepted as ordinary outcomes so a
slow CDN does not preserve a large number of losing engines.
"""

from __future__ import annotations

import argparse
import difflib
import json
import math
import os
import subprocess
import tempfile
import unicodedata
from datetime import datetime
from pathlib import Path
from typing import Any

try:
    import cv2  # type: ignore
    import numpy as np  # type: ignore
except ImportError as exc:  # pragma: no cover
    raise SystemExit(
        "jackpot detection requires numpy and opencv-python-headless: " + str(exc)
    ) from exc


SKILL_DIR = Path(__file__).resolve().parents[1]
BASELINE_DIR = SKILL_DIR / "assets" / "baselines"
VISION_SOURCE = SKILL_DIR / "scripts" / "vision_ocr.swift"
VISION_BINARY = Path(tempfile.gettempdir()) / "aov-reroll-vision-ocr"
CANVAS = (2560, 1440)
ORDINARY_ACTIVITY_NAMES = ("次元徽章", "賞金牌", "LSR雙倍卡包")

# Coordinates are on the native 2560x1440 ADB frame, never on the Mac window.
# A is the central blue reward panel from screenshot 01.  B is the upper-left
# item under 稀有獎勵包含 from screenshot 03.
CONDITIONS: dict[str, dict[str, Any]] = {
    "activity_rewards": {
        "screenshot_suffix": "_01_activity_rewards.png",
        "roi": (205, 330, 2150, 750),
        "metric_size": (860, 300),
        "same_thresholds": {"correlation": 0.960, "mad": 0.030, "changed_ratio": 0.080},
        "baselines": [
            "activity_rewards_normal.png",
            "activity_rewards_downloading.png",
        ],
    },
    "rare_reward_top_left": {
        "screenshot_suffix": "_03_41_draw_rewards.png",
        "roi": (105, 450, 190, 190),
        "metric_size": (190, 190),
        "same_thresholds": {"correlation": 0.955, "mad": 0.040, "changed_ratio": 0.140},
        "baselines": [
            "rare_reward_top_left_normal.png",
            "rare_reward_top_left_downloading.png",
        ],
    },
}


def now_iso() -> str:
    return datetime.now().astimezone().isoformat(timespec="seconds")


def atomic_json(path: Path, value: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(value, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    os.replace(temporary, path)


def read_image(path: Path) -> Any:
    image = cv2.imread(str(path), cv2.IMREAD_COLOR)
    if image is None:
        raise RuntimeError(f"Image cannot be read: {path}")
    return image


def normalize_text(value: str) -> str:
    value = unicodedata.normalize("NFKC", value).casefold()
    return "".join(character for character in value if character.isalnum())


def ensure_vision_ocr() -> Path:
    if not VISION_SOURCE.is_file():
        raise RuntimeError(f"Vision OCR source is missing: {VISION_SOURCE}")
    if (
        VISION_BINARY.is_file()
        and VISION_BINARY.stat().st_mtime >= VISION_SOURCE.stat().st_mtime
    ):
        return VISION_BINARY
    subprocess.run(
        [
            "xcrun",
            "swiftc",
            "-O",
            "-framework",
            "Vision",
            "-framework",
            "AppKit",
            str(VISION_SOURCE),
            "-o",
            str(VISION_BINARY),
        ],
        check=True,
        capture_output=True,
        text=True,
        timeout=120,
    )
    return VISION_BINARY


def ocr_image(path: Path) -> list[dict[str, Any]]:
    completed = subprocess.run(
        [str(ensure_vision_ocr()), str(path)],
        check=True,
        capture_output=True,
        text=True,
        timeout=30,
    )
    value = json.loads(completed.stdout)
    if not isinstance(value, list):
        raise RuntimeError("Vision OCR returned a non-list result")
    return [item for item in value if isinstance(item, dict)]


def activity_name_evidence(observations: list[dict[str, Any]]) -> dict[str, Any]:
    """Classify the item-name row while ignoring icons that may be unloaded."""
    candidates: list[dict[str, Any]] = []
    for item in observations:
        text = str(item.get("text", "")).strip()
        normalized = normalize_text(text)
        y = float(item.get("y", 0.0))
        confidence = float(item.get("confidence", 0.0))
        # On the 2150x750 evidence crop, item labels occupy this narrow band.
        if not (470.0 <= y <= 590.0) or confidence < 0.25 or not normalized:
            continue
        if normalized.isdigit() or normalized == normalize_text("限時"):
            continue
        candidates.append({"text": text, "normalized": normalized, "confidence": confidence})

    found: dict[str, str] = {}
    extra: list[str] = []
    for item in candidates:
        observed = str(item["normalized"])
        best_name = ""
        best_score = 0.0
        for expected in ORDINARY_ACTIVITY_NAMES:
            target = normalize_text(expected)
            if target in observed or observed in target:
                score = min(len(observed), len(target)) / max(len(observed), len(target))
            else:
                score = difflib.SequenceMatcher(None, target, observed).ratio()
            if score > best_score:
                best_name, best_score = expected, score
        if best_score >= 0.68:
            found[best_name] = str(item["text"])
        else:
            extra.append(str(item["text"]))
    return {
        "ordinary_names_expected": list(ORDINARY_ACTIVITY_NAMES),
        "ordinary_names_found": found,
        "extra_item_names": extra,
        "raw_item_label_observations": candidates,
        "all_ordinary_names_found": len(found) == len(ORDINARY_ACTIVITY_NAMES),
    }


def normalize_canvas(image: Any) -> Any:
    target_width, target_height = CANVAS
    height, width = image.shape[:2]
    ratio = width / height
    target_ratio = target_width / target_height
    if abs(ratio - target_ratio) > 0.025:
        raise RuntimeError(
            f"Screenshot aspect ratio {width}x{height} is not close to {target_width}x{target_height}"
        )
    if (width, height) == CANVAS:
        return image
    interpolation = cv2.INTER_AREA if width > target_width else cv2.INTER_CUBIC
    return cv2.resize(image, CANVAS, interpolation=interpolation)


def crop(image: Any, rectangle: tuple[int, int, int, int]) -> Any:
    x, y, width, height = rectangle
    return image[y : y + height, x : x + width]


def _correlation(a: Any, b: Any) -> float:
    a_flat = a.reshape(-1).astype(np.float32)
    b_flat = b.reshape(-1).astype(np.float32)
    if float(a_flat.std()) < 1e-6 or float(b_flat.std()) < 1e-6:
        return 1.0 if np.array_equal(a_flat, b_flat) else 0.0
    result = float(np.corrcoef(a_flat, b_flat)[0, 1])
    return result if math.isfinite(result) else 0.0


def perceptual_metrics(reference: Any, candidate: Any, metric_size: tuple[int, int]) -> dict[str, Any]:
    """Return tolerant metrics after resize, brightness correction, and ±3px alignment."""
    ref = cv2.resize(reference, metric_size, interpolation=cv2.INTER_AREA)
    test = cv2.resize(candidate, metric_size, interpolation=cv2.INTER_AREA)
    ref = cv2.GaussianBlur(ref, (5, 5), 0)
    test = cv2.GaussianBlur(test, (5, 5), 0)

    # Correct a small global brightness/color cast before structural comparison.
    color_delta = np.median(
        ref.astype(np.float32) - test.astype(np.float32), axis=(0, 1)
    )
    color_delta = np.clip(color_delta, -18.0, 18.0)
    test = np.clip(test.astype(np.float32) + color_delta, 0, 255).astype(np.uint8)
    ref_gray = cv2.cvtColor(ref, cv2.COLOR_BGR2GRAY)
    test_gray = cv2.cvtColor(test, cv2.COLOR_BGR2GRAY)

    best: dict[str, Any] | None = None
    margin = 4
    height, width = ref_gray.shape
    for dy in range(-3, 4):
        for dx in range(-3, 4):
            x0 = margin + max(0, dx)
            x1 = width - margin + min(0, dx)
            y0 = margin + max(0, dy)
            y1 = height - margin + min(0, dy)
            ref_part = ref[y0:y1, x0:x1]
            test_part = test[y0 - dy : y1 - dy, x0 - dx : x1 - dx]
            ref_gray_part = ref_gray[y0:y1, x0:x1]
            test_gray_part = test_gray[y0 - dy : y1 - dy, x0 - dx : x1 - dx]
            difference = np.abs(
                ref_part.astype(np.int16) - test_part.astype(np.int16)
            )
            pixel_delta = np.mean(difference, axis=2)
            correlation = _correlation(ref_gray_part, test_gray_part)
            mad = float(difference.mean() / 255.0)
            changed_ratio = float(np.mean(pixel_delta > 18.0))
            rank = correlation - mad * 0.35 - changed_ratio * 0.05
            current = {
                "correlation": correlation,
                "mad": mad,
                "changed_ratio": changed_ratio,
                "translation": [dx, dy],
                "rank": rank,
            }
            if best is None or current["rank"] > best["rank"]:
                best = current
    assert best is not None
    best.pop("rank", None)
    best["color_correction_bgr"] = [round(float(value), 3) for value in color_delta]
    return best


def metrics_are_same(metrics: dict[str, Any], thresholds: dict[str, float]) -> bool:
    return bool(
        float(metrics["correlation"]) >= thresholds["correlation"]
        and float(metrics["mad"]) <= thresholds["mad"]
        and float(metrics["changed_ratio"]) <= thresholds["changed_ratio"]
    )


def find_screenshot(run_dir: Path, suffix: str) -> Path | None:
    candidates = sorted(run_dir.glob(f"*{suffix}"))
    if not candidates:
        # During a killed worker, prefix_result_screenshots may not have run yet.
        unprefixed = run_dir / suffix.removeprefix("_")
        if unprefixed.is_file():
            return unprefixed
        return None
    return candidates[-1]


def analyze_condition(
    name: str,
    definition: dict[str, Any],
    run_dir: Path,
    evidence_dir: Path,
) -> dict[str, Any]:
    screenshot = find_screenshot(run_dir, str(definition["screenshot_suffix"]))
    if screenshot is None:
        return {
            "name": name,
            "status": "missing",
            "is_jackpot": False,
            "reason": "required screenshot was not reached",
        }
    try:
        frame = normalize_canvas(read_image(screenshot))
        observed_crop = crop(frame, tuple(definition["roi"]))
    except Exception as exc:
        return {
            "name": name,
            "status": "invalid",
            "is_jackpot": False,
            "screenshot": screenshot.name,
            "reason": str(exc),
        }

    evidence_dir.mkdir(parents=True, exist_ok=True)
    evidence_path = evidence_dir / f"condition_{name}.png"
    if not cv2.imwrite(str(evidence_path), observed_crop):
        raise RuntimeError(f"Could not save evidence crop: {evidence_path}")

    comparisons: list[dict[str, Any]] = []
    for baseline_name in definition["baselines"]:
        baseline_path = BASELINE_DIR / str(baseline_name)
        baseline = read_image(baseline_path)
        metrics = perceptual_metrics(
            baseline,
            observed_crop,
            tuple(definition["metric_size"]),
        )
        comparisons.append(
            {
                "baseline": baseline_path.name,
                "same": metrics_are_same(metrics, definition["same_thresholds"]),
                **metrics,
            }
        )

    closest = max(
        comparisons,
        key=lambda item: (
            float(item["correlation"]),
            -float(item["mad"]),
            -float(item["changed_ratio"]),
        ),
    )
    matched = next((item for item in comparisons if item["same"]), None)
    name_evidence: dict[str, Any] | None = None
    if name == "activity_rewards":
        try:
            name_evidence = activity_name_evidence(ocr_image(evidence_path))
        except Exception as exc:
            name_evidence = {"ocr_error": str(exc)}

    if name_evidence and name_evidence.get("extra_item_names"):
        status = "changed"
        is_jackpot = True
        reason = "OCR found additional item name(s) beyond the three ordinary rewards"
    elif name_evidence and name_evidence.get("all_ordinary_names_found"):
        status = "normal"
        is_jackpot = False
        reason = "OCR found exactly the three ordinary item names; unloaded icons are ignored"
    elif matched is not None:
        status = "normal"
        is_jackpot = False
        reason = f"matches accepted baseline {matched['baseline']}"
    else:
        # Screenshot checkpoints are gated by the workflow.  Once no ordinary
        # or known-download baseline matches, the visible difference is the
        # user's requested candidate condition.
        status = "changed"
        is_jackpot = True
        reason = "visibly different from all accepted ordinary baselines"
    return {
        "name": name,
        "status": status,
        "is_jackpot": is_jackpot,
        "reason": reason,
        "screenshot": screenshot.name,
        "evidence_crop": evidence_path.name,
        "roi_native": list(definition["roi"]),
        "thresholds": definition["same_thresholds"],
        "closest_baseline": closest["baseline"],
        "comparisons": comparisons,
        "name_evidence": name_evidence,
    }


def analyze_run(run_dir: Path, *, write: bool = True) -> dict[str, Any]:
    run_dir = run_dir.expanduser().resolve()
    if not run_dir.is_dir():
        raise RuntimeError(f"Run directory does not exist: {run_dir}")
    evidence_dir = run_dir / "jackpot-evidence"
    conditions = [
        analyze_condition(name, definition, run_dir, evidence_dir)
        for name, definition in CONDITIONS.items()
    ]
    keep = any(bool(item["is_jackpot"]) for item in conditions)
    result = {
        "schema_version": 1,
        "analyzed_at": now_iso(),
        "run_dir": str(run_dir),
        "keep_engine": keep,
        "decision": "keep" if keep else "delete",
        "decision_rule": "activity_rewards_changed OR rare_reward_top_left_changed",
        "conditions": conditions,
    }
    if write:
        atomic_json(run_dir / "decision.json", result)
    return result


def parser() -> argparse.ArgumentParser:
    result = argparse.ArgumentParser(description=__doc__)
    result.add_argument("--run-dir", required=True)
    result.add_argument("--no-write", action="store_true")
    return result


def main() -> int:
    args = parser().parse_args()
    result = analyze_run(Path(args.run_dir), write=not args.no_write)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


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