#!/usr/bin/env python3
# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""Qwen3-TTS FastAPI Server for UmbraAudio.

Provides REST endpoints for voice design and voice cloning.
Expected by umbraaudio.voices.qwen_tts.QwenTTSEngine.
"""

from __future__ import annotations

import asyncio
import io
import json
import logging
import os
import time

import anyio
import numpy as np
import soundfile as sf
import torch
from dataclasses import asdict
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import Response
from pathlib import Path
from pydantic import BaseModel
from qwen_tts import Qwen3TTSModel
from qwen_tts.inference.qwen3_tts_model import VoiceClonePromptItem

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Directory where precomputed voice-clone prompts (the "8k matrix" ref_code +
# speaker embedding) are cached, one .pt file per speaker slug. Populated by
# POST /v1/voice-prompt and consumed by /v1/tts via the prompt_slug field so
# every utterance for a speaker reuses the exact same voice identity.
PROMPT_DIR = Path(
    os.getenv(
        "QWEN_PROMPT_DIR",
        "/Volumes/1TB_APL_Storage/UmbraAudio/DealHunterShip_4.0e_whole-to-audio/.umbra/voice_prompts",
    )
)

app = FastAPI(title="Qwen3-TTS Server", version="0.1.0")


async def _watchdog_loop() -> None:
    """Background task: if a generation has been stuck past MAX_GEN, restart."""
    while True:
        await asyncio.sleep(min(_MAX_GEN, 30))
        s = _gen_start_at
        if s is not None and (time.monotonic() - s) > _MAX_GEN:
            _restart_server(
                f"generation hung >{_MAX_GEN:.0f}s (poisoned CUDA context?)"
            )
            await asyncio.sleep(3600)  # never reached after execv


@app.on_event("startup")
async def _start_watchdog() -> None:
    app.state._wd_task = asyncio.create_task(_watchdog_loop())

# Rainbow Passage — used as default ref_text for ICL voice cloning mode.
# All voice refs were generated with this text via Edge-TTS.
VOICE_REFERENCE_TEXT = (
    "When the sunlight strikes raindrops in the air, they act as a prism and form a rainbow. "
    "The rainbow is a division of white light into many beautiful colors. "
    "These take the shape of a long round arch, with its path high above, "
    "and its two ends apparently beyond the horizon. "
    "There is, according to legend, a boiling pot of gold at one end. "
    "People look, but no one ever finds it. "
    "When a man looks for something beyond his reach, "
    "his friends say he is looking for the pot of gold at the end of the rainbow. "
    "When the sunlight strikes raindrops in the air, they act as a prism and form a rainbow. "
    "The rainbow is a division of white light into many beautiful colors. "
    "These take the shape of a long round arch, with its path high above, "
    "and its two ends apparently beyond the horizon. "
    "There is, according to legend, a boiling pot of gold at one end. "
    "People look, but no one ever finds it. "
    "When a man looks for something beyond his reach, "
    "his friends say he is looking for the pot of gold at the end of the rainbow."
)

# Global model state
_vd_model: Qwen3TTSModel | None = None
_base_model: Qwen3TTSModel | None = None
_current_model: str | None = None  # "voicedesign" or "base"

# Serializes all model access. Model loads and generations are blocking CPU/GPU
# calls; they run in worker threads (anyio.to_thread) under this lock so the
# event loop stays responsive and concurrent /v1/tts requests cannot wedge it.
_model_lock = asyncio.Lock()

# Hang watchdog. On CUDA, a kernel error (e.g. OOM) can poison the context and
# wedge the generation thread forever while the event loop stays responsive
# (health works, every /v1/tts hangs). The only safe recovery is to restart the
# process. We track when a generation started and, if it exceeds MAX_GEN_SECONDS,
# we relaunch ourselves via os.execv to get a clean CUDA context.
_MAX_GEN = float(os.getenv("QWEN_MAX_GEN_S", "1800"))
_gen_start_at: float | None = None  # time.monotonic() when a generation began


def _mark_gen_started() -> None:
    global _gen_start_at
    _gen_start_at = time.monotonic()


def _mark_gen_done() -> None:
    global _gen_start_at
    _gen_start_at = None


def _restart_server(reason: str) -> None:
    """Relaunch the current process (os.execv) to clear a wedged GPU context."""
    import sys as _sys

    logger.warning("Restarting TTS server: %s", reason)
    try:
        for s in (_sys.stdout, _sys.stderr):
            try:
                s.flush()
            except Exception:
                pass
    except Exception:
        pass
    try:
        _sys.stdout = open(os.devnull, "w")
        _sys.stderr = open(os.devnull, "w")
        _sys.stdin = open(os.devnull)
        os.execv(
            _sys.executable,
            [_sys.executable, os.path.abspath(__file__), *list(_sys.argv[1:])],
        )
    except Exception as e:  # pragma: no cover
        logger.error("os.execv failed: %s", e)
        raise SystemExit(1)

# Model paths from environment
VOICEDESIGN_PATH = os.getenv(
    "QWEN_VOICEDESIGN_MODEL",
    "/Volumes/1TB_APL_Storage/UmbraAudio/models/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
)
BASE_PATH = os.getenv(
    "QWEN_BASE_MODEL",
    "/Volumes/1TB_APL_Storage/UmbraAudio/models/Qwen3-TTS-12Hz-1.7B-Base",
)
DEVICE = os.getenv("QWEN_DEVICE", "mps")
DTYPE = torch.bfloat16 if DEVICE in ("cuda", "mps") else torch.float32
# Attention implementation. "eager" is very slow for Qwen3-TTS autoregressive
# codec decoding but is the only one that reliably runs create_voice_clone_prompt
# (prompt extraction). "sdpa" (built into torch) is far faster for generation.
# Runtime-switchable via POST /attn (reloads the model); this is read at model
# load time, so switching reloads the active model with the new implementation.
_ATTN_IMPL: str = os.getenv("QWEN_ATTN", "sdpa")


class LoadRequest(BaseModel):
    model_type: str  # "voicedesign" or "base"


class GenerateRequest(BaseModel):
    text: str
    sample_rate: int = 24000
    voice_description: str | None = None


def _load_model(model_type: str) -> Qwen3TTSModel:
    """Load a model, unloading the other if needed."""
    global _vd_model, _base_model, _current_model

    if _current_model == model_type:
        return _vd_model if model_type == "voicedesign" else _base_model

    # Unload current
    if _current_model == "voicedesign" and _vd_model is not None:
        del _vd_model
        _vd_model = None
    elif _current_model == "base" and _base_model is not None:
        del _base_model
        _base_model = None

    if torch.mps.is_available() or torch.cuda.is_available():
        torch.mps.empty_cache() if torch.mps.is_available() else torch.cuda.empty_cache()

    # Load requested
    path = VOICEDESIGN_PATH if model_type == "voicedesign" else BASE_PATH
    logger.info(f"Loading {model_type} model from {path}")

    model = Qwen3TTSModel.from_pretrained(
        path,
        device_map=DEVICE,
        dtype=DTYPE,
        attn_implementation=_ATTN_IMPL,
    )

    if model_type == "voicedesign":
        _vd_model = model
    else:
        _base_model = model

    _current_model = model_type
    logger.info(f"{model_type} model loaded successfully")
    return model


def _prompt_path(slug: str) -> Path:
    safe = Path(slug).name
    PROMPT_DIR.mkdir(parents=True, exist_ok=True)
    return PROMPT_DIR / f"{safe}.pt"


def _save_prompt(slug: str, items: list[VoiceClonePromptItem]) -> Path:
    path = _prompt_path(slug)
    payload = {"items": [asdict(it) for it in items]}
    torch.save(payload, str(path))
    logger.info("Saved ICL voice prompt for %s to %s", slug, path)
    return path


def _load_prompt_items(slug: str) -> list[VoiceClonePromptItem] | None:
    path = _prompt_path(slug)
    if not path.is_file():
        return None
    try:
        payload = torch.load(str(path), map_location="cpu", weights_only=True)
    except Exception as e:
        logger.warning("Failed to load prompt %s: %s", slug, e)
        return None
    raw_items = payload.get("items", [])
    items: list[VoiceClonePromptItem] = []
    for d in raw_items:
        ref_code = d.get("ref_code")
        if ref_code is not None and not torch.is_tensor(ref_code):
            ref_code = torch.tensor(ref_code)
        ref_spk = d.get("ref_spk_embedding")
        if ref_spk is None:
            logger.warning("Prompt %s missing ref_spk_embedding", slug)
            return None
        if not torch.is_tensor(ref_spk):
            ref_spk = torch.tensor(ref_spk)
        items.append(
            VoiceClonePromptItem(
                ref_code=ref_code,
                ref_spk_embedding=ref_spk,
                x_vector_only_mode=bool(d.get("x_vector_only_mode", False)),
                icl_mode=bool(d.get("icl_mode", not bool(d.get("x_vector_only_mode", False)))),
                ref_text=d.get("ref_text"),
            )
        )
    return items


@app.post("/v1/voice-prompt")
async def build_voice_prompt(
    slug: str = Form(...),
    ref_text: str | None = Form(None),
    ref_audio: UploadFile | None = File(None),
    x_vector_only: int = Form(0),
    force: int = Form(0),
):
    """Build and cache a speaker's vmatrix (digital voice) once.

    The vmatrix is the set of numbers inside Qwen3-TTS that IS the voice. Two capture
    modes (a speaker's cached prompt is reused on every subsequent utterance via
    prompt_slug):

      * x_vector_only=0 (DEFAULT, canonical, "full vmatrix"): extracts the FULL
        ICL codec-token matrix (ref_code, the "8k matrix") + speaker embedding.
        Generation conditions on the exact reference identity, so every utterance
        is faithful AND identical in identity across the whole pipeline.
      * x_vector_only=1: captures ONLY the coarse speaker embedding (ref_code=None).
        Legacy/experimental - caused the historical "voice changed between
        chapters" inconsistency. Not used for production consistency.

    Idempotent: if `<slug>.pt` is already cached it returns without rebuilding,
    unless `force=1` is passed (use force to switch an existing prompt mode, e.g.
    after an approved reference WAV is replaced).
    """
    path = _prompt_path(slug)
    if path.exists() and not force:
        return {"ok": True, "slug": slug, "cached": True, "path": str(path)}

    if ref_audio is None:
        raise HTTPException(400, "ref_audio is required to build a voice prompt")
    ref_bytes = await ref_audio.read()

    ref_text = (ref_text or "").strip() or VOICE_REFERENCE_TEXT
    xvec = bool(x_vector_only)

    def _build() -> list[VoiceClonePromptItem]:
        model = _load_model("base")
        ref_audio_array, ref_sr = sf.read(io.BytesIO(ref_bytes), dtype="float32")
        if ref_sr != 24000:
            import librosa
            ref_audio_array = librosa.resample(
                ref_audio_array, orig_sr=ref_sr, target_sr=24000
            )
        items = model.create_voice_clone_prompt(
            ref_audio=(ref_audio_array, 24000),
            ref_text=ref_text,
            x_vector_only_mode=xvec,
        )
        return items

    _mark_gen_started()
    async with _model_lock:
        items = await anyio.to_thread.run_sync(_build)
    _mark_gen_done()

    _save_prompt(slug, items)
    return {"ok": True, "slug": slug, "path": str(_prompt_path(slug))}


@app.get("/health")
async def health():
    return {
        "status": "ok",
        "voicedesign_loaded": _vd_model is not None,
        "base_loaded": _base_model is not None,
        "current_model": _current_model,
        "device": DEVICE,
    }


@app.post("/load")
async def load_model(request: LoadRequest):
    if request.model_type not in ("voicedesign", "base"):
        raise HTTPException(400, "model_type must be 'voicedesign' or 'base'")

    try:
        async with _model_lock:
            await anyio.to_thread.run_sync(_load_model, request.model_type)
        return {"ok": True, "model": request.model_type}
    except Exception as e:
        logger.exception("Failed to load model")
        raise HTTPException(500, f"Failed to load model: {e}") from e


def _unload_model(model_type: str) -> None:
    global _vd_model, _base_model, _current_model

    if model_type in ("voicedesign", "all") and _vd_model is not None:
        del _vd_model
        _vd_model = None
    if model_type in ("base", "all") and _base_model is not None:
        del _base_model
        _base_model = None

    if model_type == "all" or _current_model == model_type:
        _current_model = None

    if torch.mps.is_available() or torch.cuda.is_available():
        torch.mps.empty_cache() if torch.mps.is_available() else torch.cuda.empty_cache()


@app.post("/unload")
async def unload_model(request: LoadRequest):
    if request.model_type not in ("voicedesign", "base", "all"):
        raise HTTPException(400, "model_type must be 'voicedesign', 'base', or 'all'")

    async with _model_lock:
        await anyio.to_thread.run_sync(_unload_model, request.model_type)

    return {"ok": True, "unloaded": request.model_type}


class AttnRequest(BaseModel):
    impl: str  # "eager", "sdpa", "fluent", "flash_attention_2"


@app.get("/attn")
async def get_attn():
    return {"ok": True, "impl": _ATTN_IMPL}


@app.post("/attn")
async def set_attn(request: AttnRequest):
    """Switch the attention implementation and reload the active model.

    The attention implementation is fixed at model-load time, so this unloads
    the current model(s) and reloads `base` with the new implementation. Callers
    should switch to ``eager`` for prompt extraction (Phase 1) and ``sdpa`` for
    generation (Phase 2) without a process restart.
    """
    global _ATTN_IMPL
    impl = (request.impl or "").strip().lower()
    if impl not in ("eager", "sdpa", "fluent", "flash_attention_2"):
        raise HTTPException(400, "impl must be eager|sdpa|fluent|flash_attention_2")

    async with _model_lock:
        try:
            await anyio.to_thread.run_sync(_unload_model, "all")
            _ATTN_IMPL = impl
            # reload so subsequent use has the new attn; base is the main model
            await anyio.to_thread.run_sync(_load_model, "base")
        except Exception as e:
            logger.exception("Failed to switch attn to %s", impl)
            raise HTTPException(500, f"Failed to switch attn: {e}") from e

    return {"ok": True, "impl": _ATTN_IMPL}


@app.post("/v1/tts")
async def generate_tts(
    text: str = Form(...),
    sample_rate: int = Form(24000),
    # Was hardcoded to "english" in every branch, so German and French text was
    # synthesised with English phonology and prosody - an English-accented
    # foreign voice. Defaults to "english" so the existing English pipeline is
    # byte-for-byte unchanged; multilingual callers pass their own language.
    # "Auto" is also accepted and lets the model detect it.
    language: str = Form("english"),
    voice_description: str | None = Form(None),
    ref_audio: UploadFile | None = File(None),
    ref_text: str | None = Form(None),
    speaker_embedding: str | None = Form(None),
    prompt_slug: str | None = Form(None),
    temperature: float = Form(0.9),
    do_sample: bool = Form(True),
    seed: int | None = Form(None),
):
    """Generate TTS audio.

    - voice_description + text -> VoiceDesign generation
    - ref_audio + text + ref_text -> VoiceClone generation
    - speaker_embedding + text -> Direct embedding generation (fastest, most consistent)
    - prompt_slug -> ICL generation from cached voice-clone prompt (most faithful &
      consistent; the "8k matrix" built by POST /v1/voice-prompt)

    prompt_slug takes priority: if the cached prompt exists it is used, giving
    the speaker the identical reference voice on every utterance.

    voice_description also combines with a clone/ICL prompt: when ref_audio or
    prompt_slug is given, the instruction text is prepended as style direction
    (emotion, pace, delivery) while the cloned speaker identity is preserved.
    """
    gen_kwargs: dict = {"temperature": temperature, "do_sample": do_sample}

    if seed is not None:
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)

    _mark_gen_started()
    async with _model_lock:
        if prompt_slug:
            items = _load_prompt_items(prompt_slug)
            if items is None:
                raise HTTPException(
                    404, f"No cached voice prompt for slug '{prompt_slug}'"
                )

            def _icl() -> tuple[np.ndarray, int]:
                model = _load_model("base")
                kw = dict(gen_kwargs)
                if voice_description:
                    kw["instruct_ids"] = [
                        model._tokenize_texts(
                            [model._build_instruct_text(voice_description)]
                        )[0]
                    ]
                wavs, sr = model.generate_voice_clone(
                    text=text,
                    voice_clone_prompt=items,
                    language=language,
                    **kw,
                )
                return wavs, sr

            wavs, sr = await anyio.to_thread.run_sync(_icl)
        elif speaker_embedding is not None:
            emb_list = json.loads(speaker_embedding)
            emb_np = np.array(emb_list, dtype=np.float32)
            spk_tensor = torch.from_numpy(emb_np)

            def _embed() -> tuple[np.ndarray, int]:
                model = _load_model("base")
                device = next(model.model.parameters()).device
                dtype = next(model.model.parameters()).dtype
                spk = spk_tensor.to(device).to(dtype)

                wavs, sr = model.generate_voice_clone(
                    text=text,
                    voice_clone_prompt={
                        "ref_code": [None],
                        "ref_spk_embedding": [spk],
                        "x_vector_only_mode": [True],
                        "icl_mode": [False],
                    },
                    language=language,
                    **gen_kwargs,
                )
                return wavs, sr

            wavs, sr = await anyio.to_thread.run_sync(_embed)
        elif ref_audio is not None:
            ref_bytes = await ref_audio.read()

            def _clone() -> tuple[np.ndarray, int]:
                model = _load_model("base")
                ref_audio_array, ref_sr = sf.read(io.BytesIO(ref_bytes), dtype="float32")

                if ref_sr != sample_rate:
                    import librosa

                    ref_audio_array = librosa.resample(
                        ref_audio_array, orig_sr=ref_sr, target_sr=sample_rate
                    )

                kw = dict(gen_kwargs)
                if voice_description:
                    kw["instruct_ids"] = [
                        model._tokenize_texts(
                            [model._build_instruct_text(voice_description)]
                        )[0]
                    ]
                wavs, sr = model.generate_voice_clone(
                    text=text,
                    ref_audio=(ref_audio_array, sample_rate),
                    ref_text=None,
                    x_vector_only_mode=True,
                    language=language,
                    **kw,
                )
                return wavs, sr

            wavs, sr = await anyio.to_thread.run_sync(_clone)
        elif voice_description:
            def _design() -> tuple[np.ndarray, int]:
                model = _load_model("voicedesign")
                wavs, sr = model.generate_voice_design(
                    text=text,
                    instruct=voice_description,
                    language=language,
                    **gen_kwargs,
                )
                return wavs, sr

            wavs, sr = await anyio.to_thread.run_sync(_design)
        else:
            raise HTTPException(400, "Either voice_description or ref_audio required")

    _mark_gen_done()
    if not wavs:
        raise HTTPException(500, "Model returned no audio")

    # Return WAV as binary
    wav_bytes = io.BytesIO()
    sf.write(wav_bytes, wavs[0], sr, format="WAV")
    wav_bytes.seek(0)

    return Response(
        content=wav_bytes.read(),
        media_type="audio/wav",
        headers={"Content-Disposition": "attachment; filename=output.wav"},
    )


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8881)