# ruff: noqa: I001
import spaces  # Must be imported before torch for ZeroGPU CUDA emulation.

import math
import os
from pathlib import Path

import gradio as gr
import soundfile as sf
import torch
from pocket_tts import TTSModel

from private_assets import build_runtime_model_config
from ruaccent_loader import load_ruaccent
from text_processing import MAX_TEXT_CHARS, prepare_russian_text, to_model_stress_notation
from voice_catalog import (
    DEFAULT_VOICE,
    VOICE_CHOICES,
    get_builtin_voice_state,
    validate_builtin_voice_assets,
)

APP_DIR = Path(__file__).resolve().parent
MODEL_CONFIG_TEMPLATE = APP_DIR / "model" / "config.yaml"
MODEL_CONFIG = build_runtime_model_config(MODEL_CONFIG_TEMPLATE)
RUACCENT_WORKDIR = APP_DIR / "ruaccent_data"
DEVICE = os.environ.get("POCKET_TTS_DEVICE", "cuda")
MIN_VOICE_SECONDS = 1.0
MAX_VOICE_SECONDS = 30.0

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")


def _validate_assets() -> None:
    required = [MODEL_CONFIG_TEMPLATE]
    missing = [str(path.relative_to(APP_DIR)) for path in required if not path.is_file()]
    if missing:
        raise RuntimeError("Missing Space assets: " + ", ".join(missing))


_validate_assets()
validate_builtin_voice_assets()

# ZeroGPU expects CUDA placement at module level. Outside ZeroGPU, launch with
# POCKET_TTS_DEVICE=cpu for local validation.
tts_model = TTSModel.load_model(config=MODEL_CONFIG, temp=0.5, eos_threshold=-1.0).to(DEVICE)
tts_model.eval()

# RUAccent remains on CPU and runs before the decorated GPU callback.
accentizer = load_ruaccent(RUACCENT_WORKDIR)


def _validate_custom_voice(voice_path: str | None) -> None:
    if not voice_path:
        raise gr.Error("Загрузите или запишите голосовой референс.")

    try:
        info = sf.info(voice_path)
    except Exception as exc:
        raise gr.Error("Не удалось прочитать голосовой референс.") from exc

    duration = info.frames / info.samplerate
    if duration < MIN_VOICE_SECONDS:
        raise gr.Error("Голосовой референс должен быть не короче одной секунды.")
    if duration > MAX_VOICE_SECONDS:
        raise gr.Error("Голосовой референс должен быть не длиннее 30 секунд.")


def _validate_voice_source(voice_choice: str, voice_path: str | None) -> Path | str | None:
    if voice_path:
        _validate_custom_voice(voice_path)
        return None

    try:
        return get_builtin_voice_state(voice_choice)
    except ValueError as exc:
        raise gr.Error(str(exc)) from exc


def prepare_request(
    text: str, use_automatic_stress: bool, voice_choice: str, voice_path: str | None
) -> str:
    """CPU preflight: reject invalid requests and add canonical stress marks."""

    _validate_voice_source(voice_choice, voice_path)
    try:
        return prepare_russian_text(text, use_ruaccent=use_automatic_stress, accentizer=accentizer)
    except ValueError as exc:
        raise gr.Error(str(exc)) from exc


def preview_stress(text: str, use_automatic_stress: bool) -> str:
    try:
        return prepare_russian_text(text, use_ruaccent=use_automatic_stress, accentizer=accentizer)
    except ValueError as exc:
        raise gr.Error(str(exc)) from exc


def estimate_gpu_duration(
    prepared_text: str, voice_choice: str, voice_path: str | None, seed: int
) -> int:
    """Reserve realistic ZeroGPU time; short declarations also improve queue priority."""

    del voice_choice, voice_path, seed
    text_length = len(prepared_text or "")
    return min(36, max(12, 10 + 4 * math.ceil(text_length / 120)))


@spaces.GPU(duration=estimate_gpu_duration)
def synthesize(prepared_text: str, voice_choice: str, voice_path: str | None, seed: int):
    built_in_state = _validate_voice_source(voice_choice, voice_path)
    if not prepared_text:
        raise gr.Error("Сначала подготовьте текст.")

    seed = int(seed)
    torch.manual_seed(seed)
    if DEVICE == "cuda":
        torch.cuda.manual_seed_all(seed)

    if built_in_state is not None:
        voice_state = tts_model.get_state_for_audio_prompt(built_in_state)
    else:
        voice_state = tts_model.get_state_for_audio_prompt(Path(voice_path), truncate=True)
    # The demo exposes easy-to-type +vowel marks, while the tokenizer was
    # trained with a combining acute after the stressed vowel.
    model_text = to_model_stress_notation(prepared_text)
    audio = tts_model.generate_audio(voice_state, model_text, copy_state=True)
    samples = audio.detach().float().cpu().numpy()
    return tts_model.sample_rate, samples


DEFAULT_TEXT = "Ст+арый з+амок ст+оит д+орого, а тяжёлую дв+ерь закрыв+ает надёжный зам+ок."

with gr.Blocks(title="xVibePocketTTS") as demo:
    gr.Markdown(
        """
        # xVibePocketTTS

        Публичная демо-версия. Выберите один из трёх встроенных голосов
        или загрузите свой референс, введите русский текст и запустите синтез.
        Автоматические ударения ставит RUAccent. Ручное ударение обозначайте
        знаком `+` перед ударной гласной: `з+амок` или `зам+ок`. Оно всегда имеет приоритет.
        Над `ё` знак ударения не нужен.

        Telegram автора:
        <a href="https://t.me/xVibeNot" target="_blank" rel="noopener noreferrer">t.me/xVibeNot</a>
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            source_text = gr.Textbox(
                label=f"Текст, до {MAX_TEXT_CHARS} символов", value=DEFAULT_TEXT, lines=5
            )
            use_automatic_stress = gr.Checkbox(
                label="Автоматически расставлять ударения с RUAccent", value=True
            )
            prepared_text = gr.Textbox(
                label="Текст, который получит модель", lines=5, interactive=False, buttons=["copy"]
            )
        with gr.Column(scale=2):
            voice_choice = gr.Radio(label="Голос", choices=VOICE_CHOICES, value=DEFAULT_VOICE)
            gr.Markdown(
                "Если загружен или записан свой референс, он имеет приоритет "
                "над выбранным встроенным голосом."
            )
            voice_audio = gr.Audio(
                label="Свой референс (1–30 секунд)",
                sources=["upload", "microphone"],
                type="filepath",
            )
            seed = gr.Number(label="Seed", value=20260829, precision=0)

    with gr.Row():
        preview_button = gr.Button("Только проверить ударения")
        generate_button = gr.Button("Синтезировать", variant="primary")

    output_audio = gr.Audio(label="Результат", autoplay=False)

    preview_button.click(
        fn=preview_stress,
        inputs=[source_text, use_automatic_stress],
        outputs=prepared_text,
        concurrency_limit=1,
    )

    prepared_event = generate_button.click(
        fn=prepare_request,
        inputs=[source_text, use_automatic_stress, voice_choice, voice_audio],
        outputs=prepared_text,
        concurrency_limit=1,
    )
    prepared_event.then(
        fn=synthesize,
        inputs=[prepared_text, voice_choice, voice_audio, seed],
        outputs=output_audio,
        concurrency_limit=1,
    )

    gr.Markdown("Не используйте результат для имитации человека без его согласия.")

demo.queue(default_concurrency_limit=1, max_size=8)

if __name__ == "__main__":
    demo.launch()