#!/usr/bin/env python
"""Inference-only motion/audio preparation shared by the camera worker.

This module intentionally contains no offline training framework imports.  It
keeps the data, quaternion, manifest, and prepared-song contracts used by the
high-quality generator byte-for-byte compatible with the training pipeline.
"""

from __future__ import annotations

import hashlib
import json
import math
import os
import re
import sys
import tempfile
import uuid
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Iterable, Mapping

import librosa
import numpy as np


PIPELINE_VERSION = 1
AUDIO_RATE = 24_000
AUDIO_HOP = 400  # Exactly 60 audio feature frames per second.
MEL_BANDS = 32
REVISION_DATASET_PATTERN = re.compile(
    r"^(?P<base>.+)_rev_(?P<fingerprint>[0-9a-f]{16})$"
)
GENERATION_INPUT_CACHE_IDENTITY_VERSION = "generation-input-content-v1"
GENERATION_INPUT_CACHE_FORMAT = "cw-ai-generation-input-prepared"
GENERATION_INPUT_CACHE_SCHEMA_VERSION = 1
GENERATION_INPUT_PREPARATION_LOGIC_VERSION = "generation-input-preparation-v2"
AUDIO_ALIGNED_FRAME_MASK_POLICY_VERSION = "audio-aligned-half-open-v1"
GENERATION_INPUT_VOLATILE_MANIFEST_FIELDS = frozenset(
    {
        "audioAssetPath",
        "audioFeaturesFile",
        "audioFeaturesFingerprint",
        "audioFile",
        "audioFingerprint",
        "cameraFile",
        "cameraFingerprint",
        "characterPath",
        "createdUtc",
        "datasetId",
        "folderName",
        "jointsFile",
        "jointsFingerprint",
        "mainTimelineAssetPath",
        "motionTimelineAssetPath",
        "previewFile",
        "rootFile",
        "rootFingerprint",
        "scenePath",
        "shotIndexFile",
        "shotIndexFingerprint",
        "shotsFile",
        "sourceFingerprint",
        "timeFile",
        "timeFingerprint",
        "unityVersion",
    }
)

# A shared semantic skeleton for the Humanoid and Generic rigs in the current data.
CANONICAL_JOINT_ALIASES = {
    "pelvis": ("hips", "pelvis"),
    "spine": ("spine",),
    "chest": ("chest", "upperchest", "spine1"),
    "neck": ("neck",),
    "head": ("head",),
    "left_shoulder": ("leftshoulder", "lclavicle"),
    "left_upper_arm": ("leftupperarm", "lupperarm"),
    "left_lower_arm": ("leftlowerarm", "lforearm"),
    "left_hand": ("lefthand", "lhand"),
    "right_shoulder": ("rightshoulder", "rclavicle"),
    "right_upper_arm": ("rightupperarm", "rupperarm"),
    "right_lower_arm": ("rightlowerarm", "rforearm"),
    "right_hand": ("righthand", "rhand"),
    "left_upper_leg": ("leftupperleg", "lthigh"),
    "left_lower_leg": ("leftlowerleg", "lcalf"),
    "left_foot": ("leftfoot", "lfoot"),
    "left_toes": ("lefttoes", "ltoe0"),
    "right_upper_leg": ("rightupperleg", "rthigh"),
    "right_lower_leg": ("rightlowerleg", "rcalf"),
    "right_foot": ("rightfoot", "rfoot"),
    "right_toes": ("righttoes", "rtoe0"),
}
CANONICAL_JOINTS = tuple(CANONICAL_JOINT_ALIASES)


def configure_utf8_stdio() -> None:
    for stream in (sys.stdout, sys.stderr):
        reconfigure = getattr(stream, "reconfigure", None)
        if reconfigure is not None:
            reconfigure(encoding="utf-8", errors="backslashreplace")


@dataclass
class SongRecord:
    manifest_path: Path
    dataset_dir: Path
    project_root: Path
    manifest: dict
    song: dict

    @property
    def name(self) -> str:
        return self.song["songName"]

    @property
    def id(self) -> str:
        scene_path = str(self.manifest.get("scenePath", ""))
        folder_name = str(self.song["folderName"])
        value = f"{scene_path}\0{self.name}\0{folder_name}"
        if self.is_revision:
            value += f"\0revision:{self.dataset_id}"
        value = value.encode("utf-8")
        return hashlib.sha256(value).hexdigest()[:16]

    @property
    def dataset_id(self) -> str:
        return self.dataset_dir.name

    @property
    def revision_of_dataset_id(self) -> str | None:
        match = REVISION_DATASET_PATTERN.fullmatch(self.dataset_id)
        return match.group("base") if match is not None else None

    @property
    def is_revision(self) -> bool:
        return self.revision_of_dataset_id is not None

    @property
    def song_dir(self) -> Path:
        return self.dataset_dir / self.song["folderName"]

    @property
    def audio_path(self) -> Path:
        portable_audio = self.song.get("audioFile")
        if portable_audio:
            return self.song_dir / portable_audio
        project_root_override = os.environ.get("CWAI_UNITY_PROJECT_ROOT")
        if project_root_override:
            return Path(project_root_override).resolve() / self.song["audioAssetPath"]
        return self.project_root / self.song["audioAssetPath"]

    @property
    def portable_prepared_path(self) -> Path:
        file_name = self.song.get("preparedFile", "prepared_v1.npz")
        return self.song_dir / file_name

    @property
    def is_generation_input(self) -> bool:
        """Whether this record is a camera-free generation target.

        Generation manifests deliberately live outside the authored-camera
        training catalog.  Accept the fields at either manifest or song scope
        so the single-song Unity exporter can keep common metadata at the
        document root while older tooling can still emit per-song metadata.
        """

        input_kind = self.song.get(
            "inputKind",
            self.manifest.get("inputKind"),
        )
        has_authored_camera = self.song.get(
            "hasAuthoredCamera",
            self.manifest.get("hasAuthoredCamera"),
        )
        return (
            str(input_kind or "").strip().lower() == "generation_input"
            and has_authored_camera is False
        )


