#!/usr/bin/env python
"""Versioned, integrity-checked cache bundles for hybrid camera generation.

The production hybrid generator has several expensive preparation stages that
do not depend on the random generation seed once the training policy and target
evaluation group are fixed.  This module stores those results as an immutable,
content-addressed directory that can be published atomically and loaded before
any pickle-backed payload is deserialized.

Cache bundles are trusted build artifacts, not a format for accepting files
from untrusted users.  SHA-256 validation detects corruption and stale/mixed
members; it is not a substitute for an artifact signature.
"""

from __future__ import annotations

import hashlib
import json
import math
import os
import shutil
import tempfile
import time
import uuid
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Iterable, Iterator, Mapping

import numpy as np

import camera_runtime_data as runtime_data


CACHE_FORMAT = "cw-ai-hybrid-preparation"
CACHE_SCHEMA_VERSION = 1
CACHE_BUILDER_VERSION = "2"

MANIFEST_FILE = "manifest.json"
MANIFEST_HASH_FILE = "manifest.sha256"
TEMPLATE_FILE = "templates.joblib"
MODEL_FILE = "models.joblib"
BEAT_ONSET_FILE = "beat_onset.npz"

_MEMBER_NAMES = (TEMPLATE_FILE, MODEL_FILE, BEAT_ONSET_FILE)


class CacheError(RuntimeError):
    """Base exception raised for invalid preparation-cache artifacts."""


class CacheIntegrityError(CacheError):
    """Raised before deserialization when a cache bundle fails validation."""


@dataclass(frozen=True)
class CacheBundleInfo:
    """Result of publishing or reusing one immutable cache bundle."""

    path: Path
    cache_key: str
    reused: bool
    manifest: dict[str, Any]


@dataclass(frozen=True)
class LoadedHybridPreparationCache:
    """Validated payloads needed by the hybrid generation preparation stage."""

    path: Path
    manifest: dict[str, Any]
    templates: Any
    models: Any
    beat_onset: BeatOnsetCache


