import asyncio
import json
import os
import tempfile
import wave
import websockets
import aiohttp

# ========================
# 配置
# ========================
TARGET_WS_URL = os.getenv("TARGET_WS_URL", "ws://localhost:59805/ws/asr")
LOCAL_HOST = "0.0.0.0"
LOCAL_PORT = 10095

# ⚠️ 修改为你自己的 FastAPI ASR 服务地址
ASR_LOCAL_URL = os.getenv("ASR_LOCAL_URL", "http://172.18.127.124:6688/asr")

SAMPLE_RATE = 16000  # ⬅️ 关键！
SAMPLE_WIDTH = 2      # 16-bit
CHANNELS = 1

config = {
    "chunk_size": [10, 10, 10],
    "wav_name": "h5",
    "is_speaking": True,
    "wav_format": "pcm",
    "chunk_interval": 10,
    "itn": True,
    "mode": "2pass",
    "hotwords": "",
}


# ========================
# 调用本地 ASR 接口（带说话人识别）
# ========================
async def call_local_asr_with_audio(wav_path: str, session: aiohttp.ClientSession):
    """调用你自己的 /asr 接口，返回结构化结果"""
    try:
        with open(wav_path, "rb") as f:
            data = aiohttp.FormData()
            data.add_field(
                "audio_file",
                f,
                filename=os.path.basename(wav_path),
                content_type="audio/wav"
            )
            # 可选：传递参数（如 spk_diarization=True 已默认开启）
            params = {
                "spk_diarization": "true",
                "spk_threshold": "0.6",
                "output": "json"
            }
            async with session.post(ASR_LOCAL_URL, data=data, params=params, timeout=aiohttp.ClientTimeout(total=60)) as resp:
                if resp.status == 200:
                    return await resp.json()
                else:
                    text = await resp.text()
                    print(f"❌ 本地 ASR 接口错误: {resp.status} - {text}")
                    return None
    except Exception as e:
        print(f"❌ 调用本地 ASR 异常: {e}")
        return None


# ========================
# 后台处理：音频 → 本地 ASR → 推送结果
# ========================
async def process_audio_with_local_asr(audio_chunks, asr_text, websocket, session):
    if not audio_chunks:
        return

    wav_path = None
    try:
        # 1. 保存为 WAV
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
            wav_path = tmp.name
        with wave.open(wav_path, 'wb') as wf:
            wf.setnchannels(CHANNELS)
            wf.setsampwidth(SAMPLE_WIDTH)
            wf.setframerate(SAMPLE_RATE)
            wf.writeframes(b''.join(audio_chunks))

        # 2. 调用本地 ASR 接口
        asr_result = await call_local_asr_with_audio(wav_path, session)
        if not asr_result or not asr_result.get("success"):
            print("❌ 本地 ASR 处理失败")
            return

        # 3. 构建前端消息
        segments = asr_result.get("segments", [])
        full_text = asr_result.get("text", "")

        # 推送完整结果（含说话人）
        msg = {
            "type": "asr_with_speaker",
            "full_text": full_text,
            "segments": segments,  # 每段含 speaker, text, start, end 等
            "original_asr_text": asr_text  # 来自远程 ASR 的原始文本（可选）
        }

        await websocket.send(json.dumps(msg, ensure_ascii=False))
        print(f"✅ 本地 ASR + 说话人识别完成，共 {len(segments)} 段")

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


# ========================
# 主 WebSocket 转发逻辑
# ========================
async def forward_audio(websocket):
    print("🟢 前端已连接。")
    current_audio_chunks = []
    lock = asyncio.Lock()

    async with aiohttp.ClientSession() as session:
        try:
            async with websockets.connect(TARGET_WS_URL) as target_ws:
                print(f"🔗 连接到远程 ASR: {TARGET_WS_URL}")
                await target_ws.send(json.dumps(config))

                async def frontend_to_target():
                    async for message in websocket:
                        if isinstance(message, bytes):
                            async with lock:
                                current_audio_chunks.append(message)
                            await target_ws.send(message)
                        else:
                            await target_ws.send(message)

                async def target_to_frontend():
                    async for message in target_ws:
                        # 立即转发原始 ASR 消息（保持低延迟）
                        await websocket.send(message)

                        if isinstance(message, str):
                            try:
                                data = json.loads(message)
                                if data.get("mode") == "2pass-offline":
                                    asr_text = data.get("text", "").strip()
                                    print(f"🔊 触发本地 ASR 处理: '{asr_text}'")

                                    async with lock:
                                        audio_copy = current_audio_chunks.copy()
                                        current_audio_chunks.clear()

                                    # 启动后台任务：调用你自己的 /asr 接口
                                    asyncio.create_task(
                                        process_audio_with_local_asr(audio_copy, asr_text, websocket, session)
                                    )
                            except json.JSONDecodeError:
                                pass

                await asyncio.gather(frontend_to_target(), target_to_frontend())

        except Exception as e:
            print(f"❌ 连接异常: {e}")
        finally:
            print("🔴 前端断开连接。")


# ========================
# 启动服务
# ========================
async def main():
    print(f"🚀 启动 WebSocket 中转服务: ws://{LOCAL_HOST}:{LOCAL_PORT}")
    print(f"📝 本地 ASR 接口: {ASR_LOCAL_URL}")
    server = await websockets.serve(forward_audio, LOCAL_HOST, LOCAL_PORT)
    await server.wait_closed()


if __name__ == "__main__":
    asyncio.run(main())