import asyncio
import json
import os
import websockets
import redis.asyncio as redis
import re
from openai import AsyncOpenAI  # ⬅️ 引入 OpenAI 异步客户端

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

# Redis
REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0")
REDIS_KEY_PREFIX = "meeting:"

# LLM 配置 (适配 OpenAI SDK)
# 如果是 OpenAI 官方，Base URL 默认即可；如果是 DeepSeek/阿里等，需修改 BASE_URL
LLM_API_KEY = os.getenv("LLM_API_KEY", "sk-xxxx")
LLM_BASE_URL = os.getenv("LLM_BASE_URL", "http://172.18.127.124:9997/v1") 
LLM_MODEL = os.getenv("LLM_MODEL", "Qwen3-32B")

# 候选人
CANDIDATE_SPEAKERS = ["局长", "主持人", "商务专员", "市场专员"]
DEFAULT_SESSION_ID = "room_101"

# ASR 初始化配置
config = {
    "chunk_size": [10, 10, 10],
    "wav_name": "h5",
    "is_speaking": True,
    "wav_format": "pcm",
    "chunk_interval": 10,
    "itn": True,
    "mode": "2pass",
    "hotwords": "",
}

# ========================
# Redis 操作封装 (保持不变)
# ========================
async def append_to_buffer(session_id, text):
    """将新的片段追加到暂存区"""
    r = redis.from_url(REDIS_URL, decode_responses=True)
    key = f"{REDIS_KEY_PREFIX}buffer:{session_id}"
    await r.append(key, text + " ")
    current_buffer = await r.get(key)
    await r.close()
    return current_buffer

async def clear_buffer(session_id):
    """清空暂存区"""
    r = redis.from_url(REDIS_URL, decode_responses=True)
    key = f"{REDIS_KEY_PREFIX}buffer:{session_id}"
    await r.delete(key)
    await r.close()

async def save_history(session_id, role, content):
    """归档历史记录"""
    r = redis.from_url(REDIS_URL, decode_responses=True)
    key = f"{REDIS_KEY_PREFIX}history:{session_id}"
    record = json.dumps({"role": role, "content": content})
    await r.rpush(key, record)
    await r.ltrim(key, -20, -1)
    await r.close()

async def get_history(session_id):
    """获取上下文"""
    r = redis.from_url(REDIS_URL, decode_responses=True)
    key = f"{REDIS_KEY_PREFIX}history:{session_id}"
    items = await r.lrange(key, 0, -1)
    await r.close()
    return [json.loads(i) for i in items]

# ========================
# 通用 LLM 调用函数 (改为 OpenAI SDK)
# ========================
async def call_llm(client: AsyncOpenAI, system_prompt: str, user_text: str):
    """
    使用 OpenAI SDK 发起调用
    """
    try:
        response = await client.chat.completions.create(
            model=LLM_MODEL,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_text}
            ],
            temperature=0.1,
        )
        
        content = response.choices[0].message.content
        print(f"🤖 LLM 返回: {content}")
        return content
    except Exception as e:
        print(f"❌ OpenAI API 调用异常: {e}")
        return None

# ========================
# Agent A: 完结判断 (The Auditor)
# ========================
async def agent_check_completion(text_buffer, client: AsyncOpenAI):
    """
    判断当前的文本累积是否已经构成了一个完整的发言回合。
    """
    system_prompt = """
    #ROLE
        你是一个对话状态判断助手。你的任务是判断当前的发言内容是否已经结束。
    #RULE
        1. **强制标准**：如果内容中包含 "发言完毕"、"讲完了"、"汇报完毕" 等明确结束语，必须返回 true。
        2. **语义标准**：如果内容完整且看似一段独立的话结束，返回 true；如果看似还在列举或句子未完，返回 false。
        3. 如果内容仅仅是很短的一句 "嗯"、"好的"，视作未结束或无需独立归档。
        4. 请将结果用 <result> </result>包裹着。

    #EXAMPLES
    示例1:空中侦察席报告，今日预计飞行287架航班，空气湿度适宜，安全风险较小。 <result>true</result>
    示例2:控制要素席报告，我这边已经完成了所有的准备工作.<result>false<result>
    示例3:下面有请局长发言 <result>false</result>
    示例4:地面侦察席报告目前没有发现敌对情况.<result>false</result>
    """
    
    result = await call_llm(client, system_prompt, f"当前累积内容：{text_buffer}")
    if result:
        # 更健壮的版本（容忍某些格式错误）
        match = re.search(r'<result>([^<]*)(?:</result>|<result>|$)', result, re.IGNORECASE)
        if match:
            content = match.group(1).strip()
        else:
            content = False  # 或者根据你的逻辑处理未匹配的情况
        return content
    return False

