#!/usr/bin/env python3
"""
Grok LLM Agent 集成 (简化版)
使用requests直接调用Grok API，避免LangChain兼容性问题
"""

import os
import json
import requests
from typing import List, Dict, Optional, Callable


class Tool:
    """工具类"""
    def __init__(self, name: str, func: Callable, description: str):
        self.name = name
        self.func = func
        self.description = description

    def run(self, input_str: str) -> str:
        """执行工具"""
        return self.func(input_str)


class GrokAgent:
    def __init__(self, api_key: str, api_base: str):
        """
        初始化Grok Agent

        Args:
            api_key: API密钥
            api_base: API基础URL
        """
        self.api_key = api_key
        self.api_base = api_base
        self.tools = self._create_tools()
        self.conversation_history = []

    def _create_tools(self) -> List[Tool]:
        """创建Agent可用的工具"""

        def calculator(expression: str) -> str:
            """计算数学表达式"""
            try:
                # 安全的数学计算
                allowed_chars = "0123456789+-*/().,[] "
                if not all(c in allowed_chars for c in expression):
                    return "错误: 表达式包含非法字符"
                result = eval(expression)
                return f"计算结果: {result}"
            except Exception as e:
                return f"计算错误: {e}"

        def multiplication_table(number: str) -> str:
            """生成指定数字的乘法表"""
            try:
                n = int(number.strip())
                if n < 1 or n > 20:
                    return "请输入1-20之间的数字"

                result = []
                for i in range(1, n + 1):
                    row = []
                    for j in range(1, i + 1):
                        row.append(f"{j} × {i} = {j * i:2d}")
                    result.append("  ".join(row))
                return "\n" + "\n".join(result)
            except Exception as e:
                return f"生成乘法表错误: {e}"

        def search_info(query: str) -> str:
            """模拟信息搜索"""
            return f"关于'{query}'的信息: 这是一个示例搜索结果"

        tools = [
            Tool(
                name="calculator",
                func=calculator,
                description="用于计算数学表达式。输入: 有效的数学表达式，如'2+2'或'10*5+3'"
            ),
            Tool(
                name="multiplication_table",
                func=multiplication_table,
                description="生成乘法表。输入: 一个数字，如'9'表示生成9×9乘法表"
            ),
            Tool(
                name="search",
                func=search_info,
                description="搜索信息。输入: 搜索查询词"
            )
        ]

        return tools

    def _get_tools_description(self) -> str:
        """获取工具描述"""
        desc = "你可以使用以下工具:\n\n"
        for tool in self.tools:
            desc += f"- {tool.name}: {tool.description}\n"
        return desc

    def _call_grok_api(self, messages: List[Dict], temperature: float = 0.7) -> Optional[str]:
        """
        调用Grok API

        Args:
            messages: 消息列表
            temperature: 温度参数

        Returns:
            Grok的回复内容
        """
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        data = {
            "model": "grok-beta",
            "messages": messages,
            "stream": False,
            "temperature": temperature
        }

        try:
            response = requests.post(
                f"{self.api_base}/v1/chat/completions",
                headers=headers,
                json=data,
                timeout=30
            )

            if response.status_code == 200:
                result = response.json()
                if 'choices' in result and len(result['choices']) > 0:
                    return result['choices'][0]['message']['content']
            else:
                print(f"API错误: {response.status_code} - {response.text}")
                return None

        except Exception as e:
            print(f"API调用异常: {e}")
            return None

    def simple_chat(self, message: str) -> str:
        """
        简单对话（不使用工具）

        Args:
            message: 用户消息

        Returns:
            Grok的回复
        """
        messages = [{"role": "user", "content": message}]
        response = self._call_grok_api(messages)
        return response if response else "抱歉，无法获取回复"

    def chat_with_tools(self, user_message: str, max_iterations: int = 3) -> str:
        """
        使用工具的对话

        Args:
            user_message: 用户消息
            max_iterations: 最大迭代次数

        Returns:
            最终回复
        """
        # 构建系统提示
        system_prompt = f"""你是一个有帮助的AI助手。你可以使用工具来帮助回答问题。

{self._get_tools_description()}

当你需要使用工具时，请按以下格式回复:
USE_TOOL: tool_name
INPUT: tool_input

例如:
USE_TOOL: calculator
INPUT: 2+2

如果不需要工具，直接回答用户的问题。
"""

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

        iteration = 0
        while iteration < max_iterations:
            iteration += 1

            # 调用Grok API
            response = self._call_grok_api(messages)
            if not response:
                return "抱歉，无法获取回复"

            # 检查是否需要使用工具
            if "USE_TOOL:" in response and "INPUT:" in response:
                try:
                    # 解析工具调用
                    lines = response.split('\n')
                    tool_name = None
                    tool_input = None

                    for line in lines:
                        if line.startswith("USE_TOOL:"):
                            tool_name = line.replace("USE_TOOL:", "").strip()
                        elif line.startswith("INPUT:"):
                            tool_input = line.replace("INPUT:", "").strip()

                    if tool_name and tool_input:
                        # 查找并执行工具
                        tool = next((t for t in self.tools if t.name == tool_name), None)
                        if tool:
                            print(f"\n🔧 使用工具: {tool_name}")
                            print(f"📥 输入: {tool_input}")

                            tool_result = tool.run(tool_input)
                            print(f"📤 输出: {tool_result}\n")

                            # 将工具结果添加到对话中
                            messages.append({"role": "assistant", "content": response})
                            messages.append({
                                "role": "user",
                                "content": f"工具 {tool_name} 的执行结果:\n{tool_result}\n\n请基于这个结果继续回答用户的问题。"
                            })
                            continue
                        else:
                            return f"错误: 未找到工具 '{tool_name}'"
                except Exception as e:
                    return f"工具解析错误: {e}"

            # 没有工具调用，返回最终答案
            return response

        return "达到最大迭代次数，对话结束"


def main():
    """示例用法"""
    import sys

    # 从配置文件读取API信息
    try:
        with open('34ku.txt', 'r') as f:
            lines = f.readlines()
            api_key = lines[0].strip()
            api_base = lines[3].strip()
    except Exception as e:
        print(f"❌ 读取配置文件失败: {e}")
        sys.exit(1)

    print("🤖 初始化Grok Agent...")
    agent = GrokAgent(api_key=api_key, api_base=api_base)

    print("\n✅ Grok Agent已就绪!")
    print("=" * 70)

    # 测试1: 简单对话
    print("\n📝 测试1: 简单对话")
    print("-" * 70)
    response = agent.simple_chat("你好，请用一句话介绍你自己")
    print(f"Grok: {response}")

    # 测试2: 使用工具生成乘法表
    print("\n" + "=" * 70)
    print("\n📝 测试2: 使用工具生成乘法表")
    print("-" * 70)
    response = agent.chat_with_tools("请帮我生成9的乘法表")
    print(f"Grok: {response}")

    # 测试3: 使用计算器
    print("\n" + "=" * 70)
    print("\n📝 测试3: 使用计算器")
    print("-" * 70)
    response = agent.chat_with_tools("帮我计算 123 * 456 + 789")
    print(f"Grok: {response}")

    # 测试4: 复杂任务
    print("\n" + "=" * 70)
    print("\n📝 测试4: 复杂任务")
    print("-" * 70)
    response = agent.chat_with_tools("生成5的乘法表，然后计算5*5的结果")
    print(f"Grok: {response}")

    print("\n" + "=" * 70)
    print("\n✅ 所有测试完成!")


if __name__ == "__main__":
    main()