"""
Step 03: 构建简单 Agent

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

🎓 本节内容：
    1. 什么是 Agent？（LLM + Tool + Loop）
    2. Agent 的核心循环：Think → Act → Observe
    3. 如何构建一个完整的 Agent
    4. LangGraph 的基本用法

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
"""

from typing import TypedDict, List, Dict, Any, Literal, Optional
from dataclasses import dataclass
from abc import ABC, abstractmethod
import json


# ═══════════════════════════════════════════════════════════════════════════════
# 第一部分：理解 Agent 的本质
# ═══════════════════════════════════════════════════════════════════════════════

"""
Agent 的核心公式：

    Agent = LLM + Tool + Loop

为什么这个公式很重要？

    1. LLM：负责"思考"——理解用户意图，决定下一步行动
    2. Tool：负责"行动"——执行具体的操作
    3. Loop：负责"循环"——不断执行直到任务完成

没有 LLM，Agent 只能机械地执行预设命令
没有 Tool，LLM 只能思考不能行动
没有 Loop，Agent 只能执行一步

只有三者结合，才是真正的"智能代理"
"""

# ═══════════════════════════════════════════════════════════════════════════════
# 第二部分：定义 Agent State
# ═══════════════════════════════════════════════════════════════════════════════

class SimpleAgentState(TypedDict, total=False):
    """
    简单 Agent 的状态

    相比 Step 02 的状态，这个状态更精简，适合教学
    """

    # 用户输入
    user_input: str

    # 对话历史
    messages: List[dict]
    """
    格式：
        [
            {"role": "user", "content": "用户说了什么"},
            {"role": "assistant", "content": "助手说了什么"},
            {"role": "system", "content": "系统做了什么"},
        ]
    """

    # 当前行动
    current_action: str
    """
    当前正在执行的动作：
        - "thinking": LLM 正在思考
        - "using_tool": 正在使用工具
        - "finished": 完成
    """

    # 工具调用相关
    tool_calls: List[dict]     # 记录所有工具调用
    last_tool_result: Any      # 最后一个工具的执行结果
    tool_result: Any          # 当前工具执行结果

    # 生成结果
    final_output: str          # 最终输出
    generation: str           # 当前正在生成的文本

    # 状态
    status: str
    """
    状态值：
        - "input": 等待输入
        - "thinking": 思考中
        - "acting": 执行中
        - "finished": 完成
        - "error": 错误
    """
    error: str


# ═══════════════════════════════════════════════════════════════════════════════
# 第三部分：定义 Tool 接口
# ═══════════════════════════════════════════════════════════════════════════════

class BaseTool(ABC):
    """
    工具基类（简化版，来自 Step 01）
    """

    @property
    @abstractmethod
    def name(self) -> str:
        """工具名称"""
        pass

    @property
    @abstractmethod
    def description(self) -> str:
        """工具描述"""
        pass

    @abstractmethod
    def execute(self, **kwargs) -> Any:
        """执行工具"""
        pass


# ═══════════════════════════════════════════════════════════════════════════════
# 第四部分：实现具体的 Tool
# ═══════════════════════════════════════════════════════════════════════════════

class CalculatorTool(BaseTool):
    """计算器工具"""

    @property
    def name(self) -> str:
        return "calculator"

    @property
    def description(self) -> str:
        return "执行数学计算。输入一个数学表达式，返回计算结果。"

    def execute(self, expression: str) -> str:
        """执行计算"""
        try:
            result = eval(expression, {"__builtins__": {}}, {})
            return str(result)
        except Exception as e:
            return f"计算错误: {e}"