@dataclass(frozen=True)
class BeatOnsetCache:
    """Target-song audio events reused by cut-boundary selection."""

    frame_count: int
    onset_normalized: np.ndarray
    beat_frames: np.ndarray
    peak_frames: np.ndarray
    candidates: np.ndarray
    audio_start: float
    timeline_offset: int

    def __post_init__(self) -> None:
        frame_count = int(self.frame_count)
        onset = np.asarray(self.onset_normalized, dtype=np.float64)
        beat_frames = np.asarray(self.beat_frames, dtype=np.int64)
        peak_frames = np.asarray(self.peak_frames, dtype=np.int64)
        candidates = np.asarray(self.candidates, dtype=np.int64)
        audio_start = float(self.audio_start)
        timeline_offset = int(self.timeline_offset)

        if frame_count < 1:
            raise ValueError("frame_count must be positive.")
        if onset.shape != (frame_count,):
            raise ValueError(
                "onset_normalized must have shape "
                f"({frame_count},), got {onset.shape}."
            )
        if not np.all(np.isfinite(onset)):
            raise ValueError("onset_normalized must contain finite values.")
        for name, values in (
            ("beat_frames", beat_frames),
            ("peak_frames", peak_frames),
            ("candidates", candidates),
        ):
            if values.ndim != 1:
                raise ValueError(f"{name} must be one-dimensional.")
            if len(values) and np.any(np.diff(values) < 0):
                raise ValueError(f"{name} must be sorted.")
        expected_candidates = np.unique(
            np.concatenate(
                (
                    beat_frames[
                        (beat_frames > 0) & (beat_frames < frame_count)
                    ],
                    peak_frames[
                        (peak_frames > 0) & (peak_frames < frame_count)
                    ],
                )
            )
        )
        if not np.array_equal(candidates, expected_candidates):
            raise ValueError(
                "candidates must be the sorted, in-range union of beat_frames "
                "and peak_frames."
            )
        if not math.isfinite(audio_start):
            raise ValueError("audio_start must be finite.")

        object.__setattr__(self, "frame_count", frame_count)
        object.__setattr__(self, "onset_normalized", onset.copy())
        object.__setattr__(self, "beat_frames", beat_frames.copy())
        object.__setattr__(self, "peak_frames", peak_frames.copy())
        object.__setattr__(self, "candidates", candidates.copy())
        object.__setattr__(self, "audio_start", audio_start)
        object.__setattr__(self, "timeline_offset", timeline_offset)

    def cut_boundaries(
        self,
        minimum_seconds: float,
        maximum_seconds: float,
        sample_rate: int = 60,
    ) -> list[int]:
        """Reproduce the hybrid generator's beat-aware cut selection."""

        minimum_seconds = float(minimum_seconds)
        maximum_seconds = float(maximum_seconds)
        sample_rate = int(sample_rate)
        if (
            sample_rate < 1
            or not math.isfinite(minimum_seconds)
            or not math.isfinite(maximum_seconds)
            or minimum_seconds <= 0.0
            or maximum_seconds <= minimum_seconds
        ):
            raise ValueError("Invalid cut duration or sample rate.")

        minimum_frames = max(
            sample_rate,
            int(round(minimum_seconds * sample_rate)),
        )
        maximum_frames = max(
            minimum_frames + 1,
            int(round(maximum_seconds * sample_rate)),
        )
        boundaries = [0]
        while self.frame_count - boundaries[-1] > maximum_frames:
            current = boundaries[-1]
            energy_window = self.onset_normalized[
                current : min(self.frame_count, current + maximum_frames)
            ]
            energy = float(np.quantile(energy_window, 0.85))
            desired_seconds = float(
                np.clip(3.8 - 0.45 * energy, 2.9, 4.7)
            )
            desired_frame = current + int(
                round(desired_seconds * sample_rate)
            )
            lower = current + minimum_frames
            upper = min(
                self.frame_count - minimum_frames,
                current + maximum_frames,
            )
            valid = self.candidates[
                (self.candidates >= lower) & (self.candidates <= upper)
            ]
            if len(valid) == 0:
                next_boundary = min(upper, desired_frame)
            else:
                distance_penalty = (
                    np.abs(valid - desired_frame) / maximum_frames
                )
                scores = (
                    self.onset_normalized[valid]
                    - 1.15 * distance_penalty
                )
                next_boundary = int(valid[int(np.argmax(scores))])
            if next_boundary <= current:
                next_boundary = min(
                    self.frame_count,
                    current + minimum_frames,
                )
            boundaries.append(next_boundary)

        if (
            self.frame_count - boundaries[-1] < minimum_frames
            and len(boundaries) > 1
        ):
            boundaries[-1] = self.frame_count
        else:
            boundaries.append(self.frame_count)
        return boundaries


class StageTimer:
    """Accumulate named wall-clock stage durations for build reports."""

    def __init__(self, clock: Callable[[], float] = time.perf_counter) -> None:
        self._clock = clock
        self._started = float(clock())
        self._seconds: dict[str, float] = {}
        self._active: str | None = None

    @contextmanager
    def stage(self, name: str) -> Iterator[None]:
        normalized = str(name).strip()
        if not normalized:
            raise ValueError("Stage name must not be empty.")
        if self._active is not None:
            raise RuntimeError(
                f"Cannot start stage {normalized!r} while {self._active!r} "
                "is active."
            )
        self._active = normalized
        started = float(self._clock())
        try:
            yield
        finally:
            elapsed = max(0.0, float(self._clock()) - started)
            self._seconds[normalized] = (
                self._seconds.get(normalized, 0.0) + elapsed
            )
            self._active = None

    def snapshot(self) -> dict[str, Any]:
        if self._active is not None:
            raise RuntimeError(
                f"Cannot snapshot while stage {self._active!r} is active."
            )
        return {
            "stageSeconds": dict(self._seconds),
            "totalSeconds": max(0.0, float(self._clock()) - self._started),
        }


def canonical_json_bytes(value: Any) -> bytes:
    """Encode cache identity data deterministically."""

    return json.dumps(
        value,
        ensure_ascii=False,
        allow_nan=False,
        sort_keys=True,
        separators=(",", ":"),
    ).encode("utf-8")