@dataclass
class PreparedSong:
    name: str
    inputs: np.ndarray
    targets: np.ndarray
    cuts: np.ndarray
    times: np.ndarray
    body_scale: float
    source_manifest: str


def find_manifest_paths(values: Iterable[str]) -> list[Path]:
    paths: list[Path] = []
    requested = [Path(value).resolve() for value in values]
    if not requested:
        requested = [(Path.cwd() / "DatasetExports").resolve()]

    for path in requested:
        if path.is_file() and path.name == "dataset_manifest.json":
            paths.append(path)
        elif (path / "dataset_manifest.json").is_file():
            paths.append(path / "dataset_manifest.json")
        elif path.is_dir():
            paths.extend(path.glob("TimelineCamera_*/dataset_manifest.json"))
            paths.extend(path.glob("GenerationInput_*/dataset_manifest.json"))

    unique = sorted(set(path.resolve() for path in paths))
    if not unique:
        raise FileNotFoundError("No dataset_manifest.json files were found.")
    return unique


def load_records(
    manifest_paths: Iterable[Path],
    skip_unavailable: bool = False,
) -> list[SongRecord]:
    records: list[SongRecord] = []
    for manifest_path in manifest_paths:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        if manifest["sampleRate"] != 60:
            raise ValueError(f"{manifest_path} is not sampled at 60 FPS.")
        dataset_dir = manifest_path.parent
        project_root = dataset_dir.parent.parent
        for song in manifest["songs"]:
            record = SongRecord(
                manifest_path=manifest_path,
                dataset_dir=dataset_dir,
                project_root=project_root,
                manifest=manifest,
                song=song,
            )
            has_audio = record.audio_path.is_file()
            has_prepared = record.portable_prepared_path.is_file()
            unavailable_reason = ""
            if record.is_generation_input:
                required_target_paths = (
                    record.audio_path,
                    record.song_dir / "joints_world.f32",
                    record.song_dir / "root.f32",
                    record.song_dir / "time.f64",
                )
                missing_target_paths = [
                    str(path)
                    for path in required_target_paths
                    if not path.is_file()
                ]
                if missing_target_paths:
                    unavailable_reason = (
                        "Generation input is missing motion/audio sources: "
                        + ", ".join(missing_target_paths)
                    )
                else:
                    try:
                        canonical_joint_indices(manifest["sharedBones"])
                    except (KeyError, ValueError) as error:
                        unavailable_reason = str(error)
            elif not has_audio and not has_prepared:
                unavailable_reason = (
                    "Neither source audio nor portable prepared data was found: "
                    f"{record.audio_path}"
                )
            elif has_audio and not has_prepared:
                try:
                    canonical_joint_indices(manifest["sharedBones"])
                except (KeyError, ValueError) as error:
                    unavailable_reason = str(error)
            if unavailable_reason:
                if skip_unavailable:
                    continue
                raise FileNotFoundError(unavailable_reason)
            records.append(record)
    if len(records) < 2:
        raise ValueError("At least two songs are required for a song-level validation split.")
    return records


def select_record(
    records: Iterable[SongRecord],
    song_name: str = "",
    song_id: str = "",
    dataset_id: str = "",
) -> SongRecord:
    candidates = list(records)
    if song_name:
        candidates = [record for record in candidates if record.name == song_name]
    if song_id:
        candidates = [record for record in candidates if record.id == song_id]
    if dataset_id:
        candidates = [
            record for record in candidates if record.dataset_id == dataset_id
        ]

    selector_parts = []
    if song_name:
        selector_parts.append(f"name={song_name!r}")
    if song_id:
        selector_parts.append(f"id={song_id!r}")
    if dataset_id:
        selector_parts.append(f"dataset={dataset_id!r}")
    selector = ", ".join(selector_parts) or "(no selector)"

    if len(candidates) == 1:
        return candidates[0]
    if not candidates:
        raise ValueError(f"No song matched {selector}.")

    preview = "\n".join(
        (
            f"  {record.id} | {record.name} | "
            f"{record.dataset_id}/{record.song['folderName']}"
        )
        for record in candidates[:20]
    )
    suffix = ""
    if len(candidates) > 20:
        suffix = f"\n  ... and {len(candidates) - 20} more"
    raise ValueError(
        f"Song selector {selector} matched {len(candidates)} records. "
        "Add --song-id or --dataset-id to select one:\n"
        f"{preview}{suffix}"
    )


