"""
FunASR Paraformer 语音识别 FastAPI 服务

ASR 模型: speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch
VAD 模型: speech_fsmn_vad_zh-cn-16k-common-pytorch
PUNC 模型: punc_ct-transformer_zh-cn-common-vocab272727-pytorch
端口: 59805
"""

import os
import re
import json
import asyncio
import tempfile
import logging
import subprocess
import wave
from typing import Optional

import requests
import uvicorn
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse, PlainTextResponse
from funasr import AutoModel


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


_CN_NUM = {
    "零": 0,
    "一": 1,
    "幺": 1,
    "二": 2,
    "两": 2,
    "三": 3,
    "四": 4,
    "五": 5,
    "六": 6,
    "七": 7,
    "八": 8,
    "九": 9,
}

_CN_UNIT = {
    "十": 10,
    "百": 100,
    "千": 1000,
    "万": 10000,
    "亿": 100000000,
}

_CN_PATTERN = re.compile(r"[零一幺二两三四五六七八九十百千万亿]+")


def _cn2num(s: str) -> int:
    total, section, num = 0, 0, 0

    for ch in s:
        if ch in _CN_NUM:
            num = _CN_NUM[ch]
        elif ch in _CN_UNIT:
            unit = _CN_UNIT[ch]
            if unit >= 10000:
                section = (section + (num or 1)) * unit
                total += section
                section = 0
            else:
                section += (num or 1) * unit
            num = 0

    return total + section + num


def cn_to_arabic(text: str) -> str:
    return _CN_PATTERN.sub(lambda m: str(_cn2num(m.group())), text)


def read_txt_file(file_path: str) -> str:
    try:
        with open(file_path, "r", encoding="utf-8") as file:
            return file.read()
    except FileNotFoundError:
        logger.warning(f"找不到热词文件: {file_path}")
        return ""
    except Exception as e:
        logger.warning(f"读取热词文件失败: {e}")
        return ""


app = FastAPI(
    title="Paraformer ASR API",
    description="基于 FunASR Paraformer 的中文语音识别服务",
    version="1.0.0",
)


model = None
VOICE_DETECTED_URL = os.environ.get("VOICE_DETECTED_URL", "http://127.0.0.1:18000")
VOICEPRINT_STORE = os.environ.get("VOICEPRINT_STORE", "/tmp/asr_voiceprints.json")
VOICEPRINT_MATCH_THRESHOLD = float(os.environ.get("VOICEPRINT_MATCH_THRESHOLD", "0.45"))
REALTIME_PARTIAL_INTERVAL_SECONDS = float(os.environ.get("REALTIME_PARTIAL_INTERVAL_SECONDS", "2.0"))
REALTIME_PARTIAL_MIN_SECONDS = float(os.environ.get("REALTIME_PARTIAL_MIN_SECONDS", "1.5"))
REALTIME_PARTIAL_WINDOW_SECONDS = float(os.environ.get("REALTIME_PARTIAL_WINDOW_SECONDS", "6.0"))


def _voiceprint_key(database_name: str, collection_name: str) -> str:
    return f"{database_name or 'NB'}::{collection_name or 'voice'}"


def _load_voiceprints() -> dict:
    if not os.path.exists(VOICEPRINT_STORE):
        return {}
    try:
        with open(VOICEPRINT_STORE, "r", encoding="utf-8") as f:
            data = json.load(f)
        return data if isinstance(data, dict) else {}
    except Exception as e:
        logger.warning(f"读取声纹库失败，将使用空库: {e}")
        return {}


def _save_voiceprints(data: dict):
    os.makedirs(os.path.dirname(VOICEPRINT_STORE) or ".", exist_ok=True)
    tmp_path = f"{VOICEPRINT_STORE}.tmp"
    with open(tmp_path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False)
    os.replace(tmp_path, VOICEPRINT_STORE)


def _extract_embedding(audio_path: str, filename: str) -> list:
    with open(audio_path, "rb") as f:
        response = requests.post(
            f"{VOICE_DETECTED_URL.rstrip('/')}/extract_embedding",
            files={"wav_file": (filename or "voice.wav", f, "audio/wav")},
            timeout=120,
        )
    response.raise_for_status()
    payload = response.json()
    embedding = payload.get("embedding")
    if not isinstance(embedding, list):
        raise ValueError("voice-dected 返回中没有有效 embedding")
    return embedding