def sha256_bytes(value: bytes) -> str:
    return hashlib.sha256(value).hexdigest()


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        while True:
            chunk = stream.read(1024 * 1024)
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest()


def cache_key_for(build_inputs: Mapping[str, Any]) -> str:
    """Return the content address for one preparation configuration."""

    envelope = {
        "format": CACHE_FORMAT,
        "schemaVersion": CACHE_SCHEMA_VERSION,
        "builderVersion": CACHE_BUILDER_VERSION,
        "buildInputs": dict(build_inputs),
    }
    return sha256_bytes(canonical_json_bytes(envelope))


def build_inputs_for_records(
    repository_root: Path,
    camera_director_root: Path,
    training_index_path: Path,
    target_record: Any,
    training_records: Iterable[Any],
    planner_records: Iterable[Any],
    entries: Mapping[str, Mapping[str, Any]],
    library_split: str,
    preparation_model_seed: int,
    beam_width: int,
    versions: Mapping[str, Any],
) -> dict[str, Any]:
    """Build the shared cache identity without importing the generator.

    Keeping this helper in the lightweight cache module lets a future
    ``generate_hybrid.py`` integration calculate exactly the same key without
    importing the builder, which itself imports the generator.
    """

    repository_root = Path(repository_root).resolve()
    camera_director_root = Path(camera_director_root).resolve()
    training_index_path = Path(training_index_path).resolve()
    source_names = (
        "camera_runtime_data.py",
        "planner.py",
        "shot_features.py",
        "build_hybrid_preparation_cache.py",
        "hybrid_preparation_cache.py",
    )
    source_files = [camera_director_root / name for name in source_names]
    missing = [path for path in source_files if not path.is_file()]
    if missing:
        raise FileNotFoundError(
            "Preparation cache identity source is missing: "
            + ", ".join(str(path) for path in missing)
        )

    def policy_identity(record: Any) -> dict[str, Any]:
        entry = entries[str(record.id)]
        identity = {
            "songId": str(record.id),
            "datasetId": str(record.dataset_id),
            "contentFingerprint": entry.get("contentFingerprint"),
            "audioFingerprint": entry.get("audioFingerprint"),
            "evaluationGroupId": entry.get("evaluationGroupId"),
            "recommendedSplit": entry.get("recommendedSplit"),
            "role": entry.get("role"),
            "usage": entry.get("usage", {}),
            "aspectRatio": entry.get("aspectRatio"),
            "aspectSource": entry.get("aspectSource"),
        }
        if "audioAlignedFrameMaskRanges" in entry:
            frame_count = entry.get("frameCount")
            ranges = runtime_data.normalize_audio_aligned_frame_mask_ranges(
                entry["audioAlignedFrameMaskRanges"],
                frame_count,
                label=f"Training index audio mask for {record.id}",
            )
            identity["frameCount"] = frame_count
            identity["audioAlignedFrameMaskRanges"] = [
                list(frame_range) for frame_range in ranges
            ]
        return identity

    def target_identity(record: Any) -> dict[str, Any]:
        if bool(getattr(record, "is_generation_input", False)):
            return runtime_data.generation_input_content_identity(record)
        entry = entries.get(str(record.id))
        if entry is None:
            raise KeyError(
                f"Unindexed target {record.id} is not a generation input."
            )
        return policy_identity(record)

    return {
        "pipeline": "hierarchical_hybrid_retarget_v7",
        "plannerMode": "hierarchical",
        "shotLabelMode": "semantic",
        "librarySplit": str(library_split),
        "preparationModelSeed": int(preparation_model_seed),
        "beamWidth": int(beam_width),
        "target": target_identity(target_record),
        "trainingRecords": [
            policy_identity(record)
            for record in sorted(training_records, key=lambda item: item.id)
        ],
        "plannerRecordIds": sorted(
            str(record.id) for record in planner_records
        ),
        "trainingIndexSha256": sha256_file(training_index_path),
        "codeSha256": {
            path.relative_to(repository_root).as_posix(): sha256_file(path)
            for path in source_files
        },
        "versions": dict(versions),
    }


