"""LFM3-Audio-2B as an optional second engine of the Kokoro Space (for A/B tests).

Off unless the Space variable LFM3_ENABLED=1. When on, the routes below sit
beside the Kokoro ones with the same contract (a client only changes the path):

  POST /lfm3/tts/stream   {"text", "voice", "sample_rate"}  -> raw PCM16 mono
  POST /lfm3/tts/bench    {"text", "voice"}                 -> seconds per piece
  GET  /lfm3/tts/health

The model (about 8 GB in float32 on CPU) loads on the first request, not at
start-up, so Kokoro boots and answers exactly as before; transformers is not
imported until then. One LFM3 synthesis at a time. "speed" is ignored (LFM3 has
no speed control; X-Speed-Ignored says so). Voices: "female", "male" or a
Kokoro id such as "af_heart" (see pieces.parse_voice).

LiquidAI/LFM3-Audio-2B is under the LFM Open License v1.0: commercial use only
below US$10M annual revenue.
"""
import asyncio
import os
import threading
import time
from typing import Callable

import numpy as np
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse

from pieces import parse_voice, system_prompt

MODEL_ID = os.getenv("LFM3_MODEL_ID", "LiquidAI/LFM3-Audio-2B")
MAX_NEW_TOKENS = int(os.getenv("LFM3_MAX_NEW_TOKENS", "256"))  # audio frames per piece
SAMPLE_RATE = 24000  # LFM3 codec output


def enabled() -> bool:
    return os.getenv("LFM3_ENABLED", "").strip().lower() in ("1", "true", "yes")


class Lfm3Engine:
    """Lazy, thread-safe LFM3 text-to-speech."""

    def __init__(self, threads: int) -> None:
        self.threads = threads
        self.device = "cpu"
        self.dtype = "float32"
        self._model = None
        self._processor = None
        self._codec = None
        self._load_lock = threading.Lock()
        self._synth_lock = threading.Lock()

    @property
    def loaded(self) -> bool:
        return self._model is not None

    def _load(self) -> None:
        with self._load_lock:
            if self._model is not None:
                return
            import torch
            from transformers import AutoModelForMultimodalLM, AutoProcessor

            started = time.perf_counter()
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
            dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
            self.dtype = str(dtype).replace("torch.", "")
            model = AutoModelForMultimodalLM.from_pretrained(
                MODEL_ID, dtype=dtype, device_map=self.device, trust_remote_code=True
            ).eval()
            processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
            self._codec = processor.audio_tokenizer.to(device=self.device).eval()
            self._processor = processor
            self._model = model  # set last: `loaded` means everything is ready
            print(f"[LFM3] loaded {MODEL_ID} on {self.device} ({self.dtype}) in {time.perf_counter() - started:.1f}s")

    def synth(self, text: str, language: str, gender: str) -> np.ndarray:
        """Float32 mono audio at 24 kHz for one piece of text."""
        import torch

        self._load()
        conversation = [
            {"role": "system", "content": system_prompt(language, gender)},
            {"role": "user", "content": text},
            {"role": "assistant", "content": "<|audio_start|>"},
        ]
        with self._synth_lock, torch.no_grad():
            inputs = self._processor.apply_chat_template(
                conversation,
                continue_final_message=True,
                tokenize=True,
                return_dict=True,
                return_tensors="pt",
            ).to(self.device)
            outputs = self._model.generate(
                **inputs,
                max_new_tokens=MAX_NEW_TOKENS,
                audio_temperature=1.0,
                audio_top_k=128,
                audio_repetition_penalty=1.1,
                audio_repetition_window=4096,
                return_dict_in_generate=True,
            )
            waveform = self._codec.decode(
                outputs.audio_codes[:1, :, : outputs.audio_lengths[0]]
            ).audio_values[0]
        return waveform.float().cpu().numpy().reshape(-1)


def build_router(
    *,
    engine: Lfm3Engine,
    check_key: Callable[[Request], None],
    split: Callable[[str], list],
    to_pcm16: Callable[[np.ndarray, int], bytes],
    max_chars: int,
) -> APIRouter:
    """Routes that share the Kokoro Space's API-key check, text splitting and
    PCM conversion."""
    router = APIRouter(prefix="/lfm3/tts")

    async def read_body(request: Request) -> dict:
        try:
            body = await request.json()
        except Exception:
            raise HTTPException(status_code=400, detail="json body required")
        if not isinstance(body, dict):
            raise HTTPException(status_code=400, detail="json object required")
        return body

    def voice_or_400(body: dict):
        parsed = parse_voice(str(body.get("voice") or "female"))
        if parsed is None:
            raise HTTPException(
                status_code=400, detail="unknown voice: use female, male or an a?_/b?_ Kokoro id"
            )
        return parsed

    @router.get("/health")
    def health():
        return {"ok": True, "engine": "lfm3", "model": MODEL_ID, "loaded": engine.loaded,
                "device": engine.device, "dtype": engine.dtype, "threads": engine.threads}

    @router.post("/stream")
    async def stream(request: Request):
        check_key(request)
        body = await read_body(request)
        text = str(body.get("text") or "").strip()[:max_chars]
        language, gender = voice_or_400(body)
        try:
            rate = int(body.get("sample_rate") or 16000)
        except (TypeError, ValueError):
            raise HTTPException(status_code=400, detail="invalid sample_rate")
        if rate not in (16000, 24000):
            raise HTTPException(status_code=400, detail="sample_rate must be 16000 or 24000")
        if not text:
            return JSONResponse({"detail": "empty text"}, status_code=400)

        pieces = split(text)

        def generate():
            started = time.perf_counter()
            for i, piece in enumerate(pieces):
                pcm = to_pcm16(engine.synth(piece, language, gender), rate)
                if i == 0:
                    print(f"[LFM3 stream] first piece in {time.perf_counter() - started:.2f}s ({len(piece)} chars)")
                yield pcm

        return StreamingResponse(
            generate(),  # sync generator: Starlette runs it in a worker thread
            media_type="application/octet-stream",
            headers={"X-Sample-Rate": str(rate), "X-Pieces": str(len(pieces)),
                     "X-Speed-Ignored": "1", "Cache-Control": "no-store"},
        )

    @router.post("/bench")
    async def bench(request: Request):
        """Seconds spent synthesizing against seconds of audio, per piece
        (rtf below 1 is faster than real time). The first call also loads the
        model, so run it twice and read the second."""
        check_key(request)
        body = await read_body(request)
        text = str(body.get("text") or "").strip()[:max_chars]
        language, gender = voice_or_400(body)
        if not text:
            return JSONResponse({"detail": "empty text"}, status_code=400)

        def run():
            rows = []
            for piece in split(text):
                started = time.perf_counter()
                audio = engine.synth(piece, language, gender)
                took = time.perf_counter() - started
                seconds = audio.size / SAMPLE_RATE
                rows.append({"chars": len(piece), "synth_seconds": round(took, 3),
                             "audio_seconds": round(seconds, 3),
                             "rtf": round(took / seconds, 3) if seconds else None})
            return rows

        rows = await asyncio.to_thread(run)
        return {"device": engine.device, "threads": engine.threads, "pieces": rows,
                "first_piece_seconds": rows[0]["synth_seconds"] if rows else None}

    return router