def normalize_audio_aligned_frame_mask_ranges(
    value: Any,
    frame_count: int,
    *,
    label: str = "audioAlignedFrameMaskRanges",
) -> tuple[tuple[int, int], ...]:
    """Validate canonical half-open frame ranges without mutating source data.

    Ranges must be strictly ordered, non-overlapping, and non-adjacent.  Adjacent
    ranges are rejected because their canonical representation is one merged
    range.  An empty list is valid and deliberately admits no frames.
    """

    if isinstance(frame_count, bool) or not isinstance(frame_count, int):
        raise ValueError(f"{label} frame count must be an integer.")
    if frame_count < 0:
        raise ValueError(f"{label} frame count must be non-negative.")
    if not isinstance(value, list):
        raise ValueError(f"{label} must be an array of [start, end) ranges.")

    normalized: list[tuple[int, int]] = []
    previous_end: int | None = None
    for index, raw_range in enumerate(value):
        range_label = f"{label}[{index}]"
        if not isinstance(raw_range, list) or len(raw_range) != 2:
            raise ValueError(f"{range_label} must be a two-item array.")
        start, end = raw_range
        if (
            isinstance(start, bool)
            or isinstance(end, bool)
            or not isinstance(start, int)
            or not isinstance(end, int)
        ):
            raise ValueError(f"{range_label} bounds must be integers.")
        if not 0 <= start < end <= frame_count:
            raise ValueError(
                f"{range_label} must satisfy 0 <= start < end <= {frame_count}."
            )
        if previous_end is not None and start <= previous_end:
            raise ValueError(
                f"{range_label} overlaps, touches, or is out of order; "
                "canonical ranges must be merged and strictly separated."
            )
        normalized.append((start, end))
        previous_end = end
    return tuple(normalized)


def audio_aligned_frame_mask_range_index(
    start: int,
    end: int,
    ranges: tuple[tuple[int, int], ...],
) -> int | None:
    """Return the single range containing ``[start, end)``, else ``None``.

    A span crossing a range boundary or falling fully outside the mask is
    intentionally indistinguishable to consumers: both must fail closed.
    """

    if (
        isinstance(start, bool)
        or isinstance(end, bool)
        or not isinstance(start, int)
        or not isinstance(end, int)
        or start < 0
        or end <= start
    ):
        raise ValueError("Frame spans must be non-empty half-open integer ranges.")
    for index, (range_start, range_end) in enumerate(ranges):
        if range_start <= start and end <= range_end:
            return index
    return None


def audio_aligned_frame_mask_span_status(
    start: int,
    end: int,
    ranges: tuple[tuple[int, int], ...],
) -> tuple[str, int | None]:
    """Classify a source span as contained, boundary-crossing, or outside."""

    range_index = audio_aligned_frame_mask_range_index(start, end, ranges)
    if range_index is not None:
        return "contained", range_index
    intersects = any(
        start < range_end and end > range_start
        for range_start, range_end in ranges
    )
    return ("boundary_crossing" if intersects else "outside"), None


def filter_records_by_training_index(
    records: Iterable[SongRecord],
    index_path: Path,
    purpose: str,
    validation_ids: set[str],
) -> list[SongRecord]:
    payload = json.loads(index_path.read_text(encoding="utf-8"))
    entries = {
        str(entry["songId"]): entry
        for entry in payload.get("entries", [])
    }
    missing_evaluation_groups = sorted(
        song_id
        for song_id, entry in entries.items()
        if not str(entry.get("evaluationGroupId", "")).strip()
    )
    if missing_evaluation_groups:
        raise ValueError(
            "Training index entries are missing evaluationGroupId: "
            + ", ".join(missing_evaluation_groups[:20])
        )
    group_splits: dict[str, set[str]] = {}
    for entry in entries.values():
        evaluation_group = str(entry["evaluationGroupId"])
        group_splits.setdefault(evaluation_group, set()).add(
            str(entry.get("recommendedSplit", ""))
        )
    conflicting_groups = {
        group: sorted(splits)
        for group, splits in group_splits.items()
        if len(splits) > 1
    }
    if conflicting_groups:
        raise ValueError(
            "Training index has evaluation groups spanning splits: "
            + json.dumps(
                conflicting_groups,
                ensure_ascii=False,
                sort_keys=True,
            )
        )
    ineligible_validation = sorted(
        song_id
        for song_id in validation_ids
        if song_id in entries
        and entries[song_id].get("usage", {}).get(purpose) is not True
    )
    if ineligible_validation:
        raise ValueError(
            f"Validation songs are not eligible for training purpose "
            f"{purpose!r}: "
            + ", ".join(ineligible_validation)
        )
    validation_groups = {
        str(entries[song_id].get("evaluationGroupId", ""))
        for song_id in validation_ids
        if song_id in entries
        and entries[song_id].get("evaluationGroupId")
    }
    selected: list[SongRecord] = []
    missing: list[str] = []
    for record in records:
        entry = entries.get(record.id)
        if entry is None:
            missing.append(record.id)
            continue
        is_validation = record.id in validation_ids
        is_training = (
            bool(entry.get("usage", {}).get(purpose, False))
            and entry.get("recommendedSplit") == "train"
            and str(entry.get("evaluationGroupId", ""))
            not in validation_groups
        )
        if is_validation or is_training:
            selected.append(record)

    if missing:
        raise ValueError(
            f"Training index is missing {len(missing)} available song IDs. "
            "Regenerate it with tools/build_training_index.py."
        )
    selected_ids = {record.id for record in selected}
    absent_validation = validation_ids - selected_ids
    if absent_validation:
        raise ValueError(
            "Validation songs are missing from the training index: "
            + ", ".join(sorted(absent_validation))
        )
    leakage_ids = sorted(
        record.id
        for record in selected
        if record.id not in validation_ids
        and entries[record.id].get("recommendedSplit") == "train"
        and str(entries[record.id].get("evaluationGroupId", ""))
        in validation_groups
    )
    if leakage_ids:
        raise ValueError(
            "Training selection leaks a validation evaluation group through "
            "song IDs: "
            + ", ".join(leakage_ids)
        )
    if len(selected) < 2:
        raise ValueError(
            f"Training index purpose {purpose!r} selected fewer than two songs."
        )
    return selected