def report_summary_for_target(
    cached_summary: Mapping[str, Any] | None,
    target_record: Any,
) -> dict[str, Any]:
    """Overlay request provenance without mutating cached computation data."""

    summary = dict(cached_summary or {})
    summary.update(
        {
            "targetSongId": str(target_record.id),
            "targetDatasetId": str(target_record.dataset_id),
            "targetSourceManifest": str(target_record.manifest_path),
        }
    )
    return summary


def atomic_write_bytes(path: Path, value: bytes) -> None:
    """Replace one file atomically using a temporary sibling."""

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    descriptor, temporary_name = tempfile.mkstemp(
        prefix=f".{path.name}.",
        suffix=".tmp",
        dir=path.parent,
    )
    temporary_path = Path(temporary_name)
    try:
        with os.fdopen(descriptor, "wb") as stream:
            stream.write(value)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary_path, path)
    except BaseException:
        try:
            os.close(descriptor)
        except OSError:
            pass
        temporary_path.unlink(missing_ok=True)
        raise


def atomic_write_json(path: Path, value: Any) -> None:
    atomic_write_bytes(
        path,
        json.dumps(
            value,
            ensure_ascii=False,
            allow_nan=False,
            sort_keys=True,
            indent=2,
        ).encode("utf-8")
        + b"\n",
    )


def publish_cache_bundle(
    cache_root: Path,
    build_inputs: Mapping[str, Any],
    templates: Any,
    models: Any,
    beat_onset: BeatOnsetCache,
    summary: Mapping[str, Any] | None = None,
) -> CacheBundleInfo:
    """Publish a complete immutable bundle through one atomic rename."""

    cache_root = Path(cache_root).resolve()
    cache_root.mkdir(parents=True, exist_ok=True)
    build_inputs_copy = json.loads(
        canonical_json_bytes(dict(build_inputs)).decode("utf-8")
    )
    cache_key = cache_key_for(build_inputs_copy)
    final_path = cache_root / cache_key
    if final_path.exists():
        try:
            manifest = validate_cache_bundle(final_path, cache_key)
        except (CacheError, OSError, ValueError):
            quarantined = quarantine_cache_bundle(final_path, cache_key)
            if quarantined is not None:
                try:
                    # If another process published a valid bundle between our
                    # validation and rename, restore/reuse it rather than
                    # quarantining a healthy concurrent result.
                    manifest = validate_cache_bundle(quarantined, cache_key)
                except (CacheError, OSError, ValueError):
                    pass
                else:
                    try:
                        os.replace(quarantined, final_path)
                    except OSError:
                        manifest = validate_cache_bundle(final_path, cache_key)
                    return CacheBundleInfo(
                        path=final_path,
                        cache_key=cache_key,
                        reused=True,
                        manifest=manifest,
                    )
        else:
            return CacheBundleInfo(
                path=final_path,
                cache_key=cache_key,
                reused=True,
                manifest=manifest,
            )

    temporary_path = Path(
        tempfile.mkdtemp(
            prefix=f".{cache_key[:16]}.",
            suffix=".tmp",
            dir=cache_root,
        )
    )
    published = False
    try:
        _dump_joblib(temporary_path / TEMPLATE_FILE, templates)
        _dump_joblib(temporary_path / MODEL_FILE, models)
        _write_beat_onset(temporary_path / BEAT_ONSET_FILE, beat_onset)
        members = {
            name: {
                "bytes": (temporary_path / name).stat().st_size,
                "sha256": sha256_file(temporary_path / name),
            }
            for name in _MEMBER_NAMES
        }
        manifest = {
            "format": CACHE_FORMAT,
            "schemaVersion": CACHE_SCHEMA_VERSION,
            "builderVersion": CACHE_BUILDER_VERSION,
            "cacheKey": cache_key,
            "createdUtc": datetime.now(timezone.utc).isoformat(),
            "buildInputs": build_inputs_copy,
            "summary": dict(summary or {}),
            "members": members,
        }
        manifest_path = temporary_path / MANIFEST_FILE
        atomic_write_json(manifest_path, manifest)
        manifest_hash = sha256_file(manifest_path)
        atomic_write_bytes(
            temporary_path / MANIFEST_HASH_FILE,
            f"{manifest_hash}  {MANIFEST_FILE}\n".encode("ascii"),
        )
        validate_cache_bundle(temporary_path, cache_key)

        try:
            os.replace(temporary_path, final_path)
            published = True
        except OSError:
            if not final_path.is_dir():
                raise
            existing = validate_cache_bundle(final_path, cache_key)
            return CacheBundleInfo(
                path=final_path,
                cache_key=cache_key,
                reused=True,
                manifest=existing,
            )
        return CacheBundleInfo(
            path=final_path,
            cache_key=cache_key,
            reused=False,
            manifest=manifest,
        )
    finally:
        if not published and temporary_path.exists():
            shutil.rmtree(temporary_path)


