"""FunASR Server — unified vLLM-based inference service.

Provides OpenAI-compatible API (/v1/audio/transcriptions) and REST API (/asr).
Uses vLLM for Fun-ASR-Nano (GPU) or falls back to AutoModel for non-LLM models (SenseVoice/Paraformer).
"""

import io
import os
import re
import time
import logging
import tempfile
from typing import Optional

import numpy as np
import soundfile as sf

try:
    from fastapi import FastAPI, UploadFile, File, Form, HTTPException
    from fastapi.responses import JSONResponse
except ImportError:
    raise ImportError(
        "funasr-server requires additional packages. Install with: pip install vllm fastapi uvicorn python-multipart"
    )

logger = logging.getLogger("funasr.server")


def prepare_audio_for_inference(audio_data, sr, target_sr=16000):
    """Return mono float32 audio at target_sr for ASR inference."""
    audio_data = np.asarray(audio_data)
    if audio_data.ndim > 1:
        channel_axis = -1 if audio_data.shape[-1] <= audio_data.shape[0] else 0
        audio_data = audio_data.mean(axis=channel_axis)

    if sr != target_sr:
        import librosa
        audio_data = librosa.resample(audio_data, orig_sr=sr, target_sr=target_sr)
        sr = target_sr

    return audio_data.astype(np.float32), sr

