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

from dotenv import load_dotenv

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    force=True
)
logger = logging.getLogger("app")


# ========================
# 配置
# ========================
load_dotenv()
logger.info("正在加载环境变量 ")
TARGET_WS_URL = os.getenv("TARGET_WS_URL", "ws://localhost:59805/ws/asr")
LOCAL_HOST = "0.0.0.0"
LOCAL_PORT = 10095
ASR_LOCAL_URL = os.getenv("ASR_LOCAL_URL", "http://172.18.127.124:6688/asr")

SAMPLE_RATE = 16000
SAMPLE_WIDTH = 2      # 16-bit = 2 bytes
CHANNELS = 1

# 计算每秒的字节数
BYTES_PER_SECOND = SAMPLE_RATE * SAMPLE_WIDTH * CHANNELS
# 设定最小触发时长（秒）
MIN_DURATION_THRESHOLD = 3.0 

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


# ========================
# 新增：声纹会话缓存管理类 (带超时机制 + 日志)
# ========================
class SpeakerSession:
    def __init__(self):
        self.last_speaker = None
        self.buffer = []
        self.accumulated_duration = 0.0
        self.lock = asyncio.Lock()
        self.MAX_BUFFER_DURATION = 3.0 

    async def process_and_send(self, new_msg, websocket, duration=0.0, force_flush=False):
        """
        duration: 当前这段消息对应的音频时长(秒)
        """
        async with self.lock:
            # === 1. 强制清空 (连接断开时) ===
            if force_flush:
                if self.buffer:
                    logger.info(f"🧹 [最终清空] 准备推送 {len(self.buffer)} 条缓存...")
                    for msg in self.buffer:
                        text_content = msg.get("full_text", "")
                        logger.info(f"   └── 🚀 发送: {text_content}") # 打印发送内容
                        await websocket.send(json.dumps(msg, ensure_ascii=False))
                    self.buffer = []
                    self.accumulated_duration = 0.0
                return

            # 获取当前消息的说话人
            segments = new_msg.get("segments", [])
            current_speaker = segments[0].get("spk") if segments else "unknown"

            # === 2. 如果是同一个人（或者刚开始） ===
            if self.last_speaker is None or current_speaker == self.last_speaker:
                self.buffer.append(new_msg)
                self.accumulated_duration += duration
                self.last_speaker = current_speaker
                
                # 🌟 检查是否超时 20s
                if self.accumulated_duration >= self.MAX_BUFFER_DURATION:
                    logger.info(f"⏰ [超时推送] 同人 ({current_speaker}) 积攒 {self.accumulated_duration:.2f}s >= 8s")
                    
                    for msg in self.buffer:
                        text_content = msg.get("full_text", "")
                        logger.info(f"   └── 🚀 发送: {text_content}") # 打印发送内容
                        await websocket.send(json.dumps(msg, ensure_ascii=False))
                    
                    # 清空缓存，重置时长
                    self.buffer = []
                    self.accumulated_duration = 0.0
                else:
                    # 仅缓存，不发送
                    text_preview = new_msg.get("full_text", "")
                    logger.info(f"📥 [缓存中] Spk:{current_speaker} | 文本:{text_preview} | 累计:{self.accumulated_duration:.2f}s")

            # === 3. 如果换人了 ===
            else:
                logger.info(f"🔄 [换人推送] {self.last_speaker} -> {current_speaker}")
                
                # A. 推送上一位说话人的所有积攒消息
                for msg in self.buffer:
                    text_content = msg.get("full_text", "")
                    logger.info(f"   └── 🚀 发送: {text_content}") # 打印发送内容
                    await websocket.send(json.dumps(msg, ensure_ascii=False))
                
                # B. 清空并开始缓存新人的消息
                self.buffer = []
                self.buffer.append(new_msg)
                self.accumulated_duration = duration 
                self.last_speaker = current_speaker
                
                # 打印新人入队日志
                new_text = new_msg.get("full_text", "")
                logger.info(f"📥 [新缓存] Spk:{current_speaker} | 文本:{new_text}")

# ========================
# 调用本地 ASR 接口（带说话人识别）
# ========================
async def call_local_asr_with_audio(wav_path: str, session: aiohttp.ClientSession):
    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"
            )
            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()
                    logger.error(f"❌ 本地 ASR 接口错误: {resp.status} - {text}")
                    return None
    except Exception as e:
        logger.error(f"❌ 调用本地 ASR 异常: {e}")
        return None