def validate_cache_bundle(
    bundle_path: Path,
    expected_cache_key: str | None = None,
) -> dict[str, Any]:
    """Validate manifest identity and every payload hash without unpickling."""

    bundle_path = Path(bundle_path).resolve()
    if not bundle_path.is_dir():
        raise CacheIntegrityError(
            f"Preparation cache directory was not found: {bundle_path}"
        )
    manifest_path = bundle_path / MANIFEST_FILE
    manifest_hash_path = bundle_path / MANIFEST_HASH_FILE
    if not manifest_path.is_file() or not manifest_hash_path.is_file():
        raise CacheIntegrityError(
            f"Preparation cache manifest is incomplete: {bundle_path}"
        )

    expected_manifest_hash = _read_manifest_hash(manifest_hash_path)
    actual_manifest_hash = sha256_file(manifest_path)
    if actual_manifest_hash != expected_manifest_hash:
        raise CacheIntegrityError(
            "Preparation cache manifest SHA-256 mismatch."
        )
    try:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exception:
        raise CacheIntegrityError(
            "Preparation cache manifest is not valid JSON."
        ) from exception
    if manifest.get("format") != CACHE_FORMAT:
        raise CacheIntegrityError("Unexpected preparation cache format.")
    if manifest.get("schemaVersion") != CACHE_SCHEMA_VERSION:
        raise CacheIntegrityError(
            "Unsupported preparation cache schema version."
        )
    if manifest.get("builderVersion") != CACHE_BUILDER_VERSION:
        raise CacheIntegrityError(
            "Unsupported preparation cache builder version."
        )
    build_inputs = manifest.get("buildInputs")
    if not isinstance(build_inputs, dict):
        raise CacheIntegrityError(
            "Preparation cache buildInputs must be an object."
        )
    calculated_key = cache_key_for(build_inputs)
    cache_key = manifest.get("cacheKey")
    if cache_key != calculated_key:
        raise CacheIntegrityError(
            "Preparation cache key does not match its build inputs."
        )
    if expected_cache_key is not None and cache_key != expected_cache_key:
        raise CacheIntegrityError(
            "Preparation cache key does not match the requested cache key."
        )

    members = manifest.get("members")
    if not isinstance(members, dict) or set(members) != set(_MEMBER_NAMES):
        raise CacheIntegrityError(
            "Preparation cache member list is incomplete."
        )
    for name in _MEMBER_NAMES:
        member = members.get(name)
        if not isinstance(member, dict):
            raise CacheIntegrityError(
                f"Preparation cache member metadata is invalid: {name}"
            )
        path = bundle_path / name
        if not path.is_file():
            raise CacheIntegrityError(
                f"Preparation cache member is missing: {name}"
            )
        if path.stat().st_size != member.get("bytes"):
            raise CacheIntegrityError(
                f"Preparation cache member size mismatch: {name}"
            )
        if sha256_file(path) != member.get("sha256"):
            raise CacheIntegrityError(
                f"Preparation cache member SHA-256 mismatch: {name}"
            )
    return manifest


def load_cache_bundle(
    bundle_path: Path,
    expected_cache_key: str | None = None,
) -> LoadedHybridPreparationCache:
    """Validate all hashes, then deserialize the trusted cache payloads."""

    bundle_path = Path(bundle_path).resolve()
    manifest = validate_cache_bundle(bundle_path, expected_cache_key)
    templates = _load_joblib(bundle_path / TEMPLATE_FILE)
    models = _load_joblib(bundle_path / MODEL_FILE)
    beat_onset = _read_beat_onset(bundle_path / BEAT_ONSET_FILE)
    return LoadedHybridPreparationCache(
        path=bundle_path,
        manifest=manifest,
        templates=templates,
        models=models,
        beat_onset=beat_onset,
    )


