"""Leakage-aware authored cut priors for production hybrid generation.

The learned candidate ranker is deliberately trained only on the training
split.  At production time an unlabelled Timeline may nevertheless use every
approved reference selected by ``generate_hybrid``.  This module extracts two
small, auditable priors from that selected library without loading camera
trajectories:

* duplicate-aware authored shot durations; and
* exact-audio cut patterns aligned through each Timeline's ``audioStart``.

Evaluation remains leakage-safe because callers pass only the records selected
for the requested library split and target evaluation group.
"""

from __future__ import annotations

import hashlib
import math
import re
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np

import camera_runtime_data as train


SAMPLE_RATE = 60
_SHA256 = re.compile(r"^[0-9a-f]{64}$")
_CHARACTER_ID = re.compile(r"@[0-9A-Za-z_-]+")
_AUDIO_PATTERN_TICKS_PER_FRAME = 1_000_000


@dataclass(frozen=True)
class SameAudioCutPattern:
    """One unique authored edit aligned onto the target Timeline."""

    boundaries: tuple[int, ...]
    weight: float
    same_character: bool
    source_record_ids: tuple[str, ...]


@dataclass(frozen=True)
class CutReferenceCorpus:
    """Compact cut-only context passed to the global boundary planner."""

    target_audio_fingerprint: str
    duration_seconds: tuple[float, ...]
    duration_weights: tuple[float, ...]
    same_audio_patterns: tuple[SameAudioCutPattern, ...]
    eligible_record_count: int
    unique_audio_pattern_count: int
    masked_record_count: int
    retained_authored_shot_count: int
    censored_authored_shot_count: int

    def as_metadata(self) -> dict[str, Any]:
        weights = np.asarray(self.duration_weights, dtype=np.float64)
        durations = np.asarray(self.duration_seconds, dtype=np.float64)
        weighted_median = None
        if len(durations):
            order = np.argsort(durations, kind="stable")
            cumulative = np.cumsum(weights[order])
            weighted_median = float(
                durations[order][
                    np.searchsorted(cumulative, cumulative[-1] * 0.5)
                ]
            )
        return {
            "targetAudioFingerprint": self.target_audio_fingerprint,
            "eligibleRecordCount": self.eligible_record_count,
            "uniqueAudioPatternCount": self.unique_audio_pattern_count,
            "durationSampleCount": len(self.duration_seconds),
            "durationWeightedMedianSeconds": weighted_median,
            "sameAudioPatternCount": len(self.same_audio_patterns),
            "sameAudioSourceRecordIds": sorted(
                {
                    record_id
                    for pattern in self.same_audio_patterns
                    for record_id in pattern.source_record_ids
                }
            ),
            "sameCharacterPatternCount": sum(
                pattern.same_character for pattern in self.same_audio_patterns
            ),
            "audioAlignedMaskedRecordCount": self.masked_record_count,
            "audioAlignedRetainedAuthoredShotCount": (
                self.retained_authored_shot_count
            ),
            "audioAlignedCensoredAuthoredShotCount": (
                self.censored_authored_shot_count
            ),
        }

    def weighted_duration_frame_samples(
        self,
        sample_count: int = 512,
    ) -> np.ndarray:
        """Materialize a deterministic weighted prior for simple planner APIs."""

        sample_count = int(sample_count)
        if sample_count < 1:
            raise ValueError("sample_count must be positive.")
        values = np.asarray(self.duration_seconds, dtype=np.float64)
        weights = np.asarray(self.duration_weights, dtype=np.float64)
        if len(values) == 0 or values.shape != weights.shape:
            raise ValueError("The cut reference corpus has no duration prior.")
        order = np.argsort(values, kind="stable")
        values = values[order]
        cumulative = np.cumsum(weights[order])
        cumulative /= cumulative[-1]
        probabilities = (
            np.arange(sample_count, dtype=np.float64) + 0.5
        ) / sample_count
        indices = np.searchsorted(cumulative, probabilities, side="left")
        return np.maximum(
            1,
            np.rint(values[np.minimum(indices, len(values) - 1)] * SAMPLE_RATE),
        ).astype(np.int64)


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