def _convert_to_embedding_wav(audio_path: str) -> str:
    output = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    output_path = output.name
    output.close()

    try:
        subprocess.run(
            [
                "ffmpeg",
                "-y",
                "-i",
                audio_path,
                "-acodec",
                "pcm_s16le",
                "-ar",
                "16000",
                "-ac",
                "1",
                output_path,
            ],
            check=True,
            capture_output=True,
            text=True,
            timeout=60,
        )
        return output_path
    except Exception:
        if os.path.exists(output_path):
            os.unlink(output_path)
        raise


def _cosine_similarity(a: list, b: list) -> float:
    if not a or not b or len(a) != len(b):
        return -1.0
    dot = sum(float(x) * float(y) for x, y in zip(a, b))
    norm_a = sum(float(x) * float(x) for x in a) ** 0.5
    norm_b = sum(float(y) * float(y) for y in b) ** 0.5
    if norm_a == 0 or norm_b == 0:
        return -1.0
    return dot / (norm_a * norm_b)


def _match_voiceprint(audio_path: str, filename: str, database_name: str = "NB", collection_name: str = "voice"):
    store = _load_voiceprints()
    records = store.get(_voiceprint_key(database_name, collection_name), {})
    if not records:
        return None, None

    embedding_path = _convert_to_embedding_wav(audio_path)
    try:
        embedding = _extract_embedding(embedding_path, "match.wav")
    finally:
        if os.path.exists(embedding_path):
            os.unlink(embedding_path)

    best_name = None
    best_score = -1.0

    for name, record in records.items():
        score = _cosine_similarity(embedding, record.get("embedding", []))
        if score > best_score:
            best_name = name
            best_score = score

    logger.info(f"声纹匹配结果: name={best_name}, score={best_score:.4f}, threshold={VOICEPRINT_MATCH_THRESHOLD}")
    if best_name and best_score >= VOICEPRINT_MATCH_THRESHOLD:
        return best_name, best_score
    return None, best_score


def load_model():
    global model

    if model is None:
        logger.info("正在加载 Paraformer 模型...")
        model = AutoModel(
            model="/app/models/seacomodel",
            vad_model="/app/models/vadmodel",
            punc_model="/app/models/ctpuncmodel",
        )
        logger.info("模型加载完成")

    return model


@app.on_event("startup")
async def startup_event():
    load_model()


@app.get("/")
async def root():
    return {
        "status": "ok",
        "message": "Paraformer ASR 服务运行中",
    }


@app.get("/health")
async def health():
    return {"status": "healthy"}


def _write_pcm_wav(audio_bytes: bytes) -> str:
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
        tmp_path = tmp.name

    with wave.open(tmp_path, "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(16000)
        wav_file.writeframes(audio_bytes)

    return tmp_path


def _recent_pcm(audio_chunks: list[bytes], window_seconds: float) -> bytes:
    max_bytes = max(int(window_seconds * 32000), 32000)
    audio_bytes = b"".join(audio_chunks)
    return audio_bytes[-max_bytes:]


def _text_delta(previous: str, current: str) -> str:
    previous = previous or ""
    current = current or ""
    if not current or current == previous:
        return ""
    if current.startswith(previous):
        return current[len(previous):]

    max_overlap = min(len(previous), len(current), 80)
    for size in range(max_overlap, 0, -1):
        if previous[-size:] == current[:size]:
            return current[size:]
    return current


async def run_asr_websocket(websocket: WebSocket):
    await websocket.accept(subprotocol="binary")
    audio_chunks = []
    last_partial_at = 0.0
    last_partial_text = ""

    try:
        while True:
            message = await websocket.receive()

            if "bytes" in message and message["bytes"]:
                audio_chunks.append(message["bytes"])
                audio_seconds = sum(len(chunk) for chunk in audio_chunks) / 32000.0
                now = asyncio.get_running_loop().time()

                if (
                    audio_seconds >= REALTIME_PARTIAL_MIN_SECONDS
                    and now - last_partial_at >= REALTIME_PARTIAL_INTERVAL_SECONDS
                ):
                    last_partial_at = now
                    partial_path = None
                    try:
                        partial_path = _write_pcm_wav(_recent_pcm(audio_chunks, REALTIME_PARTIAL_WINDOW_SECONDS))
                        result = await run_asr(partial_path, language="zh", response_format="json")
                        partial_text = result.get("text", "") if isinstance(result, dict) else ""
                        if partial_text and partial_text != last_partial_text:
                            await websocket.send_json({
                                "mode": "2pass-online",
                                "text": partial_text,
                                "timestamp": str(int(audio_seconds * 1000)),
                                "is_full_text": True,
                                "is_final": False,
                            })
                            last_partial_text = partial_text
                    except Exception as e:
                        logger.warning(f"实时部分识别失败: {e}")
                    finally:
                        if partial_path and os.path.exists(partial_path):
                            os.unlink(partial_path)
                continue

            if "text" in message and message["text"]:
                try:
                    payload = json.loads(message["text"])
                except Exception:
                    payload = {}

                if payload.get("type") == "stop" or payload.get("is_speaking") is False:
                    break
                continue

            if message.get("type") == "websocket.disconnect":
                break
    except WebSocketDisconnect:
        pass

    text = ""
    tmp_path = None
    try:
        if audio_chunks:
            tmp_path = _write_pcm_wav(b"".join(audio_chunks))
            result = await run_asr(tmp_path, language="zh", response_format="json")
            if isinstance(result, dict):
                text = result.get("text", "")

        try:
            await websocket.send_json({
                "mode": "2pass-offline",
                "text": text,
                "is_final": True,
            })
        except Exception as e:
            logger.warning(f"Skip final transcript because websocket is closed: {e}")
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)
        try:
            await websocket.close()
        except Exception:
            pass