def _dump_joblib(path: Path, value: Any) -> None:
    import joblib

    joblib.dump(value, path, compress=0, protocol=5)
    _flush_file(path)


def _load_joblib(path: Path) -> Any:
    import joblib

    return joblib.load(path)


def _write_beat_onset(path: Path, value: BeatOnsetCache) -> None:
    with path.open("wb") as stream:
        np.savez(
            stream,
            frame_count=np.asarray(value.frame_count, dtype=np.int64),
            onset_normalized=value.onset_normalized,
            beat_frames=value.beat_frames,
            peak_frames=value.peak_frames,
            candidates=value.candidates,
            audio_start=np.asarray(value.audio_start, dtype=np.float64),
            timeline_offset=np.asarray(
                value.timeline_offset,
                dtype=np.int64,
            ),
        )
        stream.flush()
        os.fsync(stream.fileno())


def _read_beat_onset(path: Path) -> BeatOnsetCache:
    try:
        with np.load(path, allow_pickle=False) as data:
            return BeatOnsetCache(
                frame_count=int(data["frame_count"].item()),
                onset_normalized=data["onset_normalized"],
                beat_frames=data["beat_frames"],
                peak_frames=data["peak_frames"],
                candidates=data["candidates"],
                audio_start=float(data["audio_start"].item()),
                timeline_offset=int(data["timeline_offset"].item()),
            )
    except (OSError, KeyError, ValueError) as exception:
        raise CacheIntegrityError(
            "Preparation cache beat/onset payload is invalid."
        ) from exception


def _flush_file(path: Path) -> None:
    # Windows requires a writable descriptor for fsync even when no additional
    # bytes are written after joblib closes its output stream.
    with path.open("rb+") as stream:
        os.fsync(stream.fileno())


def _read_manifest_hash(path: Path) -> str:
    try:
        value = path.read_text(encoding="ascii").strip().split()[0]
    except (OSError, IndexError, UnicodeDecodeError) as exception:
        raise CacheIntegrityError(
            "Preparation cache manifest hash is invalid."
        ) from exception
    if len(value) != 64 or any(
        character not in "0123456789abcdef" for character in value
    ):
        raise CacheIntegrityError(
            "Preparation cache manifest hash is invalid."
        )
    return value


def quarantine_cache_bundle(
    bundle_path: Path,
    expected_cache_key: str,
    *,
    expected_manifest_fingerprint: str | None = None,
) -> Path | None:
    """Preserve an invalid bundle while freeing its content-addressed path."""

    bundle_path = Path(bundle_path).absolute()
    cache_key = str(expected_cache_key)
    if bundle_path.name != cache_key:
        raise ValueError(
            "Preparation cache quarantine path does not match its cache key."
        )
    if not bundle_path.exists():
        return None
    if expected_manifest_fingerprint is not None and (
        _manifest_fingerprint(bundle_path) != expected_manifest_fingerprint
    ):
        return None
    quarantine_path = bundle_path.parent / (
        f".invalid-{cache_key[:16]}-{uuid.uuid4().hex}"
    )
    try:
        os.replace(bundle_path, quarantine_path)
    except FileNotFoundError:
        return None
    if expected_manifest_fingerprint is not None and (
        _manifest_fingerprint(quarantine_path)
        != expected_manifest_fingerprint
    ):
        try:
            validate_cache_bundle(quarantine_path, cache_key)
        except (CacheError, OSError, ValueError):
            return quarantine_path
        try:
            os.replace(quarantine_path, bundle_path)
        except OSError:
            validate_cache_bundle(bundle_path, cache_key)
        return None
    return quarantine_path


def _manifest_fingerprint(bundle_path: Path) -> str | None:
    try:
        manifest = json.loads(
            (Path(bundle_path) / MANIFEST_FILE).read_text(encoding="utf-8")
        )
        return sha256_bytes(canonical_json_bytes(manifest))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError, TypeError, ValueError):
        return None