def create_app(device: str = "cuda", preload_model: str = "auto") -> FastAPI:
    if preload_model == "auto":
        preload_model = "fun-asr-nano" if device.startswith("cuda") else "sensevoice"

    app = FastAPI(title="FunASR Server", version="1.3.6")
    app.state.device = device
    app.state.engine = None
    app.state.vad_model = None
    app.state.fallback_models = {}

    # Non-LLM model configs (use AutoModel, no vLLM)
    FALLBACK_CONFIGS = {
        "sensevoice": {
            "model": "iic/SenseVoiceSmall",
            "vad_model": "fsmn-vad",
            "vad_kwargs": {"max_single_segment_time": 30000},
        },
        "paraformer": {
            "model": "paraformer-zh",
            "vad_model": "fsmn-vad",
            "punc_model": "ct-punc",
        },
    }

    def _load_vllm_engine():
        """Load Fun-ASR-Nano vLLM engine. Falls back to AutoModel if vLLM unavailable."""
        if app.state.engine is not None:
            return
        try:
            from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
            from funasr import AutoModel as _AutoModel

            logger.info("Loading Fun-ASR-Nano vLLM engine...")
            t0 = time.time()
            app.state.engine = FunASRNanoVLLM.from_pretrained(
                model="FunAudioLLM/Fun-ASR-Nano-2512",
                hub="hf",
                device=device,
                dtype="bf16",
                max_model_len=4096,
                gpu_memory_utilization=0.5,
            )
            logger.info(f"vLLM engine ready in {time.time()-t0:.1f}s")
            app.state.use_vllm = True

            logger.info("Loading VAD model...")
            app.state.vad_model = _AutoModel(model="fsmn-vad", device=device, disable_update=True)
            logger.info("VAD ready.")
        except Exception as e:
            logger.warning(f"vLLM failed ({e}), falling back to AutoModel for fun-asr-nano")
            app.state.use_vllm = False
            from funasr import AutoModel
            cfg = {
                "model": "FunAudioLLM/Fun-ASR-Nano-2512",
                "hub": "hf",
                "trust_remote_code": True,
                "vad_model": "fsmn-vad",
                "vad_kwargs": {"max_single_segment_time": 30000},
                "device": device,
                "disable_update": True,
            }
            app.state.fallback_models["fun-asr-nano"] = AutoModel(**cfg)
            logger.info("Fallback AutoModel loaded for fun-asr-nano.")

    def _load_fallback(name: str):
        """Load non-LLM model via AutoModel."""
        if name in app.state.fallback_models:
            return app.state.fallback_models[name]
        if name not in FALLBACK_CONFIGS:
            return None
        from funasr import AutoModel
        cfg = FALLBACK_CONFIGS[name].copy()
        cfg["device"] = device
        cfg["disable_update"] = True
        logger.info(f"Loading fallback model '{name}'...")
        model = AutoModel(**cfg)
        app.state.fallback_models[name] = model
        return model

    def _process_vllm(audio_data, sr, language=None, hotwords=None, use_spk=False):
        """Process audio with vLLM engine (Fun-ASR-Nano)."""
        audio_data, sr = prepare_audio_for_inference(audio_data, sr)

        # VAD
        vad_res = app.state.vad_model.generate(input=audio_data, fs=sr)
        segments = vad_res[0]["value"] if vad_res and vad_res[0].get("value") else [[0, int(len(audio_data)*1000/sr)]]

        seg_audios = []
        seg_times = []
        for seg in segments:
            s0 = int(seg[0] * sr / 1000)
            s1 = int(seg[1] * sr / 1000)
            seg_audio = audio_data[s0:s1]
            if len(seg_audio) > sr * 0.3:
                seg_audios.append(seg_audio)
                seg_times.append((seg[0], seg[1]))

        if not seg_audios:
            return {"text": "", "segments": [], "duration": len(audio_data)/sr}

        # repetition_penalty is left at the neutral 1.0: the Fun-ASR-Nano vLLM
        # engine runs in prompt-embeds mode, where any other value crashes the
        # CUDA kernel (see issue #2948 and fun_asr_nano.vllm_utils).
        gen_kwargs = {"max_new_tokens": 500, "repetition_penalty": 1.0}
        if language:
            gen_kwargs["language"] = language
        if hotwords:
            gen_kwargs["hotwords"] = hotwords

        results = app.state.engine.generate(inputs=seg_audios, **gen_kwargs)

        output_segments = []
        full_text_parts = []
        for r, (start_ms, end_ms) in zip(results, seg_times):
            text = r["text"]
            seg_info = {"text": text, "start": start_ms/1000, "end": end_ms/1000}
            if "timestamps" in r:
                offset = start_ms / 1000
                seg_info["words"] = [
                    {"word": ts["token"], "start": ts["start_time"]+offset, "end": ts["end_time"]+offset}
                    for ts in r["timestamps"]
                ]
            output_segments.append(seg_info)
            full_text_parts.append(text)

        return {
            "text": "".join(full_text_parts),
            "segments": output_segments,
            "duration": len(audio_data) / sr,
        }

    def _process_fallback(model_name, audio_path, language=None):
        """Process with non-LLM model (SenseVoice/Paraformer)."""
        model = _load_fallback(model_name)
        kwargs = {"input": audio_path, "batch_size": 1}
        if language:
            kwargs["language"] = language
        result = model.generate(**kwargs)
        text = re.sub(r'<\|[^|]*\|>', '', result[0]["text"]).strip()
        segments = []
        if "sentence_info" in result[0]:
            for s in result[0]["sentence_info"]:
                segments.append({
                    "start": s.get("start", 0)/1000,
                    "end": s.get("end", 0)/1000,
                    "text": re.sub(r'<\|[^|]*\|>', '', s.get("text", "")).strip(),
                    "speaker": s.get("spk"),
                })
        return {"text": text, "segments": segments}

    # Pre-load
    if preload_model == "fun-asr-nano":
        _load_vllm_engine()
    else:
        _load_fallback(preload_model)

    @app.post("/v1/audio/transcriptions")
    async def transcribe(
        file: UploadFile = File(...),
        model: str = Form(default="fun-asr-nano"),
        language: Optional[str] = Form(default=None),
        response_format: Optional[str] = Form(default="json"),
        spk: bool = Form(default=False),
    ):
        content = await file.read()
        t0 = time.perf_counter()

        if model == "fun-asr-nano":
            _load_vllm_engine()
            if app.state.use_vllm:
                audio_data, sr = sf.read(io.BytesIO(content))
                result = _process_vllm(audio_data, sr, language=language, use_spk=spk)
            else:
                suffix = os.path.splitext(file.filename)[1] if file.filename else ".wav"
                with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
                    tmp.write(content)
                    tmp_path = tmp.name
                try:
                    result = _process_fallback("fun-asr-nano", tmp_path, language=language)
                finally:
                    os.unlink(tmp_path)
        elif model in FALLBACK_CONFIGS:
            suffix = os.path.splitext(file.filename)[1] if file.filename else ".wav"
            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
                tmp.write(content)
                tmp_path = tmp.name
            try:
                result = _process_fallback(model, tmp_path, language=language)
            finally:
                os.unlink(tmp_path)
        else:
            raise HTTPException(400, f"Unknown model '{model}'. Available: fun-asr-nano, {', '.join(FALLBACK_CONFIGS.keys())}")

        t1 = time.perf_counter()

        if response_format == "verbose_json":
            return JSONResponse({
                "task": "transcribe",
                "language": language or "zh",
                "duration": result.get("duration", 0),
                "text": result["text"],
                "segments": [
                    {"id": i, "start": s["start"], "end": s["end"], "text": s["text"], "words": s.get("words", [])}
                    for i, s in enumerate(result["segments"])
                ],
            })
        elif response_format == "text":
            return JSONResponse(result["text"])
        else:
            return JSONResponse({"text": result["text"]})

    @app.post("/asr")
    async def asr_endpoint(
        file: UploadFile = File(...),
        language: Optional[str] = Form(default=None),
        hotwords: str = Form(default=""),
        spk: bool = Form(default=False),
    ):
        """Full-featured ASR endpoint with timestamps and speaker diarization."""
        content = await file.read()
        _load_vllm_engine()
        hw_list = [w.strip() for w in hotwords.split(",") if w.strip()] if hotwords else None

        t0 = time.perf_counter()
        if app.state.use_vllm:
            audio_data, sr = sf.read(io.BytesIO(content))
            result = _process_vllm(audio_data, sr, language=language, hotwords=hw_list, use_spk=spk)
        else:
            suffix = ".wav"
            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
                tmp.write(content)
                tmp_path = tmp.name
            try:
                result = _process_fallback("fun-asr-nano", tmp_path, language=language)
            finally:
                os.unlink(tmp_path)
        t1 = time.perf_counter()

        result["processing_time"] = round(t1 - t0, 3)
        result["rtf"] = round((t1 - t0) / result["duration"], 4) if result.get("duration", 0) > 0 else 0
        return JSONResponse(result)

    @app.get("/v1/models")
    async def list_models():
        all_models = ["fun-asr-nano"] + list(FALLBACK_CONFIGS.keys())
        return JSONResponse({"object": "list", "data": [{"id": n, "object": "model"} for n in all_models]})

    @app.get("/health")
    async def health():
        loaded = []
        if app.state.engine is not None:
            loaded.append("fun-asr-nano (vLLM)")
        loaded.extend(app.state.fallback_models.keys())
        return {"status": "ok", "device": device, "models_loaded": loaded}

    return app