plugins/audio/api.py

import asyncio
import base64
import tempfile
from pathlib import Path

from fastapi import APIRouter, HTTPException, UploadFile, Query
from pydantic import BaseModel

from .asr_drivers import get_asr_registry, AsrDriver, AsrResult
from app.plugins.run_log import get_run_log_service, RunStatus

router = APIRouter()
plugin = None


def set_plugin_instance(plugin_instance):
    global plugin
    plugin = plugin_instance


class TranscribeBase64Request(BaseModel):
    audio_base64: str
    suffix: str = ".wav"
    language: str | None = None
    provider: str | None = None


@router.get("/status")
async def get_status():
    if plugin is None:
        return {
            "name": "audio",
            "enabled": False,
            "message": "插件未加载",
        }
    return plugin.get_status()


@router.get("/providers")
async def list_providers():
    """列出所有 ASR provider 及其可用性"""
    registry = get_asr_registry()
    availability = registry.probe_all()
    return {
        "providers": [
            {
                "name": name,
                "status": status.value,
            }
            for name, status in availability.items()
        ],
        "default": registry.get_default_provider(),
    }


def _transcribe_and_finish(
    run_id: str,
    driver: AsrDriver,
    audio_path: str,
    language: str | None,
    provider_name: str,
    availability_value: str,
) -> None:
    """同步执行转录并写入事件/完成 run，最后清理临时文件。

    供 wait 模式（to_thread 同步等待）和非 wait 模式（后台任务）共用，
    保证事件写入逻辑只有一份。
    """
    run_service = get_run_log_service()
    try:
        # 写入模型探测事件
        run_service.add_event(
            run_id=run_id,
            stage="model_probe",
            message=f"模型探测: {driver.model_id}",
            detail=f"provider={provider_name}, status={availability_value}",
        )

        # 写入转录开始事件
        run_service.add_event(
            run_id=run_id,
            stage="transcribe",
            message="开始转录",
        )

        # 执行转录（同步，可能耗时较长）
        result: AsrResult = driver.transcribe(audio_path, language)

        if result.success:
            # 写入转录完成事件
            run_service.add_event(
                run_id=run_id,
                stage="transcribe",
                message=f"转录完成: {len(result.text)} 字符",
                detail=f"timings={result.timings}",
            )
            run_service.finish_run(
                run_id=run_id,
                status=RunStatus.SUCCEEDED,
                result={
                    "text": result.text,
                    "segments": result.segments,
                    "provider": result.provider,
                    "model": result.model,
                    "timings": result.timings,
                },
            )
        else:
            # 写入转录失败事件
            run_service.add_event(
                run_id=run_id,
                stage="transcribe",
                message=f"转录失败: {result.error}",
                level="error",
            )
            run_service.finish_run(
                run_id=run_id,
                status=RunStatus.FAILED,
                error=result.error,
            )
    except Exception as e:
        # 写入异常事件
        run_service.add_event(
            run_id=run_id,
            stage="transcribe",
            message=f"转录异常: {str(e)}",
            level="error",
        )
        run_service.finish_run(
            run_id=run_id,
            status=RunStatus.FAILED,
            error=str(e),
        )
    finally:
        # 清理临时文件
        Path(audio_path).unlink(missing_ok=True)


