import io
import os
import threading
import time
from contextlib import asynccontextmanager

import torch
import torchaudio
import onnxruntime
import uvicorn
from cosyvoice.cli.cosyvoice import AutoModel
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel, Field


class SpeechRequest(BaseModel):
    model: str | None = None
    input: str = Field(min_length=1)
    voice: str | None = None
    response_format: str = "wav"
    speed: float = 1.0
    instructions: str | None = None
    instruct: str | None = None


class Runtime:
    def __init__(self) -> None:
        self.model = None
        self.loaded_at = None
        self.speech_tokenizer_providers = None
        self.error = None
        self.lock = threading.Lock()

    def load(self) -> None:
        try:
            if not torch.cuda.is_available():
                raise RuntimeError("CUDA is unavailable")
            device = os.environ["VOICE_DEVICE"]
            if device != "cuda:0":
                raise RuntimeError("VOICE_DEVICE must be cuda:0")
            if not env_bool("VOICE_FP16"):
                raise RuntimeError("VOICE_FP16 must remain enabled for core CUDA inference")
            torch.cuda.set_device(int(device.split(":", 1)[1]))
            original_inference_session = onnxruntime.InferenceSession
            speech_tokenizer_provider = os.environ["VOICE_SPEECH_TOKENIZER_PROVIDER"]
            if speech_tokenizer_provider != "CPUExecutionProvider":
                raise RuntimeError("VOICE_SPEECH_TOKENIZER_PROVIDER must be CPUExecutionProvider")

            def inference_session(*args, **kwargs):
                if kwargs.get("providers") == ["CUDAExecutionProvider"]:
                    kwargs["providers"] = [speech_tokenizer_provider]
                return original_inference_session(*args, **kwargs)

            onnxruntime.InferenceSession = inference_session
            try:
                self.model = AutoModel(
                    model_dir=os.environ["MODEL_DIR"],
                    load_trt=env_bool("VOICE_LOAD_TRT"),
                    load_vllm=env_bool("VOICE_LOAD_VLLM"),
                    fp16=env_bool("VOICE_FP16"),
                )
            finally:
                onnxruntime.InferenceSession = original_inference_session
            self.speech_tokenizer_providers = self.model.frontend.speech_tokenizer_session.get_providers()
            if self.speech_tokenizer_providers != [speech_tokenizer_provider]:
                raise RuntimeError(f"speech tokenizer provider mismatch: {self.speech_tokenizer_providers}")
            self.loaded_at = time.time()
        except Exception as error:
            self.error = f"{type(error).__name__}: {error}"
            raise


def env_bool(name: str) -> bool:
    value = os.environ[name].lower()
    if value not in {"true", "false"}:
        raise RuntimeError(f"{name} must be true or false")
    return value == "true"


runtime = Runtime()


@asynccontextmanager
async def lifespan(_: FastAPI):
    runtime.load()
    yield


app = FastAPI(title="SelfMedia Voice", version="1", lifespan=lifespan)

INSTRUCTION_END_TOKEN = "<|endofprompt|>"


def normalize_instruction(value: str | None) -> str | None:
    if value is None or not value.strip():
        return None
    normalized = value.strip()
    if INSTRUCTION_END_TOKEN in normalized:
        return normalized
    return f"You are a helpful assistant. {normalized}{INSTRUCTION_END_TOKEN}"


@app.get("/health")
def health() -> dict:
    return {
        "ok": runtime.model is not None and runtime.error is None,
        "model": os.environ["MODEL_ID"],
        "modelRef": os.environ["MODEL_REF"],
        "sourceRef": os.environ["SOURCE_REF"],
        "torch": torch.__version__,
        "torchaudio": torchaudio.__version__,
        "cuda": torch.version.cuda,
        "gpu": torch.cuda.get_device_name(torch.cuda.current_device()) if torch.cuda.is_available() else None,
        "coreInferenceDevice": os.environ["VOICE_DEVICE"],
        "coreInferenceFp16": env_bool("VOICE_FP16"),
        "loadTrt": env_bool("VOICE_LOAD_TRT"),
        "loadVllm": env_bool("VOICE_LOAD_VLLM"),
        "workers": int(os.environ["VOICE_WORKERS"]),
        "concurrency": os.environ["VOICE_CONCURRENCY"],
        "speechTokenizerProviders": runtime.speech_tokenizer_providers,
        "loadedAt": runtime.loaded_at,
        "error": runtime.error,
    }


@app.post("/v1/audio/speech")
def speech(request: SpeechRequest) -> Response:
    if request.response_format.lower() != "wav":
        raise HTTPException(status_code=400, detail="response_format must be wav")
    if not float(os.environ["VOICE_SPEED_MIN"]) <= request.speed <= float(os.environ["VOICE_SPEED_MAX"]):
        raise HTTPException(status_code=400, detail="speed is outside the configured range")
    if runtime.model is None:
        raise HTTPException(status_code=503, detail=runtime.error or "model is not ready")
    if request.voice not in {None, os.environ["REFERENCE_VOICE_ID"]}:
        raise HTTPException(status_code=400, detail="voice must select the configured reference voice")

    started = time.perf_counter()
    instruction = normalize_instruction(request.instructions or request.instruct)
    with runtime.lock:
        inference = (
            runtime.model.inference_instruct2(
                request.input,
                instruction,
                os.environ["PROMPT_WAV"],
                stream=False,
                speed=request.speed,
            )
            if instruction
            else runtime.model.inference_zero_shot(
                request.input,
                os.environ["PROMPT_TEXT"],
                os.environ["PROMPT_WAV"],
                stream=False,
                speed=request.speed,
            )
        )
        chunks = [
            output["tts_speech"].cpu()
            for output in inference
        ]
    if not chunks:
        raise HTTPException(status_code=500, detail="empty inference result")
    waveform = torch.cat(chunks, dim=1)
    output = io.BytesIO()
    torchaudio.save(output, waveform, runtime.model.sample_rate, format="wav")
    return Response(
        content=output.getvalue(),
        media_type="audio/wav",
        headers={
            "X-Inference-Seconds": f"{time.perf_counter() - started:.6f}",
            "X-Sample-Rate": str(runtime.model.sample_rate),
            "X-Inference-Mode": "instruct2" if instruction else "zero-shot",
        },
    )


if __name__ == "__main__":
    workers = int(os.environ["VOICE_WORKERS"])
    if workers != 1:
        raise RuntimeError("VOICE_WORKERS must be 1 for the single-GPU model runtime")
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ["VOICE_PORT"]), workers=workers)