def audio_fingerprint(record: train.SongRecord) -> str:
    """Return a validated portable audio hash, computing it when necessary."""

    declared = str(record.song.get("audioFingerprint", "")).strip().lower()
    if declared:
        if not _SHA256.fullmatch(declared):
            raise ValueError(
                f"Invalid audioFingerprint for cut reference {record.id}."
            )
    if record.audio_path.is_file():
        actual = sha256_file(record.audio_path)
        if declared and actual != declared:
            raise ValueError(
                f"Declared audioFingerprint does not match the materialized "
                f"audio for cut reference {record.id}."
            )
        return actual
    if declared:
        return declared
    # Portable RuntimeData omits raw audio; the catalog fingerprint is the
    # authoritative identity in that mode.
    prepared = record.song.get("preparedFile", "prepared_v1.npz")
    if (record.song_dir / prepared).is_file():
        return hashlib.sha256(
            f"prepared:{record.id}:{prepared}".encode("utf-8")
        ).hexdigest()
    raise FileNotFoundError(f"No audio or declared fingerprint for {record.id}.")


def _verified_masked_audio_fingerprint(
    record: train.SongRecord,
    indexed_fingerprint: str,
) -> str:
    """Bind an explicit frame mask to the materialized source-audio bytes."""

    audio_path = record.audio_path
    if not audio_path.is_file():
        raise ValueError(
            f"Masked cut reference {record.id} audio file is missing: {audio_path}"
        )
    lfs_pointer_prefix = b"version https://git-lfs.github.com/spec/v1"
    with audio_path.open("rb") as stream:
        if stream.read(len(lfs_pointer_prefix)) == lfs_pointer_prefix:
            raise ValueError(
                f"Masked cut reference {record.id} audio is still a Git LFS "
                f"pointer: {audio_path}"
            )
    actual = sha256_file(audio_path)
    if actual != indexed_fingerprint:
        raise ValueError(
            f"Masked cut reference {record.id} audioFingerprint does not match "
            f"the materialized audio file: {audio_path}"
        )
    return actual


def _explicit_audio_mask_ranges(
    record: train.SongRecord,
    entry: Mapping[str, Any],
    indexed_fingerprint: str,
) -> tuple[tuple[int, int], ...] | None:
    """Return a strict explicit mask, preserving legacy full-record behavior."""

    if "audioAlignedFrameMaskRanges" not in entry:
        partial = sorted(
            key
            for key in ("audioAlignedEndFrame", "audioAlignedFrameCount")
            if key in entry
        )
        if partial:
            raise ValueError(
                f"Cut reference audio mask for {record.id} is incomplete: "
                "audioAlignedFrameMaskRanges is missing while "
                + ", ".join(partial)
                + " is present."
            )
        return None

    manifest_frame_count = record.song.get("frameCount")
    indexed_frame_count = entry.get("frameCount")
    if (
        isinstance(manifest_frame_count, bool)
        or not isinstance(manifest_frame_count, int)
        or isinstance(indexed_frame_count, bool)
        or not isinstance(indexed_frame_count, int)
        or manifest_frame_count < 1
        or indexed_frame_count != manifest_frame_count
    ):
        raise ValueError(
            f"Cut reference audio mask frameCount mismatch for {record.id}."
        )
    ranges = train.normalize_audio_aligned_frame_mask_ranges(
        entry["audioAlignedFrameMaskRanges"],
        manifest_frame_count,
        label=f"Cut reference audio mask for {record.id}",
    )
    if not ranges:
        raise ValueError(
            f"Active cut reference {record.id} has an empty audio frame mask."
        )
    expected_count = sum(end - start for start, end in ranges)
    if "audioAlignedFrameCount" in entry:
        aligned_count = entry["audioAlignedFrameCount"]
        if (
            isinstance(aligned_count, bool)
            or not isinstance(aligned_count, int)
            or aligned_count != expected_count
        ):
            raise ValueError(
                f"Cut reference audioAlignedFrameCount is inconsistent for "
                f"{record.id}."
            )
    if "audioAlignedEndFrame" in entry:
        aligned_end = entry["audioAlignedEndFrame"]
        if (
            isinstance(aligned_end, bool)
            or not isinstance(aligned_end, int)
            or aligned_end != ranges[-1][1]
        ):
            raise ValueError(
                f"Cut reference audioAlignedEndFrame is inconsistent for "
                f"{record.id}."
            )
    _verified_masked_audio_fingerprint(record, indexed_fingerprint)
    return ranges