class JRXMLGeneratorTool(BaseTool):
    """JRXML 生成工具"""

    @property
    def name(self) -> str:
        return "jrxml_generator"

    @property
    def description(self) -> str:
        return "生成 JasperReports JRXML 报表模板。根据用户需求生成 XML 代码。"

    def execute(self, requirement: str, context: str = "") -> str:
        """
        执行 JRXML 生成

        实际应用中，这里会调用 LLM
        这里用模拟数据演示
        """
        # 模拟生成
        template = f'''<?xml version="1.0" encoding="UTF-8"?>
<jasperReport xmlns="http://jasperreports.sourceforge.net/jasperreports"
              name="{requirement}" pageWidth="595" pageHeight="842">
    <title>{requirement}</title>
    <queryString>
        <![CDATA[SELECT * FROM data WHERE 1=1]]>
    </queryString>
    <field name="id" class="java.lang.Integer"/>
    <field name="name" class="java.lang.String"/>
    <band height="50">
        <staticText>
            <reportElement x="0" y="0" width="100" height="20"/>
            <text>名称</text>
        </staticText>
        <textField>
            <reportElement x="0" y="20" width="100" height="20"/>
            <textFieldExpression>$F{{name}}</textFieldExpression>
        </textField>
    </band>
</jasperReport>'''

        return template


class TemplateSearchTool(BaseTool):
    """模板搜索工具"""

    @property
    def name(self) -> str:
        return "template_search"

    @property
    def description(self) -> str:
        return "搜索 JRXML 报表模板。在知识库中搜索相关的报表模板作为参考。"

    def execute(self, keyword: str, limit: int = 3) -> str:
        """搜索模板"""
        # 模拟搜索结果
        results = [
            f"模板1: 关于 '{keyword}' 的基础报表模板",
            f"模板2: '{keyword}' 的高级报表模板",
            f"模板3: '{keyword}' 的数据透视表模板",
        ]
        return "\n".join(results[:limit])


class JRXMLValidatorTool(BaseTool):
    """JRXML 验证工具"""

    @property
    def name(self) -> str:
        return "jrxml_validator"

    @property
    def description(self) -> str:
        return "验证 JRXML 报表模板的正确性。检查 XML 语法和 JasperReports 规范。"

    def execute(self, jrxml: str) -> dict:
        """验证 JRXML"""
        # 模拟验证
        issues = []

        # 检查基本结构
        if "<?xml" not in jrxml:
            issues.append("缺少 XML 声明")

        if "<jasperReport" not in jrxml:
            issues.append("缺少 <jasperReport> 根元素")

        if "<field" not in jrxml and "<band" not in jrxml:
            issues.append("报表中没有定义字段或波段")

        return {
            "valid": len(issues) == 0,
            "issues": issues
        }


# ═══════════════════════════════════════════════════════════════════════════════
# 第五部分：构建 Agent 的大脑（LLM 决策）
# ═══════════════════════════════════════════════════════════════════════════════

class AgentBrain:
    """
    Agent 的大脑——负责决定下一步做什么

    在真实系统中，这里会调用 LLM
    这里用规则引擎模拟 LLM 的决策过程

    决策逻辑：
        1. 如果用户要求生成报表 → 调用 jrxml_generator
        2. 如果用户要求搜索模板 → 调用 template_search
        3. 如果用户要求验证 → 调用 jrxml_validator
        4. 如果用户只是聊天 → 直接回答
        5. 如果工具执行完成 → 返回结果或继续调用
    """

    def __init__(self):
        # 工具描述（用于告诉 LLM 有哪些工具可用）
        self.tools = {
            "calculator": "计算器，执行数学表达式，如 2+3*4",
            "jrxml_generator": "生成 JasperReports JRXML 报表模板",
            "template_search": "搜索报表模板",
            "jrxml_validator": "验证 JRXML 报表的正确性",
        }

    def decide(self, state: SimpleAgentState) -> dict:
        """
        决定下一步行动

        返回：
            {
                "action": "use_tool" | "respond" | "finish",
                "tool_name": str,  # 如果 action == "use_tool"
                "tool_args": dict,  # 如果 action == "use_tool"
                "response": str,    # 如果 action == "respond"
            }
        """
        user_input = state.get("user_input", "").lower()
        messages = state.get("messages", [])
        tool_result = state.get("tool_result")

        # 如果有工具执行结果
        if tool_result is not None:
            return {
                "action": "respond",
                "response": f"工具执行完成，结果：{tool_result}"
            }

        # 根据用户输入决定行动
        if "生成" in user_input or "报表" in user_input:
            return {
                "action": "use_tool",
                "tool_name": "jrxml_generator",
                "tool_args": {
                    "requirement": state.get("user_input", ""),
                    "context": ""
                }
            }

        elif "搜索" in user_input or "模板" in user_input:
            return {
                "action": "use_tool",
                "tool_name": "template_search",
                "tool_args": {
                    "keyword": user_input.replace("搜索", "").replace("模板", "").strip(),
                    "limit": 3
                }
            }

        elif "验证" in user_input and "<" in user_input:
            return {
                "action": "use_tool",
                "tool_name": "jrxml_validator",
                "tool_args": {
                    "jrxml": user_input
                }
            }

        elif any(op in user_input for op in ["计算", "+", "-", "*", "/"]):
            # 提取表达式
            import re
            expr = re.search(r'[\d\+\-\*\/\(\)\.]+', user_input)
            if expr:
                return {
                    "action": "use_tool",
                    "tool_name": "calculator",
                    "tool_args": {
                        "expression": expr.group()
                    }
                }

        else:
            # 默认：直接回答
            return {
                "action": "respond",
                "response": f"我理解你的需求：{state.get('user_input', '')}。请问具体想做什么？"
            }