def normalized_leaf(path: str) -> str:
    leaf = path.replace("\\", "/").rsplit("/", 1)[-1].lower()
    leaf = leaf.replace("bip001", "")
    return re.sub(r"[^a-z0-9]", "", leaf)


def canonical_joint_indices(names: list[str]) -> list[int]:
    normalized = [normalized_leaf(name) for name in names]
    result: list[int] = []
    missing: list[str] = []
    for canonical_name, aliases in CANONICAL_JOINT_ALIASES.items():
        match = next(
            (index for index, value in enumerate(normalized) if value in aliases),
            None,
        )
        if match is None:
            missing.append(canonical_name)
        else:
            result.append(match)
    if missing:
        raise ValueError(f"Missing canonical joints: {', '.join(missing)}")
    return result


def quat_conjugate(q: np.ndarray) -> np.ndarray:
    result = q.copy()
    result[..., :3] *= -1.0
    return result


def quat_multiply(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    av, aw = a[..., :3], a[..., 3:4]
    bv, bw = b[..., :3], b[..., 3:4]
    xyz = aw * bv + bw * av + np.cross(av, bv)
    w = aw * bw - np.sum(av * bv, axis=-1, keepdims=True)
    return np.concatenate((xyz, w), axis=-1)


def quat_rotate(q: np.ndarray, v: np.ndarray) -> np.ndarray:
    qv = q[..., :3]
    qw = q[..., 3:4]
    return v + 2.0 * (
        qw * np.cross(qv, v) + np.cross(qv, np.cross(qv, v))
    )


def quat_to_rotation_6d(q: np.ndarray) -> np.ndarray:
    q = q / np.maximum(np.linalg.norm(q, axis=-1, keepdims=True), 1e-8)
    x, y, z, w = (q[..., index] for index in range(4))
    first_column = np.stack(
        (
            1.0 - 2.0 * (y * y + z * z),
            2.0 * (x * y + z * w),
            2.0 * (x * z - y * w),
        ),
        axis=-1,
    )
    second_column = np.stack(
        (
            2.0 * (x * y - z * w),
            1.0 - 2.0 * (x * x + z * z),
            2.0 * (y * z + x * w),
        ),
        axis=-1,
    )
    return np.concatenate((first_column, second_column), axis=-1)


def rotation_6d_to_quat(values: np.ndarray) -> np.ndarray:
    """Convert compact rotation columns to normalized XYZW quaternions."""

    values = np.asarray(values, dtype=np.float32)
    if values.ndim < 1 or values.shape[-1] != 6:
        raise ValueError("Rotation 6D values must have a final dimension of six.")
    first = values[..., :3]
    second = values[..., 3:6]
    first = first / np.maximum(np.linalg.norm(first, axis=-1, keepdims=True), 1e-8)
    second = second - np.sum(first * second, axis=-1, keepdims=True) * first
    second = second / np.maximum(np.linalg.norm(second, axis=-1, keepdims=True), 1e-8)
    third = np.cross(first, second)
    matrix = np.stack((first, second, third), axis=-1)
    result = np.empty(values.shape[:-1] + (4,), dtype=np.float32)
    trace = np.trace(matrix, axis1=-2, axis2=-1)
    positive = trace > 0.0
    root = np.sqrt(np.maximum(trace[positive] + 1.0, 1e-8)) * 2.0
    result[positive, 3] = 0.25 * root
    result[positive, 0] = (matrix[positive, 2, 1] - matrix[positive, 1, 2]) / root
    result[positive, 1] = (matrix[positive, 0, 2] - matrix[positive, 2, 0]) / root
    result[positive, 2] = (matrix[positive, 1, 0] - matrix[positive, 0, 1]) / root
    remaining = ~positive
    for axis in range(3):
        next_axis = (axis + 1) % 3
        final_axis = (axis + 2) % 3
        mask = remaining & (
            matrix[..., axis, axis] >= matrix[..., next_axis, next_axis]
        )
        mask &= (
            matrix[..., axis, axis] >= matrix[..., final_axis, final_axis]
        )
        if not np.any(mask):
            continue
        root = np.sqrt(
            np.maximum(
                1.0
                + matrix[mask, axis, axis]
                - matrix[mask, next_axis, next_axis]
                - matrix[mask, final_axis, final_axis],
                1e-8,
            )
        ) * 2.0
        result[mask, axis] = 0.25 * root
        result[mask, 3] = (
            matrix[mask, final_axis, next_axis]
            - matrix[mask, next_axis, final_axis]
        ) / root
        result[mask, next_axis] = (
            matrix[mask, next_axis, axis]
            + matrix[mask, axis, next_axis]
        ) / root
        result[mask, final_axis] = (
            matrix[mask, final_axis, axis]
            + matrix[mask, axis, final_axis]
        ) / root
        remaining &= ~mask
    if np.any(remaining):
        raise ValueError("Rotation 6D conversion left an unresolved matrix branch.")
    result /= np.maximum(np.linalg.norm(result, axis=-1, keepdims=True), 1e-8)
    return result


def extract_audio_features(
    path: Path,
    timeline_times: np.ndarray,
    audio_start: float,
) -> np.ndarray:
    waveform, _ = librosa.load(path, sr=AUDIO_RATE, mono=True)
    mel = librosa.feature.melspectrogram(
        y=waveform,
        sr=AUDIO_RATE,
        n_fft=1024,
        hop_length=AUDIO_HOP,
        n_mels=MEL_BANDS,
        fmin=30,
        fmax=AUDIO_RATE // 2,
        center=True,
        power=2.0,
    )
    mel = np.log1p(mel).T
    rms = librosa.feature.rms(
        y=waveform,
        frame_length=1024,
        hop_length=AUDIO_HOP,
        center=True,
    ).T
    onset = librosa.onset.onset_strength(
        y=waveform,
        sr=AUDIO_RATE,
        hop_length=AUDIO_HOP,
        center=True,
    )[:, None]
    feature_frames = min(len(mel), len(rms), len(onset))
    audio_features = np.concatenate(
        (mel[:feature_frames], rms[:feature_frames], onset[:feature_frames]),
        axis=1,
    ).astype(np.float32)
    feature_times = np.arange(feature_frames, dtype=np.float64) / 60.0
    query_times = timeline_times - float(audio_start)
    aligned = np.zeros((len(timeline_times), audio_features.shape[1]), dtype=np.float32)
    valid = (query_times >= 0.0) & (query_times <= feature_times[-1])
    for feature_index in range(audio_features.shape[1]):
        aligned[valid, feature_index] = np.interp(
            query_times[valid],
            feature_times,
            audio_features[:, feature_index],
        )
    return aligned


def _normalized_generation_input_manifest_value(value: Any) -> Any:
    """Copy JSON data while dropping export-instance-only identity fields."""

    if isinstance(value, Mapping):
        return {
            str(key): _normalized_generation_input_manifest_value(item)
            for key, item in value.items()
            if str(key) not in GENERATION_INPUT_VOLATILE_MANIFEST_FIELDS
        }
    if isinstance(value, (list, tuple)):
        return [
            _normalized_generation_input_manifest_value(item)
            for item in value
        ]
    return value


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 generation_input_preparation_implementation_identity() -> dict[str, Any]:
    """Return a frozen-safe identity for derived generation-input features.

    Source-tree runs hash this module directly.  A PyInstaller runtime instead
    consumes the source hash embedded in the worker build identity, because its
    ``__file__`` can refer to an extracted bytecode artifact rather than the
    source used to build the executable.  The explicit logic version remains a
    reviewable bump point if the feature contract changes independently.
    """

    import cw_camera_runtime

    source_identifier: str | None = None
    if cw_camera_runtime.is_frozen_runtime():
        build_identity = cw_camera_runtime.load_build_identity()
        source_hashes = build_identity.get("sourceSha256", {})
        if isinstance(source_hashes, Mapping):
            matches = [
                str(value)
                for name, value in source_hashes.items()
                if str(name).replace("\\", "/").endswith(
                    "/CameraDirector/camera_runtime_data.py"
                )
                and isinstance(value, str)
                and value
            ]
            if len(matches) == 1:
                source_identifier = matches[0]
        if source_identifier is None:
            value = build_identity.get("preparationLogicIdentifier")
            if isinstance(value, str) and value:
                source_identifier = value
    else:
        source_path = Path(__file__).resolve()
        if source_path.is_file():
            source_identifier = _sha256_file(source_path)

    return {
        "logicVersion": GENERATION_INPUT_PREPARATION_LOGIC_VERSION,
        "sourceIdentifier": (
            source_identifier or GENERATION_INPUT_PREPARATION_LOGIC_VERSION
        ),
        "numpyVersion": str(np.__version__),
        "librosaVersion": str(librosa.__version__),
        "audioRate": AUDIO_RATE,
        "audioHop": AUDIO_HOP,
        "melBands": MEL_BANDS,
        "canonicalJoints": list(CANONICAL_JOINTS),
    }


def generation_input_content_identity(
    record: SongRecord,
    *,
    pipeline_version: int = PIPELINE_VERSION,
) -> dict[str, Any]:
    """Return a path- and export-time-independent generation-input identity.

    The parsed manifest remains part of the identity so changes to skeleton,
    timing, formats, aspect policy, or any other semantic field invalidate the
    cache. Export timestamps, source paths, filenames, and redundant exporter
    fingerprints are provenance only and are removed. Binary inputs are
    addressed by bytes rather than paths, mtimes, or dataset names.
    """

    if not record.is_generation_input:
        raise ValueError(
            "Generation-input content identity requires a camera-free target."
        )
    source_paths = {
        "audio": Path(record.audio_path),
        "jointsWorld": Path(record.song_dir) / "joints_world.f32",
        "root": Path(record.song_dir) / "root.f32",
        "time": Path(record.song_dir) / "time.f64",
    }
    missing = [str(path) for path in source_paths.values() if not path.is_file()]
    if missing:
        raise FileNotFoundError(
            "Generation input cache identity is missing source files: "
            + ", ".join(missing)
        )
    return {
        "contractVersion": GENERATION_INPUT_CACHE_IDENTITY_VERSION,
        "pipelineVersion": int(pipeline_version),
        "inputKind": "generation_input",
        "hasAuthoredCamera": False,
        "preparationImplementation": (
            generation_input_preparation_implementation_identity()
        ),
        "manifest": _normalized_generation_input_manifest_value(
            record.manifest
        ),
        "song": _normalized_generation_input_manifest_value(record.song),
        "sourceSha256": {
            name: _sha256_file(path)
            for name, path in sorted(source_paths.items())
        },
    }


def generation_input_content_identity_sha256(
    record: SongRecord,
    *,
    pipeline_version: int = PIPELINE_VERSION,
) -> str:
    """Return the full digest embedded in and naming a prepared target cache."""

    identity = generation_input_content_identity(
        record,
        pipeline_version=pipeline_version,
    )
    return hashlib.sha256(
        json.dumps(
            identity,
            ensure_ascii=False,
            allow_nan=False,
            sort_keys=True,
            separators=(",", ":"),
        ).encode("utf-8")
    ).hexdigest()


def cache_path_for(record: SongRecord, cache_dir: Path) -> Path:
    source_paths = [
        record.manifest_path,
        record.audio_path,
        record.song_dir / "joints_world.f32",
        record.song_dir / "camera.f32",
        record.song_dir / "root.f32",
        record.song_dir / "shot_index.i32",
        record.song_dir / "time.f64",
    ]
    signature = {
        "pipelineVersion": PIPELINE_VERSION,
        "song": record.name,
        "sources": [
            (str(path), path.stat().st_size, path.stat().st_mtime_ns)
            for path in source_paths
        ],
    }
    digest = hashlib.sha1(
        json.dumps(signature, ensure_ascii=False).encode("utf-8")
    ).hexdigest()[:12]
    safe_name = re.sub(r"[^0-9A-Za-z가-힣_-]+", "_", record.name).strip("_")
    return cache_dir / f"{safe_name}_{digest}.npz"


def generation_input_cache_path_for(
    record: SongRecord,
    cache_dir: Path,
    *,
    pipeline_version: int = PIPELINE_VERSION,
) -> Path:
    """Cache identity for a target that never includes authored camera data."""

    digest = generation_input_content_identity_sha256(
        record,
        pipeline_version=pipeline_version,
    )
    safe_name = re.sub(r"[^\w.-]+", "_", record.name).strip("_")
    return cache_dir / f"{safe_name}_generation_{digest[:24]}.npz"


def prepare_song_inputs(
    record: SongRecord,
) -> tuple[np.ndarray, np.ndarray, float, np.ndarray, np.ndarray]:
    """Prepare motion/audio inputs shared by training and camera-free targets."""

    frames = int(record.song["frameCount"])
    joint_count = int(record.song["sharedBoneCount"])
    source_joints = np.fromfile(
        record.song_dir / "joints_world.f32", dtype="<f4"
    ).reshape(frames, joint_count, 3)
    root = np.fromfile(
        record.song_dir / "root.f32", dtype="<f4"
    ).reshape(frames, 14)
    times = np.fromfile(record.song_dir / "time.f64", dtype="<f8")
    if len(times) != frames:
        raise ValueError(
            f"Timeline time count does not match frameCount for {record.name}: "
            f"{len(times)} != {frames}"
        )

    indices = canonical_joint_indices(record.manifest["sharedBones"])
    joints = source_joints[:, indices]
    root_position = root[:, :3]
    root_rotation = root[:, 3:7]
    inverse_root_rotation = quat_conjugate(root_rotation)

    joints_local = quat_rotate(
        inverse_root_rotation[:, None, :],
        joints - root_position[:, None, :],
    )
    pelvis_index = CANONICAL_JOINTS.index("pelvis")
    head_index = CANONICAL_JOINTS.index("head")
    body_scale = float(
        np.median(
            np.linalg.norm(
                joints_local[:, head_index] - joints_local[:, pelvis_index],
                axis=-1,
            )
        )
    )
    if not math.isfinite(body_scale) or body_scale < 1e-4:
        raise ValueError(f"Invalid body scale for {record.name}: {body_scale}")

    joints_local /= body_scale
    joint_velocity = np.zeros_like(joints_local)
    joint_velocity[1:] = np.diff(joints_local, axis=0) * 60.0

    audio = extract_audio_features(
        record.audio_path,
        times,
        float(record.song["audioStart"]),
    )
    inputs = np.concatenate(
        (
            joints_local.reshape(frames, -1),
            joint_velocity.reshape(frames, -1),
            audio,
        ),
        axis=1,
    ).astype(np.float32)
    if not (np.isfinite(inputs).all() and np.isfinite(times).all()):
        raise ValueError(f"Non-finite input data found in {record.name}.")
    return inputs, times, body_scale, root_position, inverse_root_rotation


def prepare_song(record: SongRecord) -> PreparedSong:
    frames = int(record.song["frameCount"])
    (
        inputs,
        times,
        body_scale,
        root_position,
        inverse_root_rotation,
    ) = prepare_song_inputs(record)
    camera = np.fromfile(record.song_dir / "camera.f32", dtype="<f4").reshape(frames, 9)
    shot_index = np.fromfile(record.song_dir / "shot_index.i32", dtype="<i4")
    if len(shot_index) != frames:
        raise ValueError(
            f"Shot index count does not match frameCount for {record.name}: "
            f"{len(shot_index)} != {frames}"
        )

    camera_local_position = quat_rotate(
        inverse_root_rotation,
        camera[:, :3] - root_position,
    ) / body_scale
    camera_relative_rotation = quat_multiply(
        inverse_root_rotation,
        camera[:, 3:7],
    )
    camera_rotation_6d = quat_to_rotation_6d(camera_relative_rotation)
    targets = np.concatenate(
        (
            camera_local_position,
            camera_rotation_6d,
            camera[:, 7:9],
        ),
        axis=1,
    ).astype(np.float32)

    cuts = np.zeros(frames, dtype=np.float32)
    cuts[1:] = (shot_index[1:] != shot_index[:-1]).astype(np.float32)

    if not (
        np.isfinite(inputs).all()
        and np.isfinite(targets).all()
        and np.isfinite(times).all()
    ):
        raise ValueError(f"Non-finite training data found in {record.name}.")

    return PreparedSong(
        name=record.name,
        inputs=inputs,
        targets=targets,
        cuts=cuts,
        times=times,
        body_scale=body_scale,
        source_manifest=str(record.manifest_path),
    )


def prepare_generation_input_song(record: SongRecord) -> PreparedSong:
    """Prepare a new motion/audio target without touching camera artifacts."""

    if not record.is_generation_input:
        raise ValueError(
            "Camera-free preparation requires inputKind='generation_input' "
            "and hasAuthoredCamera=false."
        )
    inputs, times, body_scale, _, _ = prepare_song_inputs(record)
    frames = len(inputs)
    return PreparedSong(
        name=record.name,
        inputs=inputs,
        # These arrays satisfy the common PreparedSong shape contract only.
        # Hybrid generation never reads target camera labels or authored cuts.
        targets=np.zeros((frames, 11), dtype=np.float32),
        cuts=np.zeros(frames, dtype=np.float32),
        times=times,
        body_scale=body_scale,
        source_manifest=str(record.manifest_path),
    )


def _validated_generation_input_prepared_song(
    record: SongRecord,
    *,
    name: Any,
    inputs: Any,
    targets: Any,
    cuts: Any,
    times: Any,
    body_scale: Any,
) -> PreparedSong:
    """Validate the complete camera-free prepared-array contract."""

    frames = int(record.song["frameCount"])
    feature_count = len(CANONICAL_JOINTS) * 6 + MEL_BANDS + 2
    name_array = np.asarray(name)
    if name_array.shape != () or name_array.dtype.kind not in {"U", "S"}:
        raise ValueError("Generation input cache name must be a string scalar.")
    cached_name = str(name_array.item())
    if cached_name != record.name:
        raise ValueError(
            "Generation input cache song name does not match the target."
        )

    arrays = {
        "inputs": np.asarray(inputs),
        "targets": np.asarray(targets),
        "cuts": np.asarray(cuts),
        "times": np.asarray(times),
    }
    expected = {
        "inputs": ((frames, feature_count), np.dtype(np.float32)),
        "targets": ((frames, 11), np.dtype(np.float32)),
        "cuts": ((frames,), np.dtype(np.float32)),
        "times": ((frames,), np.dtype(np.float64)),
    }
    for key, array in arrays.items():
        expected_shape, expected_dtype = expected[key]
        if array.shape != expected_shape:
            raise ValueError(
                f"Generation input cache {key} shape is invalid: "
                f"{array.shape!r} != {expected_shape!r}."
            )
        if array.dtype != expected_dtype:
            raise ValueError(
                f"Generation input cache {key} dtype is invalid: "
                f"{array.dtype} != {expected_dtype}."
            )
        if not np.all(np.isfinite(array)):
            raise ValueError(
                f"Generation input cache {key} contains non-finite values."
            )
    if np.any(arrays["targets"] != 0.0) or np.any(arrays["cuts"] != 0.0):
        raise ValueError(
            "Generation input cache contains authored-camera placeholder data."
        )
    if frames > 1 and np.any(np.diff(arrays["times"]) <= 0.0):
        raise ValueError("Generation input cache timeline times are not increasing.")

    body_scale_array = np.asarray(body_scale)
    if (
        body_scale_array.shape != ()
        or body_scale_array.dtype != np.dtype(np.float32)
    ):
        raise ValueError(
            "Generation input cache body_scale must be a float32 scalar."
        )
    body_scale_value = float(body_scale_array.item())
    if not math.isfinite(body_scale_value) or body_scale_value < 1e-4:
        raise ValueError("Generation input cache body_scale is invalid.")

    return PreparedSong(
        name=cached_name,
        inputs=np.array(arrays["inputs"], copy=True),
        targets=np.array(arrays["targets"], copy=True),
        cuts=np.array(arrays["cuts"], copy=True),
        times=np.array(arrays["times"], copy=True),
        body_scale=body_scale_value,
        source_manifest=str(record.manifest_path),
    )


def _load_generation_input_cache(
    cache_path: Path,
    record: SongRecord,
    expected_identity_sha256: str,
) -> PreparedSong:
    required = {
        "cache_format",
        "cache_schema_version",
        "content_identity_sha256",
        "preparation_kind",
        "name",
        "inputs",
        "targets",
        "cuts",
        "times",
        "body_scale",
    }
    with np.load(cache_path, allow_pickle=False) as data:
        missing = sorted(required.difference(data.files))
        if missing:
            raise ValueError(
                "Generation input cache is missing fields: "
                + ", ".join(missing)
            )

        cache_format = np.asarray(data["cache_format"])
        schema_version = np.asarray(data["cache_schema_version"])
        identity_sha256 = np.asarray(data["content_identity_sha256"])
        preparation_kind = np.asarray(data["preparation_kind"])
        scalar_strings = {
            "cache_format": cache_format,
            "content_identity_sha256": identity_sha256,
            "preparation_kind": preparation_kind,
        }
        for key, value in scalar_strings.items():
            if value.shape != () or value.dtype.kind not in {"U", "S"}:
                raise ValueError(
                    f"Generation input cache {key} must be a string scalar."
                )
        if str(cache_format.item()) != GENERATION_INPUT_CACHE_FORMAT:
            raise ValueError("Unexpected generation input cache format.")
        if (
            schema_version.shape != ()
            or not np.issubdtype(schema_version.dtype, np.integer)
            or int(schema_version.item())
            != GENERATION_INPUT_CACHE_SCHEMA_VERSION
        ):
            raise ValueError("Unsupported generation input cache schema.")
        if str(identity_sha256.item()) != expected_identity_sha256:
            raise ValueError(
                "Generation input cache identity does not match its target."
            )
        if str(preparation_kind.item()) != "generation_input":
            raise ValueError("Unexpected generation input cache kind.")

        return _validated_generation_input_prepared_song(
            record,
            name=data["name"],
            inputs=data["inputs"],
            targets=data["targets"],
            cuts=data["cuts"],
            times=data["times"],
            body_scale=data["body_scale"],
        )


def _atomic_save_generation_input_cache(
    cache_path: Path,
    arrays: Mapping[str, np.ndarray],
) -> None:
    """Publish one NPZ without ever exposing a partial shared cache file."""

    cache_path = Path(cache_path)
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    descriptor, temporary_name = tempfile.mkstemp(
        prefix=f".{cache_path.name}.tmp-",
        suffix=".npz",
        dir=cache_path.parent,
    )
    temporary_path = Path(temporary_name)
    try:
        with os.fdopen(descriptor, "w+b") as stream:
            descriptor = -1
            np.savez(stream, **arrays)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary_path, cache_path)
    finally:
        if descriptor >= 0:
            os.close(descriptor)
        try:
            temporary_path.unlink()
        except FileNotFoundError:
            pass