# ========================
# 后台处理：音频 → 本地 ASR → Session 缓存判断
# ========================
async def process_audio_with_local_asr(audio_bytes, asr_text, websocket, session, speaker_session):
    if not audio_bytes:
        return

    wav_path = None
    try:
        # 1. 计算这段音频的时长，用于 Session 判断超时
        audio_duration = len(audio_bytes) / BYTES_PER_SECOND

        # 2. 保存为 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(audio_bytes)

        # 3. 调用本地 ASR 接口
        asr_result = await call_local_asr_with_audio(wav_path, session)
        if not asr_result or not asr_result.get("success"):
            logger.error("❌ 本地 ASR (带声纹) 处理失败")
            return

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

        msg = {
            "mode": "offline-speaker", 
            "full_text": full_text,
            "segments": segments,
            "original_asr_text": asr_text,
            "type": "asr_with_speaker",
        }

        # 5. 🌟 传入 duration 参数
        await speaker_session.process_and_send(msg, websocket, duration=audio_duration)

    except websockets.exceptions.ConnectionClosed:
        logger.warning("⚠️ WebSocket 连接已关闭")
    except Exception as e:
        logger.error(f"❌ 后台处理异常: {e}")
    finally:
        if wav_path and os.path.exists(wav_path):
            os.unlink(wav_path)


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

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

                # ---------------------------
                # 任务1: 前端 -> 远程 ASR
                # ---------------------------
                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)

                # ---------------------------
                # 任务2: 远程 ASR -> 前端
                # ---------------------------
                async def target_to_frontend():
                    async for message in target_ws:
                        # 实时推送
                        await websocket.send(message) 

                        if isinstance(message, str):
                            try:
                                data = json.loads(message)
                                mode = data.get("mode")

                                if mode == "2pass-offline":
                                    asr_text = data.get("text", "").strip()
                                    
                                    async with lock:
                                        chunk_bytes = b"".join(current_audio_chunks)
                                        current_audio_chunks.clear()
                                    
                                    local_accumulated_audio.extend(chunk_bytes)
                                    current_duration = len(local_accumulated_audio) / BYTES_PER_SECOND
                                    
                                    if current_duration < MIN_DURATION_THRESHOLD:
                                        logger.info(f"⏳ 音频拼接中: {current_duration:.2f}s")
                                    else:
                                        logger.info(f"⚡ 触发本地声纹识别: {current_duration:.2f}s")
                                        
                                        audio_to_send = bytes(local_accumulated_audio)
                                        local_accumulated_audio.clear()

                                        asyncio.create_task(
                                            process_audio_with_local_asr(audio_to_send, asr_text, websocket, session, speaker_session)
                                        )

                                elif mode == "2pass-online":
                                    pass

                            except json.JSONDecodeError:
                                pass

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

        except websockets.exceptions.ConnectionClosedOK:
            logger.info("🔗 远程 ASR 连接正常关闭。")
        except Exception as e:
            logger.error(f"❌ 连接异常: {e}")
        finally:
            logger.info("🔴 前端断开连接，清理资源...")
            
            if len(local_accumulated_audio) > 0:
                final_duration = len(local_accumulated_audio) / BYTES_PER_SECOND
                logger.info(f"⚠️ 处理断开前的剩余音频: {final_duration:.2f}s")
                try:
                    audio_to_send = bytes(local_accumulated_audio)
                    local_accumulated_audio.clear()
                    await process_audio_with_local_asr(audio_to_send, "Final-Flush", websocket, session, speaker_session)
                except Exception as e:
                    logger.error(f"❌ 剩余音频处理异常: {e}")

            # 强制推送最后残留缓存
            try:
                await speaker_session.process_and_send(None, websocket, force_flush=True)
            except Exception as e:
                logger.error(f"❌ 最终缓存推送失败: {e}")


# ========================
# 启动服务
# ========================
async def main():
    logger.info(f"🚀 启动 WebSocket 中转服务: ws://{LOCAL_HOST}:{LOCAL_PORT}")
    logger.info("📝 逻辑: 声纹识别缓存，同人累积满8s或换人时推送")
    server = await websockets.serve(forward_audio, LOCAL_HOST, LOCAL_PORT)
    await server.wait_closed()


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