#!/usr/bin/env python3
"""Short CPU proof for the Qwen engine used by Portable PRO."""

from __future__ import annotations

import json
import platform
import subprocess
import time
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
from qwen_tts import Qwen3TTSModel


ROOT = Path(__file__).resolve().parent
REFERENCE = ROOT / "references" / "my-voice-20s-24k-mono.wav"
TEXT = "Браузер передаёт микрофон по WebRTC. Core превращает звук в пакеты и отправляет их через NATS на распознавание."


started = time.perf_counter()
model = Qwen3TTSModel.from_pretrained(
    "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
    device_map="cpu",
    dtype=torch.bfloat16,
    attn_implementation="eager",
    low_cpu_mem_usage=True,
    local_files_only=True,
)
loaded = time.perf_counter()
prompt = model.create_voice_clone_prompt(
    ref_audio=str(REFERENCE),
    ref_text=None,
    x_vector_only_mode=True,
)
prompted = time.perf_counter()
wavs, sample_rate = model.generate_voice_clone(
    text=TEXT,
    language="Russian",
    voice_clone_prompt=prompt,
    max_new_tokens=180,
    do_sample=True,
)
generated = time.perf_counter()

raw = ROOT / "raw" / "qwen-0.6b-base-my-voice-cpu.wav"
sf.write(raw, np.asarray(wavs[0], dtype=np.float32), sample_rate, subtype="PCM_16")
mp3 = ROOT / "demos" / "qwen-0.6b-base-my-voice-cpu.mp3"
subprocess.run(
    [
        "ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", str(raw),
        "-af", "loudnorm=I=-18:TP=-2:LRA=11", "-codec:a", "libmp3lame", "-q:a", "3", str(mp3),
    ],
    check=True,
)
finished = time.perf_counter()
duration = len(wavs[0]) / sample_rate
metrics = {
    "model": "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
    "portable_project": "timoncool/Qwen3-TTS_portable_rus",
    "portable_ui_executed": False,
    "engine_execution": "official Qwen Python package, same upstream checkpoint family used by Portable PRO",
    "mode": "voice_clone_xvector_only",
    "voice_reference": "user-provided my voice.flac, normalized to 20 s / 24 kHz / mono",
    "device": "cpu",
    "dtype": "bfloat16",
    "host": platform.platform(),
    "text": TEXT,
    "text_characters": len(TEXT),
    "audio_duration_s": round(duration, 3),
    "sample_rate_hz": sample_rate,
    "model_load_s": round(loaded - started, 3),
    "prompt_encode_s": round(prompted - loaded, 3),
    "generation_s": round(generated - prompted, 3),
    "total_s": round(finished - started, 3),
    "rtf_generation": round((generated - prompted) / duration, 3),
    "max_new_tokens": 180,
}
(ROOT / "evidence" / "qwen-0.6b-base-my-voice-cpu.json").write_text(
    json.dumps(metrics, ensure_ascii=False, indent=2) + "\n",
    encoding="utf-8",
)
print(json.dumps(metrics, ensure_ascii=False, indent=2))