"""Command-line orchestration: tokenize, release the encoder, then render."""

import argparse
import gc
import os
from pathlib import Path

import numpy as np
import torch

from . import __version__
from .audio import load_audio, save_wav_atomic
from .quantizer import WannabeTokenizer
from .renderer import YuE2Renderer


def positive_int(value: str) -> int:
    try:
        parsed = int(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("must be an integer") from exc
    if str(parsed) != value or parsed < 1:
        raise argparse.ArgumentTypeError("must be a positive integer without extra characters")
    return parsed


def save_tokens_atomic(path: Path, tokens: np.ndarray) -> None:
    temporary = path.with_name(path.name + ".tmp")
    try:
        with temporary.open("wb") as f:
            np.save(f, tokens, allow_pickle=False)
        os.replace(temporary, path)
    finally:
        if temporary.exists():
            temporary.unlink()


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Tokenize one audio file with WannabeTokenizer and reconstruct it with YuE2."
    )
    parser.add_argument("input", type=Path, help="Input audio file")
    parser.add_argument("--tokenizer", required=True, type=Path, help="Unified WannabeTokenizer .safetensors file")
    parser.add_argument("--yue2-checkpoint", required=True, type=Path, help="YuE2 bf16 all-in-one .safetensors file")
    parser.add_argument("--output-dir", type=Path, default=Path("output"), help="Destination folder (default: output)")
    parser.add_argument("--device", default="auto", help="Torch device, for example auto, cuda, cuda:0, or cpu")
    parser.add_argument("--seed", type=int, default=0, help="NAR noise seed (default: 0)")
    parser.add_argument("--steps", type=positive_int, default=32, help="Flow-matching steps (default: 32)")
    parser.add_argument("--batch-windows", type=positive_int, default=8, help="Tokenizer windows per batch (default: 8)")
    parser.add_argument("--overwrite", action="store_true", help="Replace existing tokens.npy and reconstructed.wav")
    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
    return parser


def main() -> None:
    args = build_parser().parse_args()
    audio_path = args.input.expanduser()
    tokenizer_path = args.tokenizer.expanduser()
    yue2_path = args.yue2_checkpoint.expanduser()
    output_dir = args.output_dir.expanduser()
    for label, path in (
        ("Input audio", audio_path),
        ("Tokenizer checkpoint", tokenizer_path),
        ("YuE2 checkpoint", yue2_path),
    ):
        if not path.is_file():
            raise FileNotFoundError(f"{label} not found: {path}")

    device = "cuda" if args.device == "auto" and torch.cuda.is_available() else args.device
    if device == "auto":
        device = "cpu"
    if device.startswith("cuda") and not torch.cuda.is_available():
        raise RuntimeError("A CUDA device was requested, but CUDA is not available")
    if device == "cpu":
        print("Warning: YuE2 reconstruction on CPU is very slow and requires substantial RAM.")

    tokens_path = output_dir / "tokens.npy"
    audio_output_path = output_dir / "reconstructed.wav"
    outputs = (tokens_path, audio_output_path)
    if not args.overwrite:
        existing = [str(path) for path in outputs if path.exists()]
        if existing:
            raise FileExistsError("Output already exists (use --overwrite to replace): " + ", ".join(existing))
    output_dir.mkdir(parents=True, exist_ok=True)

    waveform = load_audio(audio_path)
    print(f"Tokenizing on {device}...")
    tokenizer = WannabeTokenizer(tokenizer_path, device=device)
    tokens = tokenizer.tokenize(waveform, batch_windows=args.batch_windows)
    save_tokens_atomic(tokens_path, tokens)
    print(f"Wrote {tokens_path} ({len(tokens)} frames, {len(tokens) / 25:.2f} seconds)")

    # The encoder and renderer do not need to coexist in accelerator memory.
    del tokenizer
    gc.collect()
    if device.startswith("cuda"):
        torch.cuda.empty_cache()

    print(f"Reconstructing on {device} with seed {args.seed} and {args.steps} steps...")
    renderer = YuE2Renderer(yue2_path, device=device)
    reconstructed = renderer.reconstruct(tokens, seed=args.seed, steps=args.steps)
    save_wav_atomic(audio_output_path, reconstructed)
    print(f"Wrote {audio_output_path} (48 kHz stereo WAV)")