def quarantine_generation_input_cache(cache_path: Path) -> Path | None:
    """Move one invalid prepared NPZ aside before safely rebuilding it."""

    cache_path = Path(cache_path)
    if not cache_path.is_file():
        return None
    quarantine_path = cache_path.with_name(
        f".invalid-{cache_path.stem}-{uuid.uuid4().hex}.npz"
    )
    try:
        os.replace(cache_path, quarantine_path)
    except FileNotFoundError:
        return None
    return quarantine_path


def load_or_prepare_generation_input_song(
    record: SongRecord,
    cache_dir: Path,
    rebuild: bool,
    *,
    pipeline_version: int = PIPELINE_VERSION,
    prepare_song_fn: Callable[[SongRecord], PreparedSong] | None = None,
) -> PreparedSong:
    """Load or atomically rebuild a validated target-only prepared cache."""

    cache_dir.mkdir(parents=True, exist_ok=True)
    identity_sha256 = generation_input_content_identity_sha256(
        record,
        pipeline_version=pipeline_version,
    )
    safe_name = re.sub(r"[^\w.-]+", "_", record.name).strip("_")
    cache_path = (
        cache_dir / f"{safe_name}_generation_{identity_sha256[:24]}.npz"
    )
    if cache_path.is_file() and not rebuild:
        try:
            return _load_generation_input_cache(
                cache_path,
                record,
                identity_sha256,
            )
        except (
            EOFError,
            KeyError,
            OSError,
            TypeError,
            ValueError,
            zipfile.BadZipFile,
        ):
            quarantine_generation_input_cache(cache_path)

    prepare = prepare_song_fn or prepare_generation_input_song
    prepared = prepare(record)
    validated = _validated_generation_input_prepared_song(
        record,
        name=np.asarray(prepared.name),
        inputs=prepared.inputs,
        targets=prepared.targets,
        cuts=prepared.cuts,
        times=prepared.times,
        body_scale=np.asarray(prepared.body_scale, dtype=np.float32),
    )
    current_identity_sha256 = generation_input_content_identity_sha256(
        record,
        pipeline_version=pipeline_version,
    )
    if current_identity_sha256 != identity_sha256:
        raise RuntimeError(
            "Generation input sources changed during preparation; retry after "
            "the Unity export is stable. No prepared cache was published."
        )
    arrays = {
        "cache_format": np.asarray(GENERATION_INPUT_CACHE_FORMAT),
        "cache_schema_version": np.asarray(
            GENERATION_INPUT_CACHE_SCHEMA_VERSION,
            dtype=np.int64,
        ),
        "content_identity_sha256": np.asarray(identity_sha256),
        "preparation_kind": np.asarray("generation_input"),
        "name": np.asarray(validated.name),
        "inputs": validated.inputs,
        "targets": validated.targets,
        "cuts": validated.cuts,
        "times": validated.times,
        "body_scale": np.asarray(validated.body_scale, dtype=np.float32),
    }
    _atomic_save_generation_input_cache(cache_path, arrays)
    return validated


