"""A small, local CNN that learns which gameplay frames are worth reviewing.

This is intentionally independent from Oasis.  Its job is to prioritize and
quality-check collected frames; it never supplies inputs to an Oasis checkpoint
or changes an Oasis model.
"""
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Callable


IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
Progress = Callable[[str], None]


@dataclass(slots=True)
class ReviewerTrainingResult:
    checkpoint: str
    kept_examples: int
    rejected_examples: int
    epochs: int


@dataclass(slots=True)
class FrameScore:
    path: str
    keep_probability: float
    suggestion: str


def image_paths(folder: str | Path, *, limit: int = 5_000) -> list[Path]:
    root = Path(folder).expanduser()
    if not root.is_dir():
        return []
    return [
        path for path in sorted(root.rglob("*"))
        if path.is_file() and path.suffix.casefold() in IMAGE_SUFFIXES
    ][:limit]


def reviewer_checkpoint(root: str | Path, dataset_folder: str | Path) -> Path:
    """Return an ADAM-owned checkpoint path, separate from dataset and Oasis."""
    import hashlib

    dataset = str(Path(dataset_folder).expanduser().resolve()).encode("utf-8")
    identifier = hashlib.sha1(dataset).hexdigest()[:12]
    return Path(root).expanduser().resolve() / "data" / "cnn_reviewers" / f"{identifier}.pt"


def _torch():
    try:
        import torch
        from torch import nn
    except ImportError as exc:  # pragma: no cover - controlled by application install
        raise RuntimeError("CNN Reviewer needs PyTorch. Install the ADAM requirements first.") from exc
    return torch, nn


def make_reviewer_model():
    """Build a deliberately small binary CNN for local frame triage."""
    _torch_module, nn = _torch()
    return nn.Sequential(
        nn.Conv2d(3, 16, kernel_size=5, stride=2, padding=2), nn.ReLU(),
        nn.Conv2d(16, 32, kernel_size=3, stride=2, padding=1), nn.ReLU(),
        nn.Conv2d(32, 48, kernel_size=3, stride=2, padding=1), nn.ReLU(),
        nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(48, 1),
    )


def _load_image(path: Path):
    torch, _nn = _torch()
    from PIL import Image

    with Image.open(path) as image:
        image = image.convert("RGB").resize((128, 72))
        # No torchvision transform is required, which keeps this feature portable.
        pixels = image.get_flattened_data() if hasattr(image, "get_flattened_data") else image.getdata()
        values = torch.tensor(list(pixels), dtype=torch.float32)
    return values.reshape(72, 128, 3).permute(2, 0, 1).div_(255.0)


def _labeled_paths(decisions: dict[str, str], *, per_class_limit: int = 250) -> tuple[list[Path], list[float]]:
    grouped: dict[str, list[Path]] = {"keep": [], "reject": []}
    for raw_path, decision in decisions.items():
        if decision not in grouped:
            continue
        path = Path(raw_path).expanduser()
        if path.is_file() and path.suffix.casefold() in IMAGE_SUFFIXES:
            grouped[decision].append(path)
    kept = sorted(grouped["keep"])[:per_class_limit]
    rejected = sorted(grouped["reject"])[:per_class_limit]
    return kept + rejected, [1.0] * len(kept) + [0.0] * len(rejected)


def train_reviewer(
    root: str | Path,
    dataset_folder: str | Path,
    decisions: dict[str, str],
    *,
    epochs: int = 8,
    progress: Progress | None = None,
) -> ReviewerTrainingResult:
    """Train a frame-quality CNN from explicit Keep and Reject review decisions."""
    torch, nn = _torch()
    paths, labels = _labeled_paths(decisions)
    keep_count = int(sum(labels))
    reject_count = len(labels) - keep_count
    if min(keep_count, reject_count) < 8:
        raise ValueError("Review at least 8 Keep and 8 Reject frames before training the CNN reviewer.")
    if progress:
        progress(f"Loading {len(paths)} reviewed frame(s)…")
    images = []
    valid_labels = []
    for path, label in zip(paths, labels):
        try:
            images.append(_load_image(path))
            valid_labels.append(label)
        except Exception:
            continue
    keep_count = int(sum(valid_labels))
    reject_count = len(valid_labels) - keep_count
    if min(keep_count, reject_count) < 8:
        raise ValueError("Some reviewed images could not be read; at least 8 valid Keep and Reject frames are needed.")
    torch.manual_seed(7)
    inputs = torch.stack(images)
    targets = torch.tensor(valid_labels, dtype=torch.float32).unsqueeze(1)
    model = make_reviewer_model()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    loss_fn = nn.BCEWithLogitsLoss()
    model.train()
    for epoch in range(max(1, min(int(epochs), 50))):
        order = torch.randperm(len(inputs))
        loss_value = 0.0
        batches = 0
        for start in range(0, len(order), 16):
            batch = order[start:start + 16]
            optimizer.zero_grad()
            loss = loss_fn(model(inputs[batch]), targets[batch])
            loss.backward()
            optimizer.step()
            loss_value += float(loss.detach())
            batches += 1
        if progress:
            progress(f"CNN reviewer epoch {epoch + 1}/{epochs} · loss {loss_value / max(1, batches):.3f}")
    checkpoint = reviewer_checkpoint(root, dataset_folder)
    checkpoint.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"state_dict": model.state_dict(), "image_size": [128, 72]}, checkpoint)
    return ReviewerTrainingResult(str(checkpoint), keep_count, reject_count, epochs)


def score_frames(checkpoint: str | Path, folder: str | Path, *, limit: int = 5_000, progress: Progress | None = None) -> list[FrameScore]:
    """Return review suggestions without modifying the dataset or its decisions."""
    torch, _nn = _torch()
    saved = torch.load(Path(checkpoint), map_location="cpu", weights_only=True)
    model = make_reviewer_model()
    model.load_state_dict(saved["state_dict"])
    model.eval()
    paths = image_paths(folder, limit=limit)
    scores: list[FrameScore] = []
    with torch.no_grad():
        for index, path in enumerate(paths, 1):
            try:
                probability = float(torch.sigmoid(model(_load_image(path).unsqueeze(0))).item())
            except Exception:
                continue
            suggestion = "keep" if probability >= 0.70 else "reject" if probability <= 0.30 else "review"
            scores.append(FrameScore(str(path.resolve()), probability, suggestion))
            if progress and (index % 100 == 0 or index == len(paths)):
                progress(f"CNN reviewer scored {index}/{len(paths)} frame(s)…")
    return scores