@router.post("/upload")
async def transcribe_upload(
    file: UploadFile,
    language: str | None = None,
    provider: str | None = None,
    wait: bool = Query(False, description="是否同步等待转录完成"),
):
    """上传音频文件并转录

    默认异步返回 run_id，前端轮询 run 事件获取实时进度；
    设置 wait=true 同步等待结果。
    """
    if plugin is None or not plugin.enabled:
        raise HTTPException(status_code=400, detail="插件未启用")

    # 获取 ASR registry
    registry = get_asr_registry()
    provider_name = provider or registry.get_default_provider()
    driver = registry.get_driver(provider_name)
    if driver is None:
        raise HTTPException(status_code=400, detail=f"未知的 ASR provider: {provider_name}")

    # 检查 provider 可用性（同步快速失败）
    availability = driver.check_availability()
    if availability.value != "available":
        raise HTTPException(
            status_code=400,
            detail=f"ASR provider {provider_name} 不可用: {availability.value}",
        )

    # 创建 run
    run_service = get_run_log_service()
    run = run_service.create_run("audio")

    # 写入上传事件
    run_service.add_event(
        run_id=run.run_id,
        stage="upload",
        message=f"音频文件已上传: {file.filename}",
        detail=f"provider={provider_name}, language={language}",
    )

    # 保存上传文件到临时目录
    content = await file.read()
    suffix = "." + file.filename.rsplit(".", 1)[-1] if "." in file.filename else ".wav"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
        tmp_file.write(content)
        tmp_path = tmp_file.name

    if wait:
        # 同步等待：在线程池执行转录，不阻塞事件循环
        await asyncio.to_thread(
            _transcribe_and_finish,
            run.run_id, driver, tmp_path, language, provider_name, availability.value,
        )
        # 读取最终结果返回
        finished = run_service.get_run(run.run_id)
        return {
            "run_id": run.run_id,
            "status": finished.status.value if finished else "failed",
            "provider": provider_name,
            "model": driver.model_id,
            "text": (finished.result or {}).get("text", "") if finished else "",
            "segments": (finished.result or {}).get("segments", []) if finished else [],
            "timings": (finished.result or {}).get("timings") if finished else None,
            "error": finished.error if finished else "运行记录缺失",
        }

    # 异步模式：启动后台任务，立即返回 run_id
    asyncio.create_task(
        asyncio.to_thread(
            _transcribe_and_finish,
            run.run_id, driver, tmp_path, language, provider_name, availability.value,
        )
    )
    return {
        "run_id": run.run_id,
        "status": "running",
        "provider": provider_name,
        "model": driver.model_id,
    }


@router.post("/audio-base64")
async def transcribe_base64(request: TranscribeBase64Request):
    """Base64 音频转录（兼容旧接口，同步返回结果）"""
    if plugin is None or not plugin.enabled:
        raise HTTPException(status_code=400, detail="插件未启用")

    # 获取 ASR registry
    registry = get_asr_registry()
    provider_name = request.provider or registry.get_default_provider()
    driver = registry.get_driver(provider_name)
    if driver is None:
        raise HTTPException(status_code=400, detail=f"未知的 ASR provider: {provider_name}")

    # 检查 provider 可用性
    availability = driver.check_availability()
    if availability.value != "available":
        raise HTTPException(
            status_code=400,
            detail=f"ASR provider {provider_name} 不可用: {availability.value}",
        )

    # 创建 run
    run_service = get_run_log_service()
    run = run_service.create_run("audio")

    # 写入上传事件
    run_service.add_event(
        run_id=run.run_id,
        stage="upload",
        message="Base64 音频已接收",
        detail=f"provider={provider_name}, language={request.language}",
    )

    # 解码 Base64 并保存到临时文件
    try:
        audio_bytes = base64.b64decode(request.audio_base64)
    except Exception as e:
        run_service.finish_run(run_id=run.run_id, status=RunStatus.FAILED, error=f"Base64 解码失败: {e}")
        raise HTTPException(status_code=400, detail=f"Base64 解码失败: {e}")

    suffix = request.suffix if request.suffix.startswith(".") else f".{request.suffix}"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
        tmp_file.write(audio_bytes)
        tmp_path = tmp_file.name

    # 同步等待：在线程池执行转录，不阻塞事件循环
    await asyncio.to_thread(
        _transcribe_and_finish,
        run.run_id, driver, tmp_path, request.language, provider_name, availability.value,
    )
    finished = run_service.get_run(run.run_id)
    return {
        "run_id": run.run_id,
        "status": finished.status.value if finished else "failed",
        "provider": provider_name,
        "model": driver.model_id,
        "text": (finished.result or {}).get("text", "") if finished else "",
        "segments": (finished.result or {}).get("segments", []) if finished else [],
        "timings": (finished.result or {}).get("timings") if finished else None,
        "error": finished.error if finished else None,
    }