def character_id(record: train.SongRecord) -> str:
    value = str(record.song.get("characterPath", ""))
    match = _CHARACTER_ID.search(value)
    if match:
        return match.group(0).lower()
    leaf = value.replace("\\", "/").rstrip("/").rsplit("/", 1)[-1]
    # Older manifests predate explicit character IDs but consistently start
    # the character object with the performer name (for example 아리사_260715).
    name = leaf.split("_", 1)[0].strip().lower()
    return "name:" + name if name else ""


def authored_boundaries(record: train.SongRecord) -> tuple[int, ...]:
    """Read strict contiguous cut boundaries without touching camera samples."""

    frame_count = int(record.song["frameCount"])
    sample_rate = int(
        record.song.get("sampleRate", record.manifest.get("sampleRate", 0))
    )
    if sample_rate != SAMPLE_RATE:
        raise ValueError(
            f"Cut reference {record.id} is {sample_rate} fps; expected 60 fps."
        )
    shot_index_path = record.song_dir / "shot_index.i32"
    if shot_index_path.is_file():
        values = np.fromfile(shot_index_path, dtype="<i4")
    else:
        prepared_path = record.portable_prepared_path
        if not prepared_path.is_file():
            raise FileNotFoundError(f"No cut source for {record.id}.")
        with np.load(prepared_path, allow_pickle=False) as data:
            cuts = np.asarray(data["cuts"], dtype=np.float32)
        values = np.cumsum(cuts > 0.5, dtype=np.int64)
    if values.shape != (frame_count,):
        raise ValueError(
            f"Cut reference {record.id} shot_index length {len(values)} != "
            f"frameCount {frame_count}."
        )
    if frame_count < 1 or np.any(values < 0):
        raise ValueError(f"Cut reference {record.id} has invalid shot indices.")
    internal = np.flatnonzero(values[1:] != values[:-1]) + 1
    boundaries = (0, *(int(value) for value in internal), frame_count)
    if any(right <= left for left, right in zip(boundaries, boundaries[1:])):
        raise ValueError(f"Cut reference {record.id} has non-positive shots.")
    return boundaries


def _align_pattern(
    boundaries: Sequence[int],
    source: train.SongRecord,
    target: train.SongRecord,
) -> tuple[int, ...]:
    target_frames = int(target.song["frameCount"])
    source_audio_start = float(source.song.get("audioStart", 0.0))
    target_audio_start = float(target.song.get("audioStart", 0.0))
    if not all(
        math.isfinite(value) for value in (source_audio_start, target_audio_start)
    ):
        raise ValueError("audioStart must be finite for exact-audio alignment.")
    mapped = [0]
    for boundary in boundaries[1:-1]:
        audio_seconds = boundary / SAMPLE_RATE - source_audio_start
        target_frame = int(
            round((audio_seconds + target_audio_start) * SAMPLE_RATE)
        )
        if 0 < target_frame < target_frames:
            mapped.append(target_frame)
    mapped.append(target_frames)
    return tuple(sorted(set(mapped)))


def _audio_relative_boundary_signature(
    boundaries: Sequence[int],
    source: train.SongRecord,
) -> tuple[int, ...]:
    """Identify one edit pattern in audio-relative micro-frame ticks.

    Raw Timeline frame boundaries are not a sufficient identity because two
    exports can share those integers while using different ``audioStart``
    offsets.  Grouping those records would union shifted cuts into false
    near-duplicate pairs.  Micro-frame ticks preserve sub-frame audio offsets
    while remaining deterministic JSON/float input material.
    """

    audio_start = float(source.song.get("audioStart", 0.0))
    if not math.isfinite(audio_start):
        raise ValueError("audioStart must be finite for cut-pattern identity.")
    audio_start_ticks = int(
        round(
            audio_start
            * SAMPLE_RATE
            * _AUDIO_PATTERN_TICKS_PER_FRAME
        )
    )
    return tuple(
        int(boundary) * _AUDIO_PATTERN_TICKS_PER_FRAME - audio_start_ticks
        for boundary in boundaries
    )


