#!/usr/bin/env python3
"""
Grok LLM Agent 集成
基于LangChain框架集成Grok API
"""

import os
from langchain_openai import ChatOpenAI
from langchain.agents import AgentExecutor, create_react_agent
from langchain.tools import Tool
from langchain import hub
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate

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

        Args:
            api_key: API密钥，如果不提供则从环境变量读取
            api_base: API基础URL，如果不提供则从环境变量读取
        """
        # 从参数或环境变量获取配置
        self.api_key = api_key or os.getenv('GROK_API_KEY')
        self.api_base = api_base or os.getenv('GROK_API_BASE', 'https://api.34ku.com/v1')

        if not self.api_key:
            raise ValueError("必须提供GROK_API_KEY")

        # 初始化Grok LLM (使用OpenAI兼容接口)
        self.llm = ChatOpenAI(
            model="grok-beta",
            openai_api_key=self.api_key,
            openai_api_base=self.api_base,
            temperature=0.7
        )

        # 初始化记忆
        self.memory = ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True
        )

        # 定义工具
        self.tools = self._create_tools()

        # 创建agent
        self.agent = self._create_agent()

    def _create_tools(self):
        """创建Agent可用的工具"""

        def calculator(expression: str) -> str:
            """计算数学表达式"""
            try:
                result = eval(expression)
                return f"计算结果: {result}"
            except Exception as e:
                return f"计算错误: {e}"

        def multiplication_table(number: str) -> str:
            """生成指定数字的乘法表"""
            try:
                n = int(number)
                result = []
                for i in range(1, n + 1):
                    row = []
                    for j in range(1, i + 1):
                        row.append(f"{j} × {i} = {j * i}")
                    result.append("  ".join(row))
                return "\n".join(result)
            except Exception as e:
                return f"生成乘法表错误: {e}"

        tools = [
            Tool(
                name="Calculator",
                func=calculator,
                description="用于计算数学表达式，输入应该是有效的Python数学表达式"
            ),
            Tool(
                name="MultiplicationTable",
                func=multiplication_table,
                description="生成乘法表，输入应该是一个数字，如9表示生成99乘法表"
            )
        ]

        return tools

    def _create_agent(self):
        """创建ReAct Agent"""
        # 使用自定义提示模板
        template = """你是一个有帮助的AI助手，可以使用各种工具来回答问题。

你可以使用以下工具:
{tools}

工具名称: {tool_names}

使用以下格式:

Question: 你需要回答的问题
Thought: 你应该思考该做什么
Action: 要采取的行动，应该是[{tool_names}]之一
Action Input: 行动的输入
Observation: 行动的结果
... (这个Thought/Action/Action Input/Observation可以重复N次)
Thought: 我现在知道最终答案了
Final Answer: 对原始问题的最终答案

开始!

Question: {input}
Thought: {agent_scratchpad}"""

        prompt = PromptTemplate.from_template(template)

        agent = create_react_agent(
            llm=self.llm,
            tools=self.tools,
            prompt=prompt
        )

        agent_executor = AgentExecutor(
            agent=agent,
            tools=self.tools,
            memory=self.memory,
            verbose=True,
            handle_parsing_errors=True,
            max_iterations=5
        )

        return agent_executor

    def chat(self, message: str) -> str:
        """
        与Agent对话

        Args:
            message: 用户消息

        Returns:
            Agent的回复
        """
        try:
            response = self.agent.invoke({"input": message})
            return response.get('output', '抱歉，我无法回答这个问题')
        except Exception as e:
            return f"错误: {e}"

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

        Args:
            message: 用户消息

        Returns:
            Grok的回复
        """
        try:
            response = self.llm.invoke(message)
            return response.content
        except Exception as e:
            return f"错误: {e}"


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("=" * 60)

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

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

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

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


if __name__ == "__main__":
    main()