# ═══════════════════════════════════════════════════════════════════════════════
# 第六部分：构建 Agent
# ═══════════════════════════════════════════════════════════════════════════════

class SimpleAgent:
    """
    简单 Agent 实现

    这个 Agent 的工作流程：
        1. 接收用户输入
        2. LLM（大脑）决定下一步行动
        3. 执行行动（使用工具或直接回答）
        4. 返回结果
        5. 循环直到用户结束
    """

    def __init__(self):
        # 初始化工具
        self.tools = {
            "calculator": CalculatorTool(),
            "jrxml_generator": JRXMLGeneratorTool(),
            "template_search": TemplateSearchTool(),
            "jrxml_validator": JRXMLValidatorTool(),
        }

        # 初始化大脑
        self.brain = AgentBrain()

        # 初始化状态
        self.state = SimpleAgentState(
            messages=[],
            tool_calls=[],
            status="input"
        )

        # 对话轮次计数器
        self.round_count = 0

    def reset(self):
        """重置 Agent 状态"""
        self.state = SimpleAgentState(
            messages=[],
            tool_calls=[],
            status="input"
        )

    def get_available_tools(self) -> list[dict]:
        """获取可用工具列表"""
        return [
            {"name": tool.name, "description": tool.description}
            for tool in self.tools.values()
        ]

    def execute_tool(self, tool_name: str, **kwargs) -> Any:
        """
        执行工具

        封装了工具执行的逻辑，包括：
            - 查找工具
            - 执行工具
            - 记录调用
            - 处理错误
        """
        if tool_name not in self.tools:
            return {"error": f"工具 '{tool_name}' 不存在"}

        tool = self.tools[tool_name]
        try:
            result = tool.execute(**kwargs)

            # 记录工具调用
            self.state["tool_calls"].append({
                "tool": tool_name,
                "args": kwargs,
                "result": result
            })

            return result

        except Exception as e:
            error_result = {"error": str(e)}
            self.state["tool_calls"].append({
                "tool": tool_name,
                "args": kwargs,
                "error": str(e)
            })
            return error_result

    def process(self, user_input: str) -> str:
        """
        处理用户输入

        这是 Agent 的主入口
        每次用户输入，调用这个方法来处理
        """
        # 1. 保存用户输入
        self.state["user_input"] = user_input
        self.state["messages"].append({
            "role": "user",
            "content": user_input
        })
        self.state["status"] = "thinking"
        self.state["tool_result"] = None
        self.round_count += 1

        # 2. 大脑决定行动
        decision = self.brain.decide(self.state)

        # 3. 执行行动
        if decision["action"] == "use_tool":
            # 使用工具
            self.state["status"] = "acting"
            tool_name = decision["tool_name"]
            tool_args = decision["tool_args"]

            result = self.execute_tool(tool_name, **tool_args)

            # 保存结果到状态
            self.state["tool_result"] = result
            self.state["current_action"] = f"使用了工具: {tool_name}"

            # 再次调用大脑，决定下一步
            decision = self.brain.decide(self.state)

        if decision["action"] == "respond":
            response = decision["response"]
            self.state["messages"].append({
                "role": "assistant",
                "content": response
            })
            self.state["final_output"] = response
            self.state["status"] = "finished"
            return response

        elif decision["action"] == "finish":
            self.state["status"] = "finished"
            return self.state.get("final_output", "任务完成")

        # 默认返回当前状态
        return self.state.get("current_action", "处理中...")

    def get_history(self) -> list[dict]:
        """获取对话历史"""
        return self.state.get("messages", [])


