#!/usr/bin/env python3
"""把本地 ASR 包成一个 OpenAI 兼容的 /v1/audio/transcriptions 端点。

存在的意义：后端 internal/skills/packages/audio_transcribe/transcribe.go 是按
OpenAI 兼容协议写的（multipart 发 model/language/file，读回
{duration,text,segments:[{speaker,start,end,text}]}）。只要本地服务说同一套协议，
**Go 侧一行都不用改**，改的是 config/ai_config.json 里那条 audio_routes 的 base_url。

用法：
    ./venv/bin/python serve.py --port 8090
然后 config/ai_config.json 加：
    "audio_route_local_whisperx": {
      "provider": "local_asr", "base_url": "http://127.0.0.1:8090/v1",
      "endpoint": "/audio/transcriptions", "model": "large-v3",
      "supports_speakers": true, "timeout_seconds": 3600
    }
⚠️ 光改 base_url 还不够：transcribe.go:85 在 APIKey 为空时直接报错，而 key 只从
ProviderSecretKey[provider] 映射来。所以还得往 internal/config/json_loader.go 的
ProviderSecretKey 加一行 "local_asr": "LOCAL_ASR_API_KEY"，并在 ai_secrets.json
里塞个占位值（本地服务不校验它，但 Go 侧要求非空）。

    curl -s -X POST http://127.0.0.1:8090/v1/audio/transcriptions \
      -H "Authorization: Bearer local" -F file=@x.mp3 -F model=large-v3 -F language=zh
"""
import argparse
import hashlib
import json
import shutil
import sys
import tempfile
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import asr_core  # noqa: E402

from fastapi import FastAPI, File, Form, Header, HTTPException, UploadFile  # noqa: E402
from fastapi.responses import JSONResponse  # noqa: E402
from starlette.concurrency import run_in_threadpool  # noqa: E402
import uvicorn  # noqa: E402

CACHE = asr_core.ROOT / "cache"
CACHE.mkdir(exist_ok=True)

app = FastAPI(title="本地 ASR（OpenAI 兼容）", version="0.1.0")
STARTED_AT = time.time()


@app.get("/health")
def health():
    return {
        "status": "ok",
        "uptime_seconds": round(time.time() - STARTED_AT, 1),
        "model": "faster-whisper-large-v3 + pyannote/speaker-diarization-3.1",
        "models_dir": str(asr_core.MODELS),
    }


@app.get("/v1/models")
def models():
    # 后端不查这个端点，但 OpenAI 兼容服务一般都有，留着方便人工确认。
    return {"object": "list",
            "data": [{"id": "large-v3", "object": "model", "owned_by": "local"}]}


@app.post("/v1/audio/transcriptions")
async def transcriptions(
    file: UploadFile = File(...),
    model: str = Form("large-v3"),
    language: str = Form(""),
    response_format: str = Form("json"),
    num_speakers: str = Form(""),
    authorization: str = Header(default=""),
):
    # 后端会带 Authorization: Bearer <key>。这里只要求「有」，不校验具体值 ——
    # 本地服务绑 127.0.0.1，本来就不对外。
    if not authorization.strip():
        raise HTTPException(401, "缺少 Authorization 头")

    raw = await file.read()
    if not raw:
        raise HTTPException(400, "上传的音频是空的")

    digest = hashlib.sha256(raw).hexdigest()
    work = CACHE / digest
    cached = work / "asr_result.json"

    # 按内容哈希缓存：同一份音频重发（比如后端重试、或我复跑）不该再烧一次 GPU。
    if cached.exists() and cached.stat().st_size > 0:
        asr_core.log(f"命中缓存 {digest[:12]}，直接返回")
        return JSONResponse(_to_openai(json.loads(cached.read_text()), model))

    work.mkdir(parents=True, exist_ok=True)
    suffix = Path(file.filename or "audio.mp3").suffix or ".mp3"
    src = work / f"source{suffix}"
    if not src.exists():
        src.write_bytes(raw)

    n = int(num_speakers) if num_speakers.strip().isdigit() else None
    asr_core.log(f"新请求 {digest[:12]}：{file.filename} "
                 f"({len(raw)/1048576:.1f} MB, language={language or 'auto'})")

    try:
        # 丢到工作线程去跑，**不能在事件循环里直接调**。
        #
        # asr_core.run 是分钟级的阻塞调用（26 分钟音频实测 3 分钟），
        # 而 uvicorn 只有一个事件循环：直接调的话，这段时间里 /v1/models
        # 一个字都回不了。平台的健康探测正是打 /v1/models 的，于是
        # 「正在转写」被读成「本地服务挂了」→ audio_route_auto 解到云端路由 →
        # 下一个任务的音频就出公网了。整个过程没有任何报错。
        #
        # 转写本身仍然是串行的（asr_core 里的 _LOCK），这里只是把「等 GPU」
        # 从事件循环里挪出去，让探活和排队中的请求还能被受理。
        result = await run_in_threadpool(asr_core.run, src, work, language or "zh", n)
    except Exception as e:
        # 不吞异常：ASR 失败必须让调用方看见 5xx，不能返回一个空稿当成功。
        # 后端那边对「空文本」也是按失败处理的（transcribe.go:203）。
        asr_core.log(f"失败：{type(e).__name__}: {e}")
        raise HTTPException(500, f"本地 ASR 失败：{type(e).__name__}: {e}")

    return JSONResponse(_to_openai(result, model))


def _to_openai(result: dict, model: str) -> dict:
    """转成 OpenAI 兼容响应。

    `usage` 里带 duration 是照着 SiliconFlow 的 diarize 模型来的 —— 后端
    ai/credits.go 按这个记用量。本地跑不要钱，但字段留着，免得下游解析时
    因为缺字段而报错。
    """
    return {
        "text": result["text"],
        "duration": result["duration"],
        "language": result.get("language", ""),
        "model": model,
        "segments": result["segments"],
        "usage": {"type": "duration", "seconds": result["duration"]},
        # 非标准字段，方便人工核对；后端只读上面那几个，多给的会被忽略。
        "_local": {"speakers": result.get("speakers", []),
                   "elapsed_seconds": result.get("elapsed_seconds")},
    }


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8090)
    args = ap.parse_args()
    asr_core.log(f"本地 ASR 服务启动于 http://{args.host}:{args.port}")
    asr_core.log(f"模型目录 {asr_core.MODELS}")
    uvicorn.run(app, host=args.host, port=args.port, log_level="warning")