#!/usr/bin/env python3
"""Run the bounded arbitrary-WAV continuous-latent continuation experiment."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Sequence

from music3lab.editing.audio_continuation_runner import (
    load_audio_continuation_config,
    run_audio_continuation,
)


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Evaluate frozen continuous encoders for local arbitrary-WAV "
            "continuation; this is not native AR/RVQ continuation."
        )
    )
    commands = parser.add_subparsers(dest="command", required=True)
    check = commands.add_parser("check-config")
    check.add_argument("--config", type=Path, required=True)
    run = commands.add_parser("run")
    run.add_argument("--config", type=Path, required=True)
    run.add_argument("--flow-config", type=Path, required=True)
    run.add_argument("--audio", type=Path, required=True)
    run.add_argument("--baseline-checkpoint", type=Path, required=True)
    run.add_argument("--specialist-checkpoint", type=Path, required=True)
    run.add_argument("--snapshot", type=Path, required=True)
    run.add_argument("--base-manifest", type=Path, required=True)
    run.add_argument("--diffusers-root", type=Path, required=True)
    run.add_argument("--output-root", type=Path, required=True)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    arguments = _parser().parse_args(argv)
    if arguments.command == "check-config":
        config, file_sha256 = load_audio_continuation_config(arguments.config)
        result = {
            "schema_version": config.schema_version,
            "config_file_sha256": file_sha256,
            "capability": config.capability,
            "specialist_champion_eligible": config.specialist_champion_eligible,
            "native_tokens_used": False,
        }
    elif arguments.command == "run":
        metrics = run_audio_continuation(
            config_path=arguments.config,
            flow_config_path=arguments.flow_config,
            source_audio_path=arguments.audio,
            baseline_checkpoint=arguments.baseline_checkpoint,
            specialist_checkpoint=arguments.specialist_checkpoint,
            snapshot=arguments.snapshot,
            base_manifest=arguments.base_manifest,
            diffusers_root=arguments.diffusers_root,
            output_root=arguments.output_root,
        )
        result = {
            "output_root": str(arguments.output_root.absolute()),
            "status": metrics["status"],
            "frozen_af2ed70_gate": metrics["encoders"]["frozen_af2ed70"][
                "gate"
            ]["passes"],
            "rejected_external_specialist_gate": metrics["encoders"][
                "rejected_external_specialist"
            ]["gate"]["passes"],
            "specialist_promoted": metrics["specialist_promoted"],
        }
    else:
        raise AssertionError("unreachable command")
    print(json.dumps(result, sort_keys=True, separators=(",", ":")))
    return 0 if result.get("status", "FUNCTIONAL_PASS") == "FUNCTIONAL_PASS" else 1


if __name__ == "__main__":
    raise SystemExit(main())