from fastapi import FastAPI, UploadFile, Query
from fastapi.responses import Response
import io
import numpy as np
import soundfile as sf
from contextlib import asynccontextmanager
import torch
import tempfile
import os

# 直接指定本地模型文件夹（绝对路径）
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ASR_MODEL_PATH = os.path.join(BASE_DIR, "modelscope", "hub", "models", "Qwen", "Qwen3-ASR-1___7B")
TTS_MODEL_PATH = os.path.join(BASE_DIR, "modelscope", "hub", "models", "OpenBMB", "VoxCPM2")

models = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    from transformers import BitsAndBytesConfig
    quantization_asr_config = BitsAndBytesConfig(load_in_8bit=True)
    
    print(f"Loading ASR model from {ASR_MODEL_PATH} ...")
    from qwen_asr import Qwen3ASRModel
    models["asr"] = Qwen3ASRModel.from_pretrained(
        ASR_MODEL_PATH,
        dtype=torch.bfloat16,
        device_map="cuda:0",
        max_inference_batch_size=4,
        attn_implementation="flash_attention_2",
        max_new_tokens=512,
        quantization_config=quantization_asr_config
    )
 
    from voxcpm import VoxCPM
    models["tts"] = VoxCPM.from_pretrained(
        TTS_MODEL_PATH,
        load_denoiser=False,
    )
    models["tts_sample_rate"] = models["tts"].tts_model.sample_rate

    print("ASR & TTS models loaded.")
    yield
    del models["asr"], models["tts"]

app = FastAPI(lifespan=lifespan)

@app.post("/asr")
async def speech_to_text(file: UploadFile):
    # 保存上传文件为临时 WAV
    suffix = os.path.splitext(file.filename)[-1] or ".wav"
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(await file.read())
        tmp_path = tmp.name

    try:
        # 调用 transcribe，这里直接使用自动语言检测
        results = models["asr"].transcribe(
            audio=[tmp_path],
            language=None,           # 自动检测，或指定 ["Chinese", "English"]
            return_time_stamps=False, # 如需时间戳可改为 True
        )
        # 取第一个结果（仅一个音频）
        if results:
            text = results[0].text
        else:
            text = ""
    finally:
        os.unlink(tmp_path)         # 清理临时文件

    return {"text": text}

@app.get("/tts")
async def text_to_speech(text: str = Query(...)):
    with torch.no_grad():
        wav = models["tts"].generate(
            text=text,
            cfg_value=2.0,
            inference_timesteps=10,
       )
    buffer = io.BytesIO()
    sf.write(buffer, wav, models["tts_sample_rate"], format="WAV")
    return Response(buffer.getvalue(), media_type="audio/wav")

@app.get("/health")
async def health():
    return {"status": "ok", "models": list(models.keys())}