"""Example: Multi-Agent System with Router and Chaining."""

import ast
import operator

from agenai.agent import (
    Agent,
    AgentConfig,
    SpecializedAgent,
    RouterAgent,
    ChainedAgent,
)
from agenai.agent.tools.base import Tool, ToolResult


# Safe calculator using AST parsing
class CalculatorTool(Tool):
    """Safe calculator that parses math expressions without eval."""

    # Supported operators
    _ops = {
        ast.Add: operator.add,
        ast.Sub: operator.sub,
        ast.Mult: operator.mul,
        ast.Div: operator.truediv,
        ast.Pow: operator.pow,
        ast.USub: operator.neg,
    }

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

    @property
    def description(self) -> str:
        return "Perform mathematical calculations. Input should be a math expression."

    @property
    def parameters(self) -> dict:
        return {
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": "Math expression to evaluate (e.g., '2 + 2 * 3')",
                }
            },
            "required": ["expression"],
        }

    def _safe_eval(self, node):
        """Safely evaluate an AST node."""
        if isinstance(node, ast.Constant):
            return node.value
        elif isinstance(node, ast.BinOp):
            op_type = type(node.op)
            if op_type not in self._ops:
                raise ValueError(f"Unsupported operator: {op_type}")
            return self._ops[op_type](
                self._safe_eval(node.left),
                self._safe_eval(node.right)
            )
        elif isinstance(node, ast.UnaryOp):
            op_type = type(node.op)
            if op_type not in self._ops:
                raise ValueError(f"Unsupported operator: {op_type}")
            return self._ops[op_type](self._safe_eval(node.operand))
        else:
            raise ValueError(f"Unsupported expression type: {type(node)}")

    def execute(self, expression: str) -> ToolResult:
        try:
            tree = ast.parse(expression, mode='eval')
            result = self._safe_eval(tree.body)
            return ToolResult.ok(str(result))
        except Exception as e:
            return ToolResult.fail(str(e))


class MathAgent(SpecializedAgent):
    """Agent for math problems."""

    name = "math"
    description = "Solves mathematical problems and calculations"

    def get_system_prompt(self) -> str:
        return """You are a math expert. Use the calculator tool for computations.
Show your work step by step."""

    def get_tools(self) -> list[Tool]:
        return [CalculatorTool()]


class WriterAgent(SpecializedAgent):
    """Agent for writing tasks."""

    name = "writer"
    description = "Writes and edits text content"

    def get_system_prompt(self) -> str:
        return """You are a professional writer.
Write clear, engaging content.
Focus on clarity and proper structure."""

    def get_tools(self) -> list[Tool]:
        return []  # Writer doesn't need tools


class SummarizerAgent(SpecializedAgent):
    """Agent for summarizing content."""

    name = "summarizer"
    description = "Summarizes long text into concise points"

    def get_system_prompt(self) -> str:
        return """You are an expert summarizer.
Create concise summaries that capture key points.
Use bullet points when appropriate."""

    def get_tools(self) -> list[Tool]:
        return []


def demo_router():
    """Demo: Router automatically selects the right agent."""
    print("=== Router Agent Demo ===\n")

    router = RouterAgent(
        agents={
            "math": MathAgent(),
            "writer": WriterAgent(),
        }
    )

    # This should route to math agent
    result = router.run("What is 15 * 23 + 7?")
    print(f"Math query result: {result.output[:100] if result.success else result.error}\n")

    # This should route to writer agent
    result = router.run("Write a haiku about programming")
    print(f"Writing query result: {result.output if result.success else result.error}\n")


def demo_chain():
    """Demo: Chain agents together for multi-step processing."""
    print("=== Chained Agents Demo ===\n")

    # First write, then summarize
    chain = ChainedAgent(
        agents=[
            WriterAgent(config=AgentConfig(max_steps=5)),
            SummarizerAgent(config=AgentConfig(max_steps=3)),
        ]
    )

    result = chain.run("Explain the benefits of modular software architecture")

    if result.success:
        print(f"Final output (written then summarized):\n{result.output}")
    else:
        print(f"Error: {result.error}")


if __name__ == "__main__":
    import os

    if not os.environ.get("ANTHROPIC_API_KEY"):
        print("Set ANTHROPIC_API_KEY environment variable")
        exit(1)

    demo_router()
    print("\n" + "=" * 50 + "\n")
    demo_chain()
