from __future__ import annotations

import importlib.util
import io
import logging
import sys
import threading
import time
import warnings
import wave
from pathlib import Path
from types import ModuleType
from typing import Any, Iterator, Protocol, cast

import gradio as gr
import numpy as np

from example_texts import EXAMPLE_TEXTS

warnings.filterwarnings("ignore", message="`torch.nn.utils.weight_norm` is deprecated")

try:
    import torch  # type: ignore
except Exception:  # pragma: no cover
    torch = None  # type: ignore

try:
    from huggingface_hub import snapshot_download  # type: ignore
    from huggingface_hub.utils import disable_progress_bars  # type: ignore

    disable_progress_bars()
except Exception:  # pragma: no cover
    snapshot_download = None  # type: ignore


class _PhonemizerWordMismatchFilter(logging.Filter):
    def filter(self, record: logging.LogRecord) -> bool:
        return "words count mismatch" not in record.getMessage()


logging.getLogger("phonemizer").addFilter(_PhonemizerWordMismatchFilter())


class _FreshStreamingBlocks(gr.Blocks):
    """Give every generated media stream a run ID that cannot be recycled."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._active_stream_runs: dict[tuple[str, int], int] = {}
        self._next_stream_run = time.time_ns()

    async def handle_streaming_outputs(
        self,
        block_fn: Any,
        data: list[Any],
        session_hash: str | None,
        run: int | None,
        root_path: str | None = None,
        final: bool = False,
    ) -> list[Any]:
        if session_hash is None or run is None:
            return await super().handle_streaming_outputs(
                block_fn,
                data,
                session_hash,
                run,
                root_path,
                final,
            )

        run_key = (session_hash, run)
        stream_run = self._active_stream_runs.get(run_key)
        existing_streams = (
            self.pending_streams[session_hash].get(stream_run)
            if stream_run is not None
            else None
        )
        if existing_streams and all(stream.ended for stream in existing_streams.values()):
            stream_run = None

        if stream_run is None:
            stream_run = self._next_stream_run
            self._next_stream_run += 1
            self._active_stream_runs[run_key] = stream_run

        try:
            return await super().handle_streaming_outputs(
                block_fn,
                data,
                session_hash,
                stream_run,
                root_path,
                final,
            )
        except BaseException:
            self._active_stream_runs.pop(run_key, None)
            raise

    def handle_streaming_diffs(
        self,
        block_fn: Any,
        data: list[Any],
        session_hash: str | None,
        run: int | None,
        final: bool,
        simple_format: bool = False,
    ) -> list[Any]:
        if session_hash is None or run is None:
            return super().handle_streaming_diffs(
                block_fn,
                data,
                session_hash,
                run,
                final,
                simple_format,
            )

        run_key = (session_hash, run)
        stream_run = self._active_stream_runs.get(run_key, run)
        try:
            return super().handle_streaming_diffs(
                block_fn,
                data,
                session_hash,
                stream_run,
                final,
                simple_format,
            )
        finally:
            if final:
                self._active_stream_runs.pop(run_key, None)


MODEL_OPTIONS = {
    "Inflect Nano v2 (4M)": {
        "model_id": "owensong/Inflect-Nano-v2",
        "cache_dir": "inflect-nano-v2",
    },
    "Inflect Micro v2 (9M)": {
        "model_id": "owensong/Inflect-Micro-v2",
        "cache_dir": "inflect-micro-v2",
    },
}
DEFAULT_MODEL = "Inflect Nano v2 (4M)"


class InflectEngine(Protocol):
    deployed_parameters: int
    device: Any
    model: Any
    sample_rate: int

    def _tokens(self, text: str) -> tuple[Any, Any]: ...


_MODEL_ENGINES: dict[str, InflectEngine] = {}
_MODEL_LOCKS = {model_name: threading.Lock() for model_name in MODEL_OPTIONS}
_INFERENCE_MODULE: ModuleType | None = None
_INFERENCE_LOCK = threading.Lock()


def _detect_device() -> str:
    return "cuda" if torch and torch.cuda.is_available() else "cpu"


def get_device_info() -> str:
    return _detect_device()


def _download_snapshot(model_name: str) -> Path:
    if snapshot_download is None:
        raise gr.Error("huggingface_hub is not installed. Install this app's requirements first.")

    model_config = MODEL_OPTIONS[model_name]
    local_dir = Path(__file__).resolve().parent / ".cache" / model_config["cache_dir"]
    return Path(
        snapshot_download(
            model_config["model_id"],
            local_dir=local_dir,
            allow_patterns=[
                "inference.py",
                "inflect_nano_v2_frontend.py",
                "inflect_vits_frontend.py",
                "runtime/**",
                "model.pth",
                "config.json",
            ],
        )
    )


def _load_inference_module(snapshot_dir: Path) -> ModuleType:
    global _INFERENCE_MODULE

    if _INFERENCE_MODULE is not None:
        return _INFERENCE_MODULE

    with _INFERENCE_LOCK:
        if _INFERENCE_MODULE is not None:
            return _INFERENCE_MODULE

        module_name = "_inflect_v2_inference"
        spec = importlib.util.spec_from_file_location(module_name, snapshot_dir / "inference.py")
        if spec is None or spec.loader is None:
            raise gr.Error("Could not load the Inflect v2 runtime from the downloaded snapshot.")

        module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = module
        root_logger = logging.getLogger()
        previous_log_level = root_logger.level
        try:
            spec.loader.exec_module(module)
        finally:
            root_logger.setLevel(previous_log_level)
        _INFERENCE_MODULE = module
        return module


def _get_model_engine(model_name: str) -> InflectEngine:
    if model_name not in MODEL_OPTIONS:
        raise gr.Error("Please select a valid Inflect v2 model.")

    existing = _MODEL_ENGINES.get(model_name)
    if existing is not None:
        return existing

    with _MODEL_LOCKS[model_name]:
        existing = _MODEL_ENGINES.get(model_name)
        if existing is not None:
            return existing

        if torch is None:
            raise gr.Error("PyTorch is not installed. Install this app's requirements first.")

        device_name = _detect_device()
        started = time.perf_counter()
        print(f"Initializing {model_name} on {device_name}")
        snapshot_dir = _download_snapshot(model_name)
        inference = _load_inference_module(snapshot_dir)
        engine = cast(InflectEngine, inference.InflectTTS(snapshot_dir, device=device_name))
        _MODEL_ENGINES[model_name] = engine

        print(
            f"{model_name} initialized. "
            f"Deployed parameters: {engine.deployed_parameters:,}. "
            f"Load time: {time.perf_counter() - started:.2f}s"
        )
        return engine


def _audio_np_to_int16(audio_np: np.ndarray) -> np.ndarray:
    audio_clipped = np.clip(audio_np, -1.0, 1.0)
    return (audio_clipped * 32767.0).astype(np.int16)


def _wav_bytes_from_int16(audio_int16: np.ndarray, sample_rate: int) -> bytes:
    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(sample_rate)
        wav_file.writeframes(audio_int16.tobytes())
    return buffer.getvalue()


def _synthesize_chunk(
    engine: InflectEngine,
    text_chunk: str,
    speed: float,
    variation: float,
    seed: int,
    max_chunk_frames: int | None,
) -> tuple[int, np.ndarray]:
    if torch is None or _INFERENCE_MODULE is None:
        raise gr.Error("The Inflect v2 runtime is not initialized.")

    with torch.inference_mode():
        tokens, lengths = engine._tokens(text_chunk)
        torch.manual_seed(seed)
        if engine.device.type == "cuda":
            torch.cuda.manual_seed_all(seed)

        waveform = engine.model.infer(
            tokens,
            lengths,
            noise_scale=variation,
            noise_scale_w=0.8,
            length_scale=1.0 / speed,
            max_len=max_chunk_frames,
        )[0][0, 0].float().cpu().numpy()

    return engine.sample_rate, _INFERENCE_MODULE.edge_fade(
        np.clip(waveform, -1.0, 1.0),
        engine.sample_rate,
    )


def _apply_pitch_shift(
    audio: np.ndarray,
    sample_rate: int,
    pitch_steps: float,
) -> np.ndarray:
    if abs(pitch_steps) < 0.01:
        return audio

    try:
        import librosa  # type: ignore
    except Exception as exc:  # pragma: no cover
        raise gr.Error(f"Pitch shifting requires librosa: {exc}")

    shifted = librosa.effects.pitch_shift(
        audio,
        sr=sample_rate,
        n_steps=pitch_steps,
        bins_per_octave=12,
        res_type="soxr_hq",
        scale=False,
    )
    return np.asarray(librosa.util.fix_length(shifted, size=audio.size), dtype=np.float32)


def _clear_audio_output() -> None:
    return None


def inflect_tts(
    model_name: str,
    text: str,
    speed: float,
    variation: float,
    pitch_steps: float,
    seed: float | None,
    max_chunk_frames: float | None,
) -> Iterator[bytes]:
    text = (text or "").strip()
    if not text:
        raise gr.Error("Please enter text to synthesize.")
    if not 0.5 <= speed <= 2.0:
        raise gr.Error("Speed must be between 0.5 and 2.0.")
    if not 0.0 <= variation <= 1.0:
        raise gr.Error("Variation must be between 0.0 and 1.0.")
    if not -2.0 <= pitch_steps <= 2.0:
        raise gr.Error("Pitch shift must be between -2 and 2 semitones.")

    seed_value = 0 if seed is None else int(seed)
    max_frames_value = 4000 if max_chunk_frames is None else int(max_chunk_frames)
    if max_frames_value < 0:
        raise gr.Error("Maximum chunk frames must be zero or greater.")
    model_max_len = max_frames_value or None

    try:
        engine = _get_model_engine(model_name)
        print(
            f"Generating speech with {model_name}: {len(text)} chars, "
            f"speed={speed:.2f}, variation={variation:.3f}, "
            f"pitch={pitch_steps:+.2f}, seed={seed_value}, max_frames={model_max_len}"
        )
        if _INFERENCE_MODULE is None:
            raise gr.Error("The Inflect v2 runtime is not initialized.")

        chunks = _INFERENCE_MODULE.split_text(text)
        for index, chunk in enumerate(chunks):
            chunk_started = time.perf_counter()
            sample_rate, audio = _synthesize_chunk(
                engine,
                chunk,
                speed=float(speed),
                variation=float(variation),
                seed=seed_value + index,
                max_chunk_frames=model_max_len,
            )
            audio = _apply_pitch_shift(audio, sample_rate, float(pitch_steps))
            if index < len(chunks) - 1:
                trailing_silence = np.zeros(
                    round(sample_rate * _INFERENCE_MODULE.boundary_pause_seconds(chunk)),
                    dtype=np.float32,
                )
                audio = np.concatenate([audio, trailing_silence])

            elapsed = time.perf_counter() - chunk_started
            print(
                f"{model_name} chunk {index + 1}/{len(chunks)} ready in {elapsed:.2f}s "
                f"({len(chunk)} chars, {audio.size / sample_rate:.2f}s audio)"
            )
            yield _wav_bytes_from_int16(_audio_np_to_int16(audio), sample_rate)
    except gr.Error:
        raise
    except Exception as exc:
        message = str(exc)
        if len(message) > 200:
            message = f"{message[:200]}..."
        raise gr.Error(f"Error during speech generation with {model_name}: {message}")


with _FreshStreamingBlocks() as demo:
    gr.HTML(
        "<h1 style='text-align: center;'>Inflect v2</h1>"
        f"<p style='text-align: center;'>Nano (4M) or Micro (9M), loaded on first use on "
        f"{get_device_info().upper()} | English single-voice 24 kHz TTS</p>"
    )

    with gr.Row(variant="panel"):
        model_name = gr.Dropdown(
            choices=list(MODEL_OPTIONS),
            value=DEFAULT_MODEL,
            label="Model",
            info="Nano prioritizes footprint; Micro prioritizes quality.",
        )
        speed = gr.Slider(
            minimum=0.5,
            maximum=2.0,
            value=1.0,
            step=0.01,
            label="Speed",
            info="Lower is slower; higher is faster.",
        )
        variation = gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.667,
            step=0.001,
            label="Variation",
            info="Lower is steadier; higher adds more latent variation.",
        )

    text_input = gr.Textbox(
        label="Input Text",
        placeholder="Enter the text you want to convert to speech here...",
        value="Wait, are you actually being for real now? I can't believe it!",
        lines=5,
        max_lines=8,
    )

    gr.Examples(
        examples=[[text] for text in EXAMPLE_TEXTS],
        inputs=[text_input],
        label="Examples",
    )

    with gr.Accordion("Advanced Options", open=False):
        with gr.Row():
            pitch_steps = gr.Slider(
                minimum=-2.0,
                maximum=2.0,
                value=0.0,
                step=0.25,
                label="Pitch Shift (Semitones)",
                info="Official v2 post-process control; changes pitch without changing speed.",
            )
            seed = gr.Number(
                value=0,
                precision=0,
                label="Seed",
                info="Use the same integer to repeat a stochastic sample on the same runtime.",
            )
            max_chunk_frames = gr.Number(
                value=4000,
                minimum=0,
                precision=0,
                label="Maximum Chunk Frames",
                info="Model generation cap per text chunk. Use 0 for no cap.",
            )
        gr.Markdown(
            "These are all effective Inflect v2 controls. Duration noise is not shown because "
            "both released checkpoints disable the stochastic duration predictor; speaker, "
            "energy, and native pitch controls are not present in v2."
        )

    generate_btn = gr.Button(
        "Generate Speech",
        variant="primary",
    )

    audio_output = gr.Audio(
        label="Generated Speech",
        streaming=True,
        autoplay=True,
        buttons=["download"],
    )

    generate_inputs = [
        model_name,
        text_input,
        speed,
        variation,
        pitch_steps,
        seed,
        max_chunk_frames,
    ]

    generate_click = generate_btn.click(
        fn=_clear_audio_output,
        outputs=audio_output,
        queue=False,
    )
    generate_click.then(
        fn=inflect_tts,
        inputs=generate_inputs,
        outputs=audio_output,
        api_name="generate_speech",
        concurrency_limit=1,
        concurrency_id="inflect_v2_tts",
    )

    text_submit = text_input.submit(
        fn=_clear_audio_output,
        outputs=audio_output,
        queue=False,
    )
    text_submit.then(
        fn=inflect_tts,
        inputs=generate_inputs,
        outputs=audio_output,
        api_name="generate_speech_enter",
        concurrency_limit=1,
        concurrency_id="inflect_v2_tts",
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(debug=False, theme="Nymbo/Nymbo_Theme")