import asyncio
import json
from typing import Optional, Dict, Callable, Awaitable
from dataclasses import dataclass, field
import aiohttp  # 新增：异步HTTP库

# 大模型配置（集中管理）
LLM_CONFIG = {
    "base_url": "http://10.10.10.202:8088/v1",
    "api_key": "app-m7HZNV1aGiheh3wr6wNVHFxX",
    "timeout": 30,  # 请求超时时间（秒）
    "default_scene": "通用聊天场景",  # 默认场景描述
    "stream_chunk_size": 1024  # 流式接收块大小
}

# 定义流式回调函数类型（异步）
LLMStreamCallback = Callable[[str, Optional[str], bool], Awaitable[None]]
"""
回调函数参数说明：
- chunk: 单次流式返回的文本片段
- conversation_id: 会话ID（首次返回，后续复用）
- is_finished: 是否结束（True=流式结束/同步返回完成）
"""

@dataclass
class LLMConversation:
    """会话对象（管理会话ID和上下文）"""
    conversation_id: Optional[str] = None
    user_id: str = ""
    scene_description: str = LLM_CONFIG["default_scene"]
    # 可选：存储会话历史（如需上下文管理）
    history: list = field(default_factory=list)

class LLMClient:
    """大模型客户端封装（异步修复版）"""
    def __init__(self):
        self.base_url = LLM_CONFIG["base_url"]
        self.headers = {
            "Authorization": f"Bearer {LLM_CONFIG['api_key']}",
            "Content-Type": "application/json"
        }
        self.timeout = aiohttp.ClientTimeout(total=LLM_CONFIG["timeout"])  # 异步超时
        self._session: Optional[aiohttp.ClientSession] = None  # 异步会话（复用连接）

    async def _get_session(self) -> aiohttp.ClientSession:
        """获取/复用异步HTTP会话"""
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession(timeout=self.timeout)
        return self._session

    async def send_message(
            self,
            query: str,
            conversation: LLMConversation,
            stream_callback: Optional[LLMStreamCallback] = None,
            response_mode: str = "streaming"
    ) -> tuple[Optional[str], str]:
        """
        发送消息到大模型（纯异步版，无线程池阻塞）
        """
        payload = {
            "query": query,
            "inputs": {"scene_description": conversation.scene_description},
            "response_mode": response_mode,
            "user": conversation.user_id
        }
        if conversation.conversation_id:
            payload["conversation_id"] = conversation.conversation_id

        url = f"{self.base_url}/chat-messages"
        full_response = ""
        res_conversation_id = conversation.conversation_id

        try:
            session = await self._get_session()
            if response_mode == "streaming":
                # 异步流式请求（无线程池，纯异步IO）
                async with session.post(url, headers=self.headers, json=payload) as response:
                    response.raise_for_status()
                    # 实时迭代流式响应
                    async for line in response.content.iter_chunked(LLM_CONFIG["stream_chunk_size"]):

                        if not line:
                            continue
                        line_data = line.decode("utf-8")
                        if line_data.startswith("data: "):
                            json_str = line_data[6:].strip()
                            if json_str == "[DONE]":
                                if stream_callback:
                                    await stream_callback("", res_conversation_id, True)
                                break
                            try:
                                data = json.loads(json_str)
                                # 更新会话ID
                                if not res_conversation_id and "conversation_id" in data:
                                    res_conversation_id = data["conversation_id"]
                                # 提取内容
                                chunk = data.get("content", data.get("answer", data.get("message", "")))
                                # print('大模型返回的', chunk)
                                if chunk:
                                    full_response += chunk
                                    if stream_callback:
                                        await stream_callback(chunk, res_conversation_id, False)
                                        await asyncio.sleep(0)  # 让出调度权
                            except json.JSONDecodeError as e:
                                # print(f"大模型解析流式数据失败: {e}")
                                continue
            else:
                # 异步非流式请求
                async with session.post(url, headers=self.headers, json=payload) as response:
                    response.raise_for_status()
                    data = await response.json()
                    res_conversation_id = data.get("conversation_id", conversation.conversation_id)
                    full_response = data.get("content", data.get("answer", data.get("message", "")))
                    if stream_callback:
                        await stream_callback(full_response, res_conversation_id, True)

            conversation.conversation_id = res_conversation_id
            return res_conversation_id, full_response

        except aiohttp.ClientError as e:
            error_msg = f"大模型请求失败: {str(e)}"
            print(error_msg)
            if stream_callback:
                await stream_callback(f"[错误] {error_msg}", res_conversation_id, True)
            return res_conversation_id, ""
        except Exception as e:
            error_msg = f"大模型处理异常: {str(e)}"
            print(error_msg)
            if stream_callback:
                await stream_callback(f"[错误] {error_msg}", res_conversation_id, True)
            return res_conversation_id, ""

    async def close(self):
        """关闭异步会话（程序退出时调用）"""
        if self._session and not self._session.closed:
            await self._session.close()

# 全局单例客户端（异步版）
llm_client = LLMClient()

# 快捷调用函数（保持原有接口不变）
async def call_llm(
        query: str,
        user_id: str,
        scene_description: str = LLM_CONFIG["default_scene"],
        conversation_id: Optional[str] = None,
        stream_callback: Optional[LLMStreamCallback] = None,
        response_mode: str = "streaming"
) -> tuple[Optional[str], str]:
    """
    快捷调用大模型（无需手动创建会话对象）
    :param query: 用户提问
    :param user_id: 用户ID
    :param scene_description: 场景描述
    :param conversation_id: 会话ID（续聊用）
    :param stream_callback: 流式回调
    :param response_mode: 响应模式
    :return: (conversation_id, 完整回复)
    """
    conversation = LLMConversation(
        conversation_id=conversation_id,
        user_id=user_id,
        scene_description=scene_description
    )
    return await llm_client.send_message(
        query=query,
        conversation=conversation,
        stream_callback=stream_callback,
        response_mode=response_mode
    )

# 可选：程序退出时关闭会话（如FastAPI的shutdown事件）
async def shutdown_llm_client():
    await llm_client.close()