def build_cut_reference_corpus(
    target: train.SongRecord,
    records: Sequence[train.SongRecord],
    index_entries: Mapping[str, Mapping[str, Any]],
) -> CutReferenceCorpus:
    """Build duplicate-safe priors from the already leakage-filtered library."""

    target_fingerprint = audio_fingerprint(target)
    target_character = character_id(target)
    pattern_records: dict[
        tuple[str, tuple[int, ...]],
        list[
            tuple[
                train.SongRecord,
                tuple[tuple[int, int], ...],
                tuple[int, ...],
            ]
        ],
    ] = defaultdict(list)
    eligible_count = 0
    masked_record_count = 0
    retained_authored_shot_count = 0
    censored_authored_shot_count = 0
    for record in records:
        entry = index_entries.get(record.id)
        if not isinstance(entry, Mapping):
            raise ValueError(f"Missing cut policy entry for {record.id}.")
        usage = entry.get("usage")
        if (
            entry.get("role") != "primary"
            or not isinstance(usage, Mapping)
            or usage.get("cutModel") is not True
        ):
            continue
        fingerprint = str(entry.get("audioFingerprint", "")).strip().lower()
        if not _SHA256.fullmatch(fingerprint):
            raise ValueError(f"Invalid indexed audio fingerprint for {record.id}.")
        boundaries = authored_boundaries(record)
        mask_ranges = _explicit_audio_mask_ranges(record, entry, fingerprint)
        if mask_ranges is None:
            retained_spans = tuple(zip(boundaries, boundaries[1:]))
            retained_cuts = boundaries[1:-1]
        else:
            masked_record_count += 1
            retained_spans = tuple(
                (left, right)
                for left, right in zip(boundaries, boundaries[1:])
                if train.audio_aligned_frame_mask_range_index(
                    left, right, mask_ranges
                )
                is not None
            )
            retained_cuts = tuple(
                boundary
                for boundary in boundaries[1:-1]
                if any(start < boundary < end for start, end in mask_ranges)
            )
        retained_authored_shot_count += len(retained_spans)
        censored_authored_shot_count += (
            len(boundaries) - 1 - len(retained_spans)
        )
        pattern_signature = _audio_relative_boundary_signature(boundaries, record)
        pattern_records[(fingerprint, pattern_signature)].append(
            (record, retained_spans, retained_cuts)
        )
        eligible_count += 1

    durations: list[float] = []
    duration_weights: list[float] = []
    same_audio_patterns: list[SameAudioCutPattern] = []
    for (fingerprint, _pattern_signature), sources in sorted(
        pattern_records.items(), key=lambda value: (value[0][0], value[0][1])
    ):
        retained_spans = sorted(
            {
                span
                for _, source_spans, _ in sources
                for span in source_spans
            }
        )
        if retained_spans:
            shot_durations = np.asarray(
                [right - left for left, right in retained_spans],
                dtype=np.float64,
            ) / SAMPLE_RATE
            # One unique audio/cut pattern contributes one unit in total,
            # regardless of how many duplicated exports or shots it contains.
            per_shot_weight = 1.0 / len(shot_durations)
            durations.extend(float(value) for value in shot_durations)
            duration_weights.extend(per_shot_weight for _ in shot_durations)
        if fingerprint != target_fingerprint:
            continue
        aligned_internal = {
            boundary
            for source, _, source_cuts in sources
            for boundary in _align_pattern(
                (0, *source_cuts, int(source.song["frameCount"])),
                source,
                target,
            )[1:-1]
        }
        if not aligned_internal:
            continue
        aligned = (0, *sorted(aligned_internal), int(target.song["frameCount"]))
        same_character = bool(
            target_character
            and any(
                character_id(source) == target_character
                for source, _, _ in sources
            )
        )
        same_audio_patterns.append(
            SameAudioCutPattern(
                boundaries=aligned,
                weight=2.0 if same_character else 1.0,
                same_character=same_character,
                source_record_ids=tuple(
                    sorted(source.id for source, _, _ in sources)
                ),
            )
        )

    if duration_weights:
        total_weight = float(sum(duration_weights))
        duration_weights = [value / total_weight for value in duration_weights]
    same_audio_patterns.sort(
        key=lambda value: (-value.weight, value.boundaries, value.source_record_ids)
    )
    return CutReferenceCorpus(
        target_audio_fingerprint=target_fingerprint,
        duration_seconds=tuple(durations),
        duration_weights=tuple(duration_weights),
        same_audio_patterns=tuple(same_audio_patterns),
        eligible_record_count=eligible_count,
        unique_audio_pattern_count=len(pattern_records),
        masked_record_count=masked_record_count,
        retained_authored_shot_count=retained_authored_shot_count,
        censored_authored_shot_count=censored_authored_shot_count,
    )