def load_or_prepare_song(
    record: SongRecord,
    cache_dir: Path,
    rebuild: bool,
) -> PreparedSong:
    if record.portable_prepared_path.is_file() and not rebuild:
        with np.load(record.portable_prepared_path, allow_pickle=False) as data:
            return PreparedSong(
                name=str(data["name"].item()),
                inputs=data["inputs"],
                targets=data["targets"],
                cuts=data["cuts"],
                times=data["times"],
                body_scale=float(data["body_scale"].item()),
                source_manifest=str(data["source_manifest"].item()),
            )

    cache_dir.mkdir(parents=True, exist_ok=True)
    cache_path = cache_path_for(record, cache_dir)
    if cache_path.is_file() and not rebuild:
        with np.load(cache_path, allow_pickle=False) as data:
            return PreparedSong(
                name=str(data["name"].item()),
                inputs=data["inputs"],
                targets=data["targets"],
                cuts=data["cuts"],
                times=data["times"],
                body_scale=float(data["body_scale"].item()),
                source_manifest=str(data["source_manifest"].item()),
            )

    prepared = prepare_song(record)
    np.savez(
        cache_path,
        name=np.asarray(prepared.name),
        inputs=prepared.inputs,
        targets=prepared.targets,
        cuts=prepared.cuts,
        times=prepared.times,
        body_scale=np.asarray(prepared.body_scale, dtype=np.float32),
        source_manifest=np.asarray(prepared.source_manifest),
    )
    return prepared
