"""Indexed, lossless recent-frame cache for CPU/GPU pipeline separation."""

from __future__ import annotations

import copy
import hashlib
import json
import mmap
from pathlib import Path
from typing import Any


FORMAT = "iclr_video.predecoded_recent4_mmap.v1"
FORMAT_V2 = "iclr_video.predecoded_causal_view_mmap.v2"
FORMAT_V3 = "iclr_video.predecoded_mixed_resolution_mmap.v3"


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(8 << 20), b""):
            digest.update(block)
    return digest.hexdigest()


class PredecodedRecent4Cache:
    """Read exact RGB frames from one mmap-backed, benchmark-wide cache."""

    def __init__(
        self,
        root: Path,
        *,
        benchmark: str,
        samples_sha256: str,
        annotation_sha256: str,
        wanted: list[str],
        frame_policy: str = "recent4",
        visual_input_mode: str = "image-list",
    ) -> None:
        root = root.resolve()
        marker = root / "DATA_READY"
        index_path = root / "index.json"
        if not marker.is_file() or not index_path.is_file():
            raise ValueError(f"predecoded cache is not DATA_READY: {root}")
        ready = json.loads(marker.read_text(encoding="utf-8"))
        index = json.loads(index_path.read_text(encoding="utf-8"))
        if ready.get("format") not in {FORMAT, FORMAT_V2, FORMAT_V3} or index.get(
            "format"
        ) != ready.get("format"):
            raise ValueError("predecoded cache format mismatch")
        if ready.get("status") != "complete":
            raise ValueError("predecoded cache marker is not complete")
        if sha256_file(index_path) != ready.get("index_sha256"):
            raise ValueError("predecoded cache index checksum mismatch")
        for key, expected in (
            ("benchmark", benchmark),
            ("samples_sha256", samples_sha256),
            ("annotation_sha256", annotation_sha256),
            ("frame_policy", frame_policy),
            ("visual_input_mode", visual_input_mode),
        ):
            if index.get(key) != expected:
                raise ValueError(f"predecoded cache {key} mismatch")
        entries = index.get("entries")
        if not isinstance(entries, dict):
            raise ValueError("predecoded cache entries are invalid")
        missing = [uid for uid in wanted if uid not in entries]
        if missing:
            raise ValueError(f"predecoded cache misses requested UIDs: {missing[:20]}")
        frames_path = root / str(index.get("frames_file") or "frames.bin")
        if not frames_path.is_file():
            raise ValueError("predecoded cache frame store is missing")
        if sha256_file(frames_path) != index.get("frames_sha256"):
            raise ValueError("predecoded cache frame checksum mismatch")
        self.root = root
        self.index = index
        self.entries = entries
        self._handle = frames_path.open("rb")
        self._mmap = mmap.mmap(self._handle.fileno(), 0, access=mmap.ACCESS_READ)

    @property
    def identity(self) -> dict[str, Any]:
        return {
            "format": self.index["format"],
            "root": str(self.root),
            "benchmark": self.index["benchmark"],
            "num_entries": int(self.index["num_entries"]),
            "frames_sha256": self.index["frames_sha256"],
        }

    def get(
        self, uid: str, *, video_path: Path, query_time: float
    ) -> tuple[list[Any], dict[str, Any]]:
        import numpy as np
        from PIL import Image

        entry = self.entries.get(uid)
        if not isinstance(entry, dict):
            raise KeyError(uid)
        if Path(str(entry["video"])).name != video_path.name:
            raise ValueError(f"predecoded cache video mismatch for {uid}")
        if float(entry["query_time"]).hex() != float(query_time).hex():
            raise ValueError(f"predecoded cache query time mismatch for {uid}")
        frame_records = entry.get("frames")
        if frame_records is not None:
            if not isinstance(frame_records, list) or not frame_records:
                raise ValueError(f"predecoded mixed-resolution frames are invalid for {uid}")
            frames = []
            for record in frame_records:
                shape = tuple(int(value) for value in record["shape"])
                if len(shape) != 3 or min(shape) <= 0 or shape[-1] != 3:
                    raise ValueError(f"predecoded cache frame shape is invalid for {uid}")
                offset = int(record["offset"])
                nbytes = int(record["nbytes"])
                expected_bytes = int(np.prod(shape))
                if (
                    nbytes != expected_bytes
                    or offset < 0
                    or offset + nbytes > len(self._mmap)
                ):
                    raise ValueError(f"predecoded cache byte range is invalid for {uid}")
                array = np.ndarray(
                    shape, dtype=np.uint8, buffer=self._mmap, offset=offset
                )
                frames.append(Image.fromarray(array.copy(), mode="RGB"))
        else:
            shape = tuple(int(value) for value in entry["shape"])
            if len(shape) != 4 or shape[0] <= 0 or shape[-1] != 3:
                raise ValueError(f"predecoded cache shape is invalid for {uid}")
            offset = int(entry["offset"])
            nbytes = int(entry["nbytes"])
            expected_bytes = int(np.prod(shape))
            if nbytes != expected_bytes or offset < 0 or offset + nbytes > len(self._mmap):
                raise ValueError(f"predecoded cache byte range is invalid for {uid}")
            array = np.ndarray(shape, dtype=np.uint8, buffer=self._mmap, offset=offset)
            frames = [Image.fromarray(frame, mode="RGB") for frame in array]
        sampling = copy.deepcopy(entry["sampling"])
        sampling["source_decode_backend"] = sampling.get("decode_backend")
        sampling["decode_backend"] = (
            "predecoded_recent4_mmap_v1"
            if self.index["frame_policy"] == "recent4"
            else f"predecoded_{self.index['frame_policy']}_mmap_v2"
        )
        sampling["predecoded_cache_hit"] = True
        return frames, sampling

    def close(self) -> None:
        self._mmap.close()
        self._handle.close()