import logging
import os
import re
import sys
import tempfile
import threading
import time
from dataclasses import dataclass
from fractions import Fraction
from pathlib import Path

import gradio as gr
import numpy as np
import soundfile as sf
import torch
from audiotsm import wsola
from audiotsm.io.array import ArrayReader, ArrayWriter
from fastapi.responses import HTMLResponse, FileResponse
from scipy.signal import resample_poly

try:
    import spaces
except ImportError:  # Local smoke tests do not require the ZeroGPU shim.
    class _Spaces:
        @staticmethod
        def GPU(*args, **kwargs):
            def decorate(function):
                return function
            return decorate
    spaces = _Spaces()


ROOT = Path(__file__).resolve().parent
RUNTIME = ROOT / "runtime"
WEB_RUNTIME = ROOT / "web_runtime"
STATIC_URL = "/web_runtime"
sys.path.insert(0, str(RUNTIME))
sys.path.insert(0, str(ROOT))

from gradio import Server
from gradio.data_classes import FileData

# Register ONNX Runtime / model asset MIME types. Windows (and some containers)
# do not map these by default, and FileResponse would otherwise serve the
# browser runtime's modules with a wrong content type.
import mimetypes
mimetypes.add_type("text/javascript", ".mjs")
mimetypes.add_type("application/wasm", ".wasm")
mimetypes.add_type("application/octet-stream", ".onnx")
mimetypes.add_type("application/json", ".config.json")

app = Server()

import commons
import utils
from inflect_vits_frontend import run_vits_frontend
from models import SynthesizerTrn
from text import cleaned_text_to_sequence
from text.symbols import symbols


@dataclass(frozen=True)
class ModelSpec:
    label: str
    short_label: str
    params: str
    checkpoint: Path
    config: Path
    revision: str


SPECS = {
    "Inflect Micro v2": ModelSpec(
        label="Inflect Micro v2",
        short_label="Micro",
        params="9.36M",
        checkpoint=ROOT / "models" / "micro" / "model.pth",
        config=ROOT / "models" / "micro" / "config.json",
        revision="3eede065",
    ),
    "Inflect Nano v2": ModelSpec(
        label="Inflect Nano v2",
        short_label="Nano",
        params="3.96M",
        checkpoint=ROOT / "models" / "nano" / "model.pth",
        config=ROOT / "models" / "nano" / "config.json",
        revision="bfca4684",
    ),
}


def split_text(text: str, limit: int = 280) -> list[str]:
    normalized = " ".join(text.split())
    sentences = [
        part.strip()
        for part in re.split(r"(?<=[.!?;:])\s+", normalized)
        if part.strip()
    ]
    chunks: list[str] = []
    for sentence in sentences or [normalized]:
        while len(sentence) > limit:
            search = sentence[: limit + 1]
            punctuation = max(search.rfind(mark) for mark in (",", ";", ":"))
            split_at = (
                punctuation + 1
                if punctuation >= limit // 2
                else sentence.rfind(" ", 0, limit + 1)
            )
            if split_at < limit // 2:
                split_at = limit
            chunks.append(sentence[:split_at].strip())
            sentence = sentence[split_at:].strip()
        if sentence:
            chunks.append(sentence)
    return chunks


def boundary_pause_seconds(chunk: str) -> float:
    ending = chunk.rstrip()[-1:] if chunk.strip() else ""
    return {
        "?": 0.28,
        "!": 0.24,
        ".": 0.22,
        ";": 0.16,
        ":": 0.13,
        ",": 0.09,
    }.get(ending, 0.08)


