"""
智谱 GLM-4.7 客户端封装

核心功能：
1. 将 Skill 的行为指导融入提示词
2. 提供同步和异步调用接口
3. 支持流式和非流式响应
4. 错误处理和重试机制
"""
from typing import List, Dict, Any, Optional, AsyncGenerator, Union
from zai import ZhipuAiClient
from app.config import settings
from app.utils.logger import get_logger

logger = get_logger(__name__)


class GLMClient:
    """
    智谱 GLM-4.7 客户端

    核心功能：将 Skill 的行为指导融入提示词，指导 LLM 创作
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        model: Optional[str] = None
    ):
        """
        初始化 GLM 客户端

        Args:
            api_key: 智谱 AI API Key
            model: 模型名称 (glm-4-flash, glm-4-plus, glm-4-0520)
        """
        self.api_key = api_key or settings.zai_api_key
        self.model = model or settings.zai_model
        self.client = ZhipuAiClient(api_key=self.api_key)

        logger.info(f"GLM 客户端初始化完成，模型: {self.model}")

    async def chat(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        top_p: float = 0.9,
        max_tokens: Optional[int] = None,
        stream: bool = False,
        **kwargs
    ) -> Dict[str, Any]:
        """
        发送聊天请求

        Args:
            messages: 消息列表 [{"role": "system", "content": "..."}, ...]
            temperature: 温度参数 (0-1)，越高越随机
            top_p: 采样参数
            max_tokens: 最大 token 数
            stream: 是否流式输出
            **kwargs: 其他参数

        Returns:
            响应结果字典
        """
        try:
            logger.debug(f"发送 GLM 请求，模型: {self.model}, 温度: {temperature}")

            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_tokens,
                stream=stream,
                **kwargs
            )

            if stream:
                return {"stream": self._process_stream_response(response)}
            else:
                return self._process_response(response)

        except Exception as e:
            logger.error(f"GLM API 调用失败: {str(e)}")
            raise

    async def chat_with_skill(
        self,
        skill_behavior: str,
        user_input: str,
        context: Optional[Dict[str, Any]] = None,
        temperature: float = 0.7,
        system_prompt_addition: Optional[str] = None,
        skill_id: Optional[str] = None,
        references: Optional[Dict[str, str]] = None
    ) -> str:
        """
        使用 Skill 行为指导进行聊天

        这是核心功能：将 Skill 的行为指导融入提示词

        Args:
            skill_behavior: Skill 的行为指导内容
            user_input: 用户输入
            context: 上下文信息 (可选)
            temperature: 温度参数
            system_prompt_addition: 额外的系统提示词 (可选)
            skill_id: Skill ID（用于加载参考文件）
            references: 参考文件内容字典 {filename: content}（可选）

        Returns:
            LLM 响应文本
        """
        # 构建包含 Skill 行为指导和参考文件的系统提示词
        system_prompt = self._build_skill_prompt(
            skill_behavior,
            context or {},
            system_prompt_addition,
            references
        )

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_input}
        ]

        try:
            response = await self.chat(messages, temperature=temperature)
            content = response["choices"][0]["message"]["content"]

            logger.info(f"GLM 响应成功，token 使用: {response.get('usage', {})}")

            return content

        except Exception as e:
            logger.error(f"chat_with_skill 调用失败: {str(e)}")
            raise

    def _build_skill_prompt(
        self,
        skill_behavior: str,
        context: Dict[str, Any],
        addition: Optional[str] = None,
        references: Optional[Dict[str, str]] = None
    ) -> str:
        """
        构建包含 Skill 行为指导和参考文件的提示词

        这是平台的核心机制：将 Skill 的行为指导和参考文件融入提示词

        Args:
            skill_behavior: Skill 的行为指导内容
            context: 上下文信息
            addition: 额外的系统提示词
            references: 参考文件内容字典

        Returns:
            完整的系统提示词
        """
        parts = [
            "你是剧集创作专家。",
            "",
            "## Skill 行为指导",
            "",
            "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━",
            skill_behavior,
            "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━",
        ]

        # 添加参考文件内容
        if references and references:
            parts.extend([
                "",
                "## 参考资料",
                "",
                "以下是 Skill 的参考文件内容，请在执行任务时参考这些资料：",
                "",
            ])

            for filename, content in references.items():
                # 判断文件类型并格式化
                if filename.endswith('.md') or filename.endswith('.txt'):
                    parts.extend([
                        f"### 📄 {filename}",
                        "```",
                        content,
                        "```",
                        ""
                    ])
                elif filename.endswith('.json') or filename.endswith('.yaml') or filename.endswith('.yml'):
                    parts.extend([
                        f"### 📋 {filename}",
                        "```json" if filename.endswith('.json') else "```yaml",
                        content,
                        "```",
                        ""
                    ])
                else:
                    # 其他文件类型，直接显示内容
                    parts.extend([
                        f"### 📎 {filename}",
                        content,
                        ""
                    ])

            parts.extend([
                "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━",
            ])

        # 添加额外的系统提示词
        if addition:
            parts.extend(["", addition])

        # 添加上下文信息
        if context:
            parts.extend([
                "",
                "## 当前上下文",
                "",
                self._format_context(context)
            ])

        return "\n".join(parts)

    def _format_context(self, context: Dict[str, Any]) -> str:
        """格式化上下文信息"""
        lines = []
        for key, value in context.items():
            if isinstance(value, dict):
                lines.append(f"- {key}:")
                for sub_key, sub_value in value.items():
                    lines.append(f"  - {sub_key}: {sub_value}")
            elif isinstance(value, list):
                lines.append(f"- {key}:")
                for item in value:
                    lines.append(f"  - {item}")
            else:
                lines.append(f"- {key}: {value}")
        return "\n".join(lines)

    def _process_response(self, response) -> Dict[str, Any]:
        """处理标准响应"""
        return {
            "id": response.id,
            "choices": [
                {
                    "message": {
                        "role": choice.message.role,
                        "content": choice.message.content
                    },
                    "finish_reason": choice.finish_reason,
                    "index": choice.index
                }
                for choice in response.choices
            ],
            "usage": {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens
            },
            "model": response.model,
            "created": response.created
        }

    async def _process_stream_response(self, response) -> AsyncGenerator[str, None]:
        """处理流式响应"""
        async for chunk in response:
            if chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content

    async def test_connection(self) -> bool:
        """测试 API 连接"""
        try:
            response = await self.chat(
                messages=[{"role": "user", "content": "你好"}],
                max_tokens=10
            )
            logger.info("GLM API 连接测试成功")
            return True
        except Exception as e:
            logger.error(f"GLM API 连接测试失败: {str(e)}")
            return False


# 全局单例
_glm_client: Optional[GLMClient] = None


def get_glm_client() -> GLMClient:
    """获取 GLM 客户端单例"""
    global _glm_client
    if _glm_client is None:
        _glm_client = GLMClient()
    return _glm_client


# 便捷导出
glm_client = get_glm_client()