@app.websocket("/ws/asr")
async def websocket_asr(websocket: WebSocket):
    await run_asr_websocket(websocket)


@app.websocket("/")
async def websocket_asr_root(websocket: WebSocket):
    await run_asr_websocket(websocket)


def make_speakr_response(text: str, language: str = "zh", speaker: str = "SPEAKER_00"):
    text = (text or "").strip()
    return {
        "text": text,
        "language": language,
        "segments": [
            {
                "start": None,
                "end": None,
                "text": text,
                "speaker": speaker or "SPEAKER_00",
            }
        ] if text else [],
    }


async def run_asr(
    audio_path: str,
    language: str = "zh",
    response_format: str = "json",
):
    try:
        hotword = read_txt_file("/app/hotword.txt")
        result = model.generate(input=audio_path, hotword=hotword)

        if result and len(result) > 0:
            text = result[0].get("text", "")
        else:
            text = ""

        text = cn_to_arabic(text)
        logger.info(f"识别结果: {text}")

        if response_format == "text":
            return PlainTextResponse(content=text)

        return {
            "text": text,
            "language": language,
        }

    except Exception as e:
        logger.error(f"识别出错: {e}")
        raise HTTPException(status_code=500, detail=f"识别失败: {str(e)}")


@app.post("/api/v1/audio/transcriptions")
async def transcribe(
    request: Request,
    file: Optional[UploadFile] = File(None),
    language: Optional[str] = Form("zh"),
    response_format: Optional[str] = Form("json"),
):
    """
    语音转文字接口，兼容两种方式：

    1. multipart/form-data 上传文件：
       - file: 音频文件
       - language: 默认 zh
       - response_format: json/text

    2. application/json 传服务器本地文件路径：
       {
         "file_path": "/app/audio/test.wav",
         "language": "zh",
         "response_format": "json"
       }
    """

    tmp_path = None
    should_delete = False

    try:
        content_type = request.headers.get("content-type", "").lower()

        if "application/json" in content_type:
            body = await request.json()

            file_path = body.get("file_path")
            language = body.get("language", language)
            response_format = body.get("response_format", response_format)

            if not file_path:
                raise HTTPException(status_code=400, detail="缺少 file_path")

            if not os.path.exists(file_path):
                raise HTTPException(status_code=400, detail=f"文件不存在: {file_path}")

            if not os.path.isfile(file_path):
                raise HTTPException(status_code=400, detail=f"不是有效文件: {file_path}")

            audio_path = file_path
            logger.info(f"处理服务器本地文件: {audio_path}")

        else:
            if file is None:
                raise HTTPException(status_code=400, detail="缺少上传文件 file")

            suffix = os.path.splitext(file.filename or "")[1] or ".wav"

            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
                content = await file.read()

                if not content:
                    raise HTTPException(status_code=400, detail="上传文件为空")

                tmp.write(content)
                tmp_path = tmp.name
                should_delete = True

            audio_path = tmp_path
            logger.info(
                f"处理上传文件: {file.filename}, 大小: {len(content)} bytes, 临时路径: {audio_path}"
            )

        return await run_asr(
            audio_path=audio_path,
            language=language or "zh",
            response_format=response_format or "json",
        )

    finally:
        if should_delete and tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)