# ========================
# Agent B: 身份识别 (The Profiler)
# ========================
async def agent_identify_speaker(text_full, session_id, client: AsyncOpenAI):
    """
    仅在 Agent A 判定结束时调用。根据全文和历史判断身份。
    """
    history = await get_history(session_id)
    history_str = "\n".join([f"{h['role']}: {h['content']}" for h in history])
    
    system_prompt = f"""
    #ROLE
        你是一个会议记录专家。请根据上下文和完整的段落内容，推断当前的发言人。
    #RULE
        1. **强制标准**：如果内容中包含 "控制要素席"、"空中侦察席"、"海上观察席"、"联合指挥中心"等明确发言人，必须直接返回名称。
        2. **语义标准**：如果内容完整且没有明确发言人，则根据语义进行判断，潜在的候选人有:{', '.join(CANDIDATE_SPEAKERS)}。
        3. 提供历史上下文：{history_str}
        4. 请将结果用 <result> </result>包裹着。
    
    #EXAMPLES
    示例1:空中侦察席报告，今日预计飞行287架航班，空气湿度适宜，安全风险较小。 <result>空中侦察席</result>
    示例2:控制要素席报告，我这边已经完成了所有的准备工作，可以开始会议了。 <result>控制要素席<result>
    示例3:下面有请局长发言 <result>主持人</result>
    """
    
    result = await call_llm(client, system_prompt, f"完整发言内容：{text_full}")
    if result:
        match = re.search(r'<result>([^<]*)(?:</result>|<result>|$)', result, re.IGNORECASE)
        if match:
            content = match.group(1).strip()
            return content
    return "未知说话人"

# ========================
# 核心编排逻辑 (多 Agent 协作)
# ========================
async def orchestrate_agents(new_fragment, websocket, client: AsyncOpenAI):
    session_id = DEFAULT_SESSION_ID
    
    # 1. 存入 Buffer (累积)
    current_buffer = await append_to_buffer(session_id, new_fragment)
    print(f"📥 [Buffer] 长度: {len(current_buffer)} | 内容片段: {current_buffer[-20:]}")

    # 2. 调用 Agent A 判断是否结束
    is_finished = await agent_check_completion(current_buffer, client)
    
    if is_finished:
        print(f"🛑 [Agent A] 判定发言结束。触发身份识别...")
        
        # 3. 调用 Agent B 识别身份
        speaker = await agent_identify_speaker(current_buffer, session_id, client)
        print(f"👤 [Agent B] 识别为: {speaker}")
        
        # 4. 存入历史并清空 Buffer
        await save_history(session_id, speaker, current_buffer)
        await clear_buffer(session_id)
        
        # 5. 推送最终结果给前端
        msg = {
            "type": "final_speech_record",
            "speaker": speaker,
            "content": current_buffer.strip(),
            "timestamp": "now"
        }
        await websocket.send(json.dumps(msg, ensure_ascii=False))
        
    else:
        # 未结束，什么都不做，继续等待下一段文本
        print(f"⏳ [Agent A] 判定未结束...")

# ========================
# WebSocket 主逻辑
# ========================
async def forward_audio(websocket):
    print("🟢 Client Connected.")
    
    # 初始化 OpenAI 客户端
    # 注意：AsyncOpenAI 是线程/协程安全的，也可以放在全局初始化
    client = AsyncOpenAI(
        api_key=LLM_API_KEY,
        base_url=LLM_BASE_URL
    )

    try:
        async with websockets.connect(TARGET_WS_URL) as target_ws:
            await target_ws.send(json.dumps(config))

            async def frontend_to_target():
                async for message in websocket:
                    await target_ws.send(message)

            async def target_to_frontend():
                async for message in target_ws:
                    # 1. 实时转发 (用于字幕上屏)
                    await websocket.send(message)

                    if isinstance(message, str):
                        try:
                            data = json.loads(message)
                            # 2. 触发 Agent 协作
                            if data.get("mode") == "2pass-offline":
                                asr_text = data.get("text", "").strip()
                                if asr_text:
                                    # 异步触发，传入 client
                                    asyncio.create_task(
                                        orchestrate_agents(asr_text, websocket, client)
                                    )
                        except json.JSONDecodeError:
                            pass

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

    except Exception as e:
        print(f"❌ Error: {e}")
    finally:
        # 关闭客户端（虽然 AsyncOpenAI 通常会自动管理，但显式关闭是好习惯）
        await client.close()
        print("🔴 Client Disconnected.")

# ========================
# 启动
# ========================
async def main():
    print(f"🚀 启动 WebSocket 智能语义分析服务: ws://{LOCAL_HOST}:{LOCAL_PORT}")
    print(f"🧠 Agent 模型: {LLM_MODEL}")
    
    # 这里的 websockets.serve 会在 main 协程内部运行，此时已有 Loop
    async with websockets.serve(forward_audio, LOCAL_HOST, LOCAL_PORT):
        # 创建一个永远 pending 的 Future，让服务一直运行不退出
        await asyncio.Future()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n🛑 服务已停止")