# ═══════════════════════════════════════════════════════════════════════════════
# 第七部分：使用 LangGraph 重构（可选进阶内容）
# ═══════════════════════════════════════════════════════════════════════════════

"""
LangGraph 是 LangChain 提供的图结构框架

核心概念：
    1. StateGraph：状态图
    2. Node：节点（执行函数）
    3. Edge：边（状态转换）
    4. Conditional Edge：条件边（根据状态决定下一步）

为什么使用 LangGraph？
    1. 代码结构更清晰
    2. 内置状态管理
    3. 支持复杂的条件分支
    4. 内置循环和回退

如果你想学习 LangGraph：
    请参考 jaspersoft 项目中的实际实现：
        - agent/graph.py：状态图定义
        - agent/nodes.py：节点实现
        - agent/state.py：状态定义

这里的 SimpleAgent 展示的是核心原理
LangGraph 只是实现方式之一
"""


# ═══════════════════════════════════════════════════════════════════════════════
# 第八部分：演示代码
# ═══════════════════════════════════════════════════════════════════════════════

def demo():
    """演示简单 Agent 的使用"""
    print("=" * 60)
    print("Step 03: 构建简单 Agent - 演示")
    print("=" * 60)

    # 创建 Agent
    agent = SimpleAgent()

    print("\n📋 可用工具:")
    for tool in agent.get_available_tools():
        print(f"  - {tool['name']}: {tool['description']}")

    # 场景 1：生成报表
    print("\n\n🔄 场景 1: 生成报表")
    print("-" * 40)
    print("用户: 生成一个销售报表")

    response = agent.process("生成一个销售报表")
    print(f"Agent: {response}")

    # 场景 2：搜索模板
    print("\n\n🔄 场景 2: 搜索模板")
    print("-" * 40)
    print("用户: 搜索采购相关的模板")

    response = agent.process("搜索采购相关的模板")
    print(f"Agent: {response}")

    # 场景 3：计算
    print("\n\n🔄 场景 3: 数学计算")
    print("-" * 40)
    print("用户: 计算 100 + 200 * 3")

    response = agent.process("计算 100 + 200 * 3")
    print(f"Agent: {response}")

    # 场景 4：多轮对话
    print("\n\n🔄 场景 4: 多轮对话")
    print("-" * 40)
    print("用户: 我想生成一个报表")

    response1 = agent.process("我想生成一个报表")
    print(f"Agent: {response1}")

    print("用户: 显示月度汇总数据")

    response2 = agent.process("显示月度汇总数据")
    print(f"Agent: {response2}")

    # 展示对话历史
    print("\n\n📜 对话历史:")
    for msg in agent.get_history():
        print(f"  [{msg['role']}]: {msg['content'][:50]}...")

    # 展示工具调用记录
    print("\n\n🔧 工具调用记录:")
    for call in agent.state.get("tool_calls", []):
        if "error" in call:
            print(f"  - {call['tool']}: ❌ {call['error']}")
        else:
            result = str(call.get("result", ""))[:30]
            print(f"  - {call['tool']}: ✓ {result}...")

    print("\n" + "=" * 60)
    print("✅ 演示完成")
    print("=" * 60)


if __name__ == "__main__":
    demo()