@app.post("/v1/audio/transcriptions")
async def transcribe_v1(
    request: Request,
    file: Optional[UploadFile] = File(None),
    language: Optional[str] = Form("zh"),
    response_format: Optional[str] = Form("json"),
):
    return await transcribe(
        request=request,
        file=file,
        language=language,
        response_format=response_format,
    )


@app.post("/asr")
async def transcribe_speakr_compatible(
    audio_file: UploadFile = File(...),
    language: Optional[str] = Form("zh"),
    database_name: Optional[str] = Form("NB"),
    collection_name: Optional[str] = Form("voice"),
):
    """
    Speakr 兼容接口：
    - 上传字段：audio_file
    - 返回字段：segments，其中每个 segment 含 text/speaker/start/end
    """

    tmp_path = None

    try:
        suffix = os.path.splitext(audio_file.filename or "")[1] or ".wav"

        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            content = await audio_file.read()

            if not content:
                raise HTTPException(status_code=400, detail="上传文件为空")

            tmp.write(content)
            tmp_path = tmp.name

        logger.info(
            f"处理 Speakr /asr 上传文件: {audio_file.filename}, 大小: {len(content)} bytes, 临时路径: {tmp_path}"
        )

        result = await run_asr(
            audio_path=tmp_path,
            language=language or "zh",
            response_format="json",
        )
        speaker_name = "SPEAKER_00"
        try:
            matched_name, score = _match_voiceprint(
                tmp_path,
                audio_file.filename or "audio.wav",
                database_name or "NB",
                collection_name or "voice",
            )
            if matched_name:
                speaker_name = matched_name
        except Exception as e:
            logger.warning(f"声纹匹配失败，使用默认说话人: {e}")

        return make_speakr_response(
            text=result.get("text", "") if isinstance(result, dict) else "",
            language=result.get("language", language or "zh") if isinstance(result, dict) else language or "zh",
            speaker=speaker_name,
        )

    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)


@app.get("/get_voice_name")
async def get_voice_name(
    database_name: str = "NB",
    collection_name: str = "voice",
):
    store = _load_voiceprints()
    records = store.get(_voiceprint_key(database_name, collection_name), {})
    return {
        "code": 0,
        "success": True,
        "message": "ok",
        "data": [
            {"name": name, "voice_name": name}
            for name in sorted(records.keys())
        ],
    }


@app.post("/voice_insert")
async def voice_insert(
    voice_name: str = Form(...),
    database_name: str = Form("NB"),
    collection_name: str = Form("voice"),
    voice_file: UploadFile = File(...),
):
    tmp_path = None

    try:
        suffix = os.path.splitext(voice_file.filename or "")[1] or ".wav"
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            content = await voice_file.read()
            if not content:
                raise HTTPException(status_code=400, detail="上传声纹文件为空")
            tmp.write(content)
            tmp_path = tmp.name

        embedding = _extract_embedding(tmp_path, voice_file.filename or f"{voice_name}.wav")
        store = _load_voiceprints()
        key = _voiceprint_key(database_name, collection_name)
        records = store.setdefault(key, {})
        records[voice_name] = {
            "name": voice_name,
            "database_name": database_name,
            "collection_name": collection_name,
            "embedding": embedding,
        }
        _save_voiceprints(store)

        return {
            "code": 0,
            "success": True,
            "message": "voice inserted",
            "data": {"name": voice_name, "embedding_size": len(embedding)},
        }
    except requests.RequestException as e:
        logger.error(f"调用 voice-dected 失败: {e}")
        raise HTTPException(status_code=502, detail=f"voice-dected 调用失败: {e}")
    except Exception as e:
        logger.error(f"写入声纹失败: {e}")
        raise HTTPException(status_code=500, detail=f"写入声纹失败: {e}")
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)


@app.delete("/delete_voice")
async def delete_voice(
    name: str,
    database_name: str = "NB",
    collection_name: str = "voice",
):
    store = _load_voiceprints()
    key = _voiceprint_key(database_name, collection_name)
    records = store.setdefault(key, {})
    existed = name in records
    records.pop(name, None)
    _save_voiceprints(store)

    return {
        "code": 0,
        "success": True,
        "message": "voice deleted" if existed else "voice not found",
        "data": {"name": name, "deleted": existed},
    }


if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=59805,
        reload=False,
        workers=1,
        log_level="info",
    )