#!/usr/bin/env python3
"""Run a reproducible CPU voice-clone check for ESpeech or F5-TTS Russian."""

from __future__ import annotations

import argparse
import json
import os
import platform
import resource
import subprocess
import time
from importlib.resources import files
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
from f5_tts.infer.utils_infer import (
    infer_process,
    load_model,
    load_vocoder,
    preprocess_ref_audio_text,
)
from f5_tts.model import DiT
from huggingface_hub import hf_hub_download
from omegaconf import OmegaConf


ROOT = Path(__file__).resolve().parent
MODEL_CACHE = Path(os.environ.get("TTS_LOCAL_MODEL_DIR", Path.home() / ".cache/tts-local/models"))
REFERENCE = ROOT / "references" / "my-voice-11.5s-16k-mono.wav"
REFERENCE_TEXT = (
    "Представим простой запрос. Покажи ближайшие аэропорты. "
    "Пользователь говорит в микрофон, а WebRTC Комаров должен превратить эту фразу "
    "в понятное действие и показать результат на карте."
)
GEN_TEXT = (
    "Браузер передаёт микрофон по WebRTC. Core превращает звук в короткие пакеты "
    "и отправляет их через NATS на распознавание."
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("model", choices=("espeech", "f5-russian"))
    parser.add_argument("--nfe-step", type=int, default=16)
    parser.add_argument("--threads", type=int, default=8)
    parser.add_argument("--seed", type=int, default=1234)
    return parser.parse_args()


def checkpoint(model: str) -> tuple[str, str, dict, str]:
    if model == "espeech":
        repo = "ESpeech/ESpeech-TTS-1_RL-V2"
        local_ckpt = MODEL_CACHE / "espeech_tts_rlv2.pt"
        ckpt = str(local_ckpt) if local_ckpt.exists() else hf_hub_download(repo, filename="espeech_tts_rlv2.pt")
        vocab = hf_hub_download(repo, filename="vocab.txt")
        model_cfg = {
            "dim": 1024,
            "depth": 22,
            "heads": 16,
            "ff_mult": 2,
            "text_dim": 512,
            "conv_layers": 4,
        }
        slug = "espeech-tts-1-rl-v2-my-voice-cpu"
    else:
        repo = "Misha24-10/F5-TTS_RUSSIAN"
        local_ckpt = MODEL_CACHE / "f5-tts-russian-v2-model_last_inference.safetensors"
        ckpt = (
            str(local_ckpt)
            if local_ckpt.exists()
            else hf_hub_download(repo, filename="F5TTS_v1_Base_v2/model_last_inference.safetensors")
        )
        vocab = hf_hub_download(repo, filename="F5TTS_v1_Base/vocab.txt")
        config = OmegaConf.load(str(files("f5_tts").joinpath("configs/F5TTS_v1_Base.yaml")))
        model_cfg = OmegaConf.to_container(config.model.arch, resolve=True)
        slug = "f5-tts-russian-v2-my-voice-cpu"
    return str(ckpt), str(vocab), model_cfg, slug


def main() -> None:
    args = parse_args()
    torch.set_num_threads(args.threads)
    torch.set_num_interop_threads(1)
    torch.manual_seed(args.seed)

    started = time.perf_counter()
    ckpt, vocab, model_cfg, slug = checkpoint(args.model)
    downloaded = time.perf_counter()

    vocoder = load_vocoder(device="cpu")
    vocoder_loaded = time.perf_counter()
    model = load_model(DiT, model_cfg, ckpt, vocab_file=vocab, device="cpu")
    model_loaded = time.perf_counter()
    ref_audio, ref_text = preprocess_ref_audio_text(str(REFERENCE), REFERENCE_TEXT, show_info=print)
    prepared = time.perf_counter()
    wave, sample_rate, _ = infer_process(
        ref_audio,
        ref_text,
        GEN_TEXT,
        model,
        vocoder,
        nfe_step=args.nfe_step,
        speed=1.0,
        device="cpu",
        show_info=print,
    )
    generated = time.perf_counter()

    raw_path = ROOT / "raw" / f"{slug}.wav"
    mp3_path = ROOT / "demos" / f"{slug}.mp3"
    raw_path.parent.mkdir(exist_ok=True)
    mp3_path.parent.mkdir(exist_ok=True)
    sf.write(raw_path, np.asarray(wave, dtype=np.float32), sample_rate, subtype="PCM_16")
    subprocess.run(
        [
            "ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", str(raw_path),
            "-af", "loudnorm=I=-18:TP=-2:LRA=11", "-codec:a", "libmp3lame", "-q:a", "3", str(mp3_path),
        ],
        check=True,
    )
    finished = time.perf_counter()

    duration = len(wave) / sample_rate
    metrics = {
        "model": args.model,
        "checkpoint": Path(ckpt).name,
        "reference": "private user audio (not published)",
        "reference_duration_s": 11.5,
        "reference_text": REFERENCE_TEXT,
        "text": GEN_TEXT,
        "text_characters": len(GEN_TEXT),
        "device": "cpu",
        "dtype": "float32",
        "threads": args.threads,
        "seed": args.seed,
        "nfe_step": args.nfe_step,
        "host": platform.platform(),
        "sample_rate_hz": sample_rate,
        "audio_duration_s": round(duration, 3),
        "checkpoint_resolution_s": round(downloaded - started, 3),
        "vocoder_load_s": round(vocoder_loaded - downloaded, 3),
        "model_load_s": round(model_loaded - vocoder_loaded, 3),
        "reference_prepare_s": round(prepared - model_loaded, 3),
        "generation_s": round(generated - prepared, 3),
        "write_s": round(finished - generated, 3),
        "total_s": round(finished - started, 3),
        "rtf_generation": round((generated - prepared) / duration, 3),
        "max_rss_kib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
        "output": str(mp3_path.relative_to(ROOT)),
    }
    metrics["max_rss_gib"] = round(metrics["max_rss_kib"] / 1024 / 1024, 3)
    evidence = ROOT / "evidence" / f"{slug}.json"
    evidence.write_text(json.dumps(metrics, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(metrics, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()