#!/usr/bin/env python
"""Build immutable preparation caches for the hybrid camera generator."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

import generate_hybrid
import hybrid_preparation_cache as preparation_cache
import planner
import camera_runtime_data as train
import cw_camera_runtime


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--song", default="")
    parser.add_argument("--song-id", default="")
    parser.add_argument("--dataset-id", default="")
    parser.add_argument("--data", nargs="+", default=["DatasetExports"])
    parser.add_argument(
        "--training-index",
        type=Path,
        default=Path("reports/training_index.json"),
    )
    parser.add_argument(
        "--library-split",
        choices=("train", "all"),
        default="train",
    )
    parser.add_argument(
        "--preparation-model-seed",
        "--seed",
        dest="preparation_model_seed",
        type=int,
        default=generate_hybrid.PREPARATION_MODEL_SEED,
        help=(
            "Stable RF/planner seed. --seed remains a compatibility alias."
        ),
    )
    parser.add_argument("--beam-width", type=int, default=6)
    parser.add_argument(
        "--cache-dir",
        type=Path,
        default=Path(".cache/camera_director"),
        help="Prepared-song fallback cache used when portable arrays are absent.",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        default=Path(".cache/hybrid_preparation"),
    )
    parser.add_argument(
        "--report-json",
        type=Path,
        default=None,
    )
    return parser.parse_args()


def compute_beat_onset_cache(
    record: train.SongRecord,
    prepared: train.PreparedSong,
) -> preparation_cache.BeatOnsetCache:
    return generate_hybrid.build_beat_onset_cache(record, prepared)


def run(args: argparse.Namespace) -> dict:
    train.configure_utf8_stdio()
    timer = preparation_cache.StageTimer()
    repository_root, camera_director_root = (
        cw_camera_runtime.source_tree_roots(
            Path(__file__).resolve().parents[2],
            Path(__file__).resolve().parent,
        )
    )
    training_index_path = args.training_index.resolve()
    runtime_versions = generate_hybrid.preparation_cache_versions()

    with timer.stage("catalog_and_policy"):
        training_index, entries = generate_hybrid.load_training_policy(
            training_index_path
        )
        manifest_paths = train.find_manifest_paths(
            generate_hybrid.generation_data_inputs(
                args.data,
                training_index_path,
                training_index,
            )
        )
        records = train.load_records(
            manifest_paths,
            skip_unavailable=True,
        )
        target_record = train.select_record(
            records,
            song_name=args.song,
            song_id=args.song_id,
            dataset_id=args.dataset_id,
        )
        (
            training_records,
            planner_records,
            target_entry,
            target_group_available,
        ) = generate_hybrid.select_generation_records(
            records,
            entries,
            target_record,
            args.library_split,
        )
        training_mask_audit = (
            generate_hybrid.audio_aligned_template_mask_audit(
                training_records,
                entries,
            )
        )
        planner_mask_audit = generate_hybrid.audio_aligned_template_mask_audit(
            planner_records,
            entries,
        )
        build_inputs = preparation_cache.build_inputs_for_records(
            repository_root,
            camera_director_root,
            training_index_path,
            target_record,
            training_records,
            planner_records,
            entries,
            args.library_split,
            args.preparation_model_seed,
            args.beam_width,
            runtime_versions,
        )
        cache_key = preparation_cache.cache_key_for(build_inputs)
        expected_path = args.output_root.resolve() / cache_key

    cached_manifest = None
    with timer.stage("cache_lookup"):
        if expected_path.is_dir():
            cached_manifest = preparation_cache.validate_cache_bundle(
                expected_path,
                cache_key,
            )
    if cached_manifest is not None:
        report = {
            "cacheHit": True,
            "cacheKey": cache_key,
            "cachePath": str(expected_path),
            "targetSongId": target_record.id,
            "targetDatasetId": target_record.dataset_id,
            "summary": preparation_cache.report_summary_for_target(
                cached_manifest.get("summary", {}),
                target_record,
            ),
            **timer.snapshot(),
        }
        _write_report(args.report_json, report)
        return report

    required_records = {
        record.id: record
        for record in [target_record, *training_records]
    }
    with timer.stage("load_prepared_songs"):
        prepared_by_id = {
            record.id: (
                train.load_or_prepare_generation_input_song(
                    record,
                    args.cache_dir.resolve(),
                    False,
                )
                if record.id == target_record.id
                and record.is_generation_input
                else train.load_or_prepare_song(
                    record,
                    args.cache_dir.resolve(),
                    False,
                )
            )
            for record in required_records.values()
        }

    with timer.stage("build_template_library"):
        aspect_by_id = {
            song_id: generate_hybrid.aspect_for_entry(entry)
            for song_id, entry in entries.items()
        }
        templates = generate_hybrid.build_template_library(
            training_records,
            prepared_by_id,
            aspect_by_id=aspect_by_id,
            policy_by_id=entries,
            shot_label_mode="semantic",
        )
        planner_record_ids = {record.id for record in planner_records}
        planner_templates = [
            template
            for template in templates
            if template.source_id in planner_record_ids
            and template.shot_type_confidence_tier != "low"
            and template.motion_confidence >= 0.45
        ]

    evaluation_group_by_id = {
        song_id: str(entry["evaluationGroupId"])
        for song_id, entry in entries.items()
    }
    with timer.stage("fit_shot_classifier"):
        shot_classifier, cross_validation_accuracy = (
            generate_hybrid.train_shot_classifier(
                planner_templates,
                evaluation_group_by_id,
            )
        )

    with timer.stage("fit_hierarchical_planner"):
        planner_samples = generate_hybrid.build_planner_samples(
            planner_templates,
            evaluation_group_by_id,
        )
        hierarchical_planner = planner.HierarchicalShotPlanner(
            random_state=args.preparation_model_seed,
            beam_width=args.beam_width,
        )
        hierarchical_planner.fit(planner_samples)

    with timer.stage("build_beat_onset_cache"):
        beat_onset = compute_beat_onset_cache(
            target_record,
            prepared_by_id[target_record.id],
        )

    with timer.stage("verify_build_inputs_stable"):
        current_build_inputs = preparation_cache.build_inputs_for_records(
            repository_root,
            camera_director_root,
            training_index_path,
            target_record,
            training_records,
            planner_records,
            entries,
            args.library_split,
            args.preparation_model_seed,
            args.beam_width,
            runtime_versions,
        )
        if preparation_cache.cache_key_for(current_build_inputs) != cache_key:
            raise RuntimeError(
                "Hybrid generator source or training policy changed while the "
                "preparation cache was being built. Retry after edits finish."
            )

    feature_matrix = np.stack([template.features for template in templates])
    feature_mean = feature_matrix.mean(axis=0)
    feature_std = np.maximum(feature_matrix.std(axis=0), 1e-5)
    models = {
        "shotClassifier": shot_classifier,
        "hierarchicalPlanner": hierarchical_planner,
        "featureMean": feature_mean.astype(np.float32),
        "featureStd": feature_std.astype(np.float32),
        "shotClassifierCrossValidationAccuracy": (
            cross_validation_accuracy
        ),
    }
    summary = {
        "targetSongId": target_record.id,
        "targetDatasetId": target_record.dataset_id,
        "targetEvaluationGroupId": generate_hybrid.target_reference_evaluation_group(
            entries,
            target_record,
        ),
        "targetTrainingIndexMembership": target_entry is not None,
        "targetInputKind": (
            "generation_input"
            if target_record.is_generation_input
            else "authored_camera_record"
        ),
        "targetEvaluationGroupAvailableRecordCount": (
            target_group_available
        ),
        "preparationModelSeed": args.preparation_model_seed,
        "trainingSongCount": len(training_records),
        "plannerSongCount": len(planner_records),
        "templateCount": len(templates),
        "plannerTemplateCount": len(planner_templates),
        "plannerSampleCount": len(planner_samples),
        "frameCount": beat_onset.frame_count,
        "beatCount": len(beat_onset.beat_frames),
        "peakCount": len(beat_onset.peak_frames),
        "cutCandidateCount": len(beat_onset.candidates),
        "shotClassifierCrossValidationAccuracy": (
            cross_validation_accuracy
        ),
        "buildTimingBeforePublish": timer.snapshot(),
    }
    training_mask_audit = (
        generate_hybrid.finalized_audio_aligned_template_mask_audit(
            training_mask_audit,
            templates,
        )
    )
    planner_mask_audit = (
        generate_hybrid.finalized_audio_aligned_template_mask_audit(
            planner_mask_audit,
            planner_templates,
        )
    )
    if training_mask_audit is not None:
        summary["audioAlignedFrameMasking"] = {
            "trainingTemplates": training_mask_audit,
            "shotPlanner": planner_mask_audit,
        }

    with timer.stage("publish_and_verify"):
        bundle = preparation_cache.publish_cache_bundle(
            args.output_root,
            build_inputs,
            templates,
            models,
            beat_onset,
            summary,
        )
        preparation_cache.validate_cache_bundle(
            bundle.path,
            bundle.cache_key,
        )

    report = {
        "cacheHit": bundle.reused,
        "cacheKey": bundle.cache_key,
        "cachePath": str(bundle.path),
        "targetSongId": target_record.id,
        "targetDatasetId": target_record.dataset_id,
        "summary": preparation_cache.report_summary_for_target(
            summary,
            target_record,
        ),
        **timer.snapshot(),
    }
    _write_report(args.report_json, report)
    return report


def _write_report(path: Path | None, report: dict) -> None:
    if path is not None:
        preparation_cache.atomic_write_json(path.resolve(), report)


def main() -> None:
    report = run(parse_args())
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