def edge_fade(waveform: np.ndarray, sample_rate: int, milliseconds: float = 5.0) -> np.ndarray:
    frames = min(round(sample_rate * milliseconds / 1000.0), waveform.size // 2)
    if frames <= 0:
        return waveform
    output = waveform.copy()
    ramp = np.linspace(0.0, 1.0, frames, endpoint=True, dtype=np.float32)
    output[:frames] *= ramp
    output[-frames:] *= ramp[::-1]
    return output


def pcm16(waveform: np.ndarray) -> np.ndarray:
    """Return the exact PCM format Gradio writes, without its native conversion path."""
    clipped = np.clip(waveform, -1.0, 1.0)
    return np.ascontiguousarray(np.rint(clipped * 32767.0), dtype=np.int16)


def pitch_shift_speech(
    waveform: np.ndarray,
    semitones: float,
    frame_length: int = 1024,
) -> np.ndarray:
    """Shift speech pitch without changing duration using WSOLA and resampling."""
    waveform = np.asarray(waveform, dtype=np.float32)
    if waveform.size < frame_length or abs(semitones) < 0.01:
        return waveform

    factor = 2.0 ** (float(semitones) / 12.0)
    target_stretched_length = max(frame_length, round(len(waveform) * factor))

    # Padding lets WSOLA flush its final overlap without truncating speech.
    reader = ArrayReader(np.pad(waveform, (0, frame_length * 4))[None, :])
    writer = ArrayWriter(1)
    wsola(
        1,
        speed=1.0 / factor,
        frame_length=frame_length,
        tolerance=frame_length // 4,
    ).run(reader, writer)

    stretched = writer.data[0]
    if len(stretched) < target_stretched_length:
        stretched = np.pad(
            stretched,
            (0, target_stretched_length - len(stretched)),
        )
    else:
        stretched = stretched[:target_stretched_length]

    ratio = Fraction(factor).limit_denominator(1000)
    shifted = resample_poly(stretched, ratio.denominator, ratio.numerator)
    if len(shifted) < len(waveform):
        shifted = np.pad(shifted, (0, len(waveform) - len(shifted)))
    return np.asarray(shifted[: len(waveform)], dtype=np.float32)


class Engine:
    def __init__(self, spec: ModelSpec) -> None:
        if not spec.config.is_file():
            raise FileNotFoundError(f"Missing release configuration for {spec.label}")
        if not spec.checkpoint.is_file():
            raise FileNotFoundError(f"Missing release weights for {spec.label}")
        self.spec = spec
        self.hps = utils.get_hparams_from_file(str(spec.config))
        self.model = SynthesizerTrn(
            len(symbols),
            self.hps.data.filter_length // 2 + 1,
            self.hps.train.segment_size // self.hps.data.hop_length,
            **self.hps.model,
        ).eval()
        root_logger = logging.getLogger()
        previous_level = root_logger.level
        try:
            root_logger.setLevel(logging.WARNING)
            utils.load_checkpoint(str(spec.checkpoint), self.model, None)
        finally:
            root_logger.setLevel(previous_level)
        self.device = torch.device("cpu")
        self.sample_rate = int(self.hps.data.sampling_rate)
        self.lock = threading.Lock()

    def tokens(self, text: str) -> tuple[torch.Tensor, torch.Tensor]:
        phonemes = run_vits_frontend(text).phoneme_text
        sequence = cleaned_text_to_sequence(phonemes)
        if self.hps.data.add_blank:
            sequence = commons.intersperse(sequence, 0)
        if not sequence:
            raise ValueError("The phoneme frontend produced no speakable tokens.")
        tokens = torch.LongTensor(sequence).to(self.device).unsqueeze(0)
        lengths = torch.LongTensor([tokens.size(1)]).to(self.device)
        return tokens, lengths

    @torch.inference_mode()
    def synthesize(
        self,
        text: str,
        speed: float,
        variation: float,
        pitch_steps: float,
        seed: int,
    ) -> tuple[int, np.ndarray]:
        chunks = split_text(text)
        waveforms: list[np.ndarray] = []
        with self.lock:
            self.device = torch.device("cuda")
            self.model.to(self.device)
            try:
                for index, chunk in enumerate(chunks):
                    if index:
                        waveforms.append(
                            np.zeros(
                                round(self.sample_rate * boundary_pause_seconds(chunks[index - 1])),
                                dtype=np.float32,
                            )
                        )
                    tokens, lengths = self.tokens(chunk)
                    torch.manual_seed(seed + index)
                    torch.cuda.manual_seed_all(seed + index)
                    waveform = self.model.infer(
                        tokens,
                        lengths,
                        noise_scale=variation,
                        noise_scale_w=0.8,
                        length_scale=1.0 / speed,
                        max_len=4000,
                    )[0][0, 0].float().cpu().numpy()
                    waveforms.append(edge_fade(waveform, self.sample_rate))
            finally:
                self.model.to("cpu")
                self.device = torch.device("cpu")
                torch.cuda.empty_cache()
        audio = np.concatenate(waveforms)
        if abs(pitch_steps) >= 0.01:
            audio = pitch_shift_speech(audio, pitch_steps)
        return self.sample_rate, pcm16(audio)


ENGINES = {label: Engine(spec) for label, spec in SPECS.items()}


def validate(
    text: str,
    speed: float,
    variation: float,
    pitch_steps: float,
    seed: int,
) -> tuple[str, float, float, float, int]:
    text = " ".join((text or "").split())
    if not text:
        raise gr.Error("Enter something for Inflect to say.")
    return text, float(speed), float(variation), float(pitch_steps), int(seed)


def render_wav(sample_rate: int, samples: np.ndarray) -> FileData:
    """Persist PCM16 samples to a WAV and return a Gradio FileData handle.

    The return-type annotation is what tells gradio.Server which output
    component to serialize, so the JS client receives a {url, ...} blob
    rather than the function's return value being dropped.
    """
    handle, out_path = tempfile.mkstemp(suffix=".wav", prefix="inflect_")
    os.close(handle)
    sf.write(out_path, samples, sample_rate, subtype="PCM_16")
    return FileData(path=out_path)


@app.api(concurrency_limit=1)
@spaces.GPU(duration=120)
def synthesize(
    text: str,
    model_name: str,
    speed: float,
    variation: float,
    pitch_steps: float,
    seed: int,
) -> tuple[FileData, str]:
    text, speed, variation, pitch_steps, seed = validate(
        text, speed, variation, pitch_steps, seed
    )
    engine = ENGINES[model_name]
    started = time.perf_counter()
    sample_rate, samples = engine.synthesize(text, speed, variation, pitch_steps, seed)
    seconds = len(samples) / sample_rate
    wall = time.perf_counter() - started
    rtf = wall / seconds
    chunks = len(split_text(text))
    status = (
        f"**{model_name}** · fixed English voice  \n"
        f"`{engine.spec.params}` parameters · `{seconds:.2f}s` audio · "
        f"`{wall:.2f}s` generation · `RTF {rtf:.3f}` · "
        f"`{len(text)}` characters · `{chunks}` chunk{'s' if chunks != 1 else ''} · "
        f"pitch `{pitch_steps:+.2f} st` · "
        f"seed `{seed}` · weights `{engine.spec.revision}`"
    )
    return render_wav(sample_rate, samples), status


@app.api(concurrency_limit=1)
@spaces.GPU(duration=120)
def compare(
    text: str,
    speed: float,
    variation: float,
    pitch_steps: float,
    seed: int,
) -> tuple[FileData, FileData, str]:
    text, speed, variation, pitch_steps, seed = validate(
        text, speed, variation, pitch_steps, seed
    )
    started = time.perf_counter()
    micro_sr, micro_pcm = ENGINES["Inflect Micro v2"].synthesize(
        text, speed, variation, pitch_steps, seed
    )
    micro_wall = time.perf_counter() - started
    started = time.perf_counter()
    nano_sr, nano_pcm = ENGINES["Inflect Nano v2"].synthesize(
        text, speed, variation, pitch_steps, seed
    )
    nano_wall = time.perf_counter() - started
    micro_seconds = len(micro_pcm) / micro_sr
    nano_seconds = len(nano_pcm) / nano_sr
    status = (
        f"Same text and seed `{seed}` · "
        f"Micro `{micro_wall:.2f}s / {micro_seconds:.2f}s audio` · "
        f"Nano `{nano_wall:.2f}s / {nano_seconds:.2f}s audio`"
    )
    return render_wav(micro_sr, micro_pcm), render_wav(nano_sr, nano_pcm), status


@app.middleware("http")
async def _cross_origin_isolation(request, call_next):
    """Enable crossOriginIsolated so the browser ONNX runtime can use WebGPU.

    COEP `credentialless` (rather than `require-corp`) keeps the Gradio JS
    client module from the CDN loadable while still enabling shared memory.
    """
    response = await call_next(request)
    response.headers["Cross-Origin-Opener-Policy"] = "same-origin"
    response.headers["Cross-Origin-Embedder-Policy"] = "credentialless"
    return response


@app.get(STATIC_URL + "/{path:path}")
async def serve_web_runtime(path: str):
    full = (WEB_RUNTIME / path).resolve()
    if not full.is_file():
        return HTMLResponse("Not found", status_code=404)
    return FileResponse(str(full))


@app.get("/")
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(f.read())


if __name__ == "__main__":
    app.launch(show_error=True)