"""Base classes and utilities for building specialized agents."""

from __future__ import annotations

from abc import abstractmethod
from typing import Any

from agenai.agent.config import AgentConfig
from agenai.agent.context import AgentContext
from agenai.agent.core import Agent, AgentResult, StepResult, StepType
from agenai.agent.hooks import HookRegistry
from agenai.agent.llm import LLMProvider
from agenai.agent.memory import ConversationMemory
from agenai.agent.retry import RetryConfig
from agenai.agent.tools import ToolRegistry


class SpecializedAgent(Agent):
    """Base class for creating specialized agents.

    Provides a cleaner interface for customization by exposing
    commonly-overridden methods with clearer semantics.

    Example:
        class ResearchAgent(SpecializedAgent):
            name = "research"
            description = "Researches topics using web search"

            def get_system_prompt(self) -> str:
                return '''You are a research assistant.
                Use web_search to find information.
                Always cite your sources.'''

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

    # Override these in subclasses
    name: str = "specialized"
    description: str = "A specialized agent"

    def __init__(
        self,
        config: AgentConfig | None = None,
        llm: LLMProvider | None = None,
        hooks: HookRegistry | None = None,
        retry_config: RetryConfig | None = None,
        **kwargs: Any,
    ):
        # Build tools from get_tools()
        tools = ToolRegistry()
        for tool in self.get_tools():
            tools.register(tool)

        # Build config with system prompt
        if config is None:
            config = AgentConfig(system_prompt=self.get_system_prompt())
        elif not config.system_prompt or config.system_prompt == "You are a helpful AI assistant.":
            config = config.model_copy(update={"system_prompt": self.get_system_prompt()})

        super().__init__(
            config=config,
            llm=llm,
            tools=tools,
            memory=ConversationMemory(),
            hooks=hooks,
            retry_config=retry_config,
        )

        # Store additional kwargs for subclass use
        self.init_kwargs = kwargs

        # Allow subclass initialization
        self.setup()

    def get_system_prompt(self) -> str:
        """Return the system prompt for this agent. Override in subclass."""
        return f"You are a {self.description}."

    def get_tools(self) -> list[Any]:
        """Return the tools available to this agent. Override in subclass."""
        return []

    def setup(self) -> None:
        """Called after initialization. Override for custom setup logic."""
        pass

    def on_start(self, context: AgentContext) -> None:
        """Called when the agent starts processing. Override for custom logic."""
        pass

    def on_complete(self, result: AgentResult) -> None:
        """Called when the agent completes. Override for custom logic."""
        pass

    def on_error(self, error: Exception) -> None:
        """Called when an error occurs. Override for custom logic."""
        pass

    def run(
        self,
        user_input: str,
        context: AgentContext | None = None,
    ) -> AgentResult:
        """Run the agent with lifecycle hooks."""
        # Create context if not provided (so on_start gets it)
        if context is None:
            import uuid
            context = AgentContext(
                run_id=str(uuid.uuid4()),
                user_input=user_input,
            )

        try:
            # Pre-run hook
            self.on_start(context)

            # Run agent
            result = super().run(user_input, context)

            # Post-run hook
            self.on_complete(result)

            return result

        except Exception as e:
            self.on_error(e)
            raise


class ChainedAgent:
    """Agent that chains multiple specialized agents together.

    Each agent in the chain processes the output of the previous one.
    Note: This is an orchestrator, not an Agent subclass.
    """

    def __init__(
        self,
        agents: list[Agent],
        config: AgentConfig | None = None,
    ):
        self.agents = agents
        self.config = config or AgentConfig()

    def run(
        self,
        user_input: str,
        context: AgentContext | None = None,
    ) -> AgentResult:
        """Run agents in sequence, passing output to next input."""
        current_input = user_input
        all_steps: list[StepResult] = []
        total_tokens = 0

        for i, agent in enumerate(self.agents):
            # Create child context
            agent_context = AgentContext(
                run_id=f"{context.run_id if context else 'chain'}-{i}",
                user_input=current_input,
                parent=context,
                metadata={"chain_position": i, "agent_name": getattr(agent, "name", f"agent_{i}")},
            )

            result = agent.run(current_input, agent_context)
            all_steps.extend(result.steps)
            total_tokens += result.total_tokens

            if not result.success:
                return AgentResult(
                    success=False,
                    output="",
                    steps=all_steps,
                    error=f"Agent {i} failed: {result.error}",
                    total_tokens=total_tokens,
                    context=context,
                )

            # Pass output to next agent
            current_input = result.output

        return AgentResult(
            success=True,
            output=current_input,
            steps=all_steps,
            total_tokens=total_tokens,
            context=context,
        )


class RouterAgent(Agent):
    """Agent that routes requests to specialized sub-agents.

    Uses the LLM to determine which agent should handle a request.
    """

    def __init__(
        self,
        agents: dict[str, Agent],
        config: AgentConfig | None = None,
        llm: LLMProvider | None = None,
    ):
        # Build routing prompt
        agent_descriptions = "\n".join(
            f"- {name}: {getattr(agent, 'description', 'No description')}"
            for name, agent in agents.items()
        )

        routing_prompt = f"""You are a routing agent. Your job is to determine which specialized agent should handle a user request.

Available agents:
{agent_descriptions}

Respond with ONLY the agent name, nothing else."""

        if config is None:
            config = AgentConfig(system_prompt=routing_prompt, max_steps=1)
        else:
            config = config.model_copy(update={"system_prompt": routing_prompt, "max_steps": 1})

        super().__init__(config=config, llm=llm)
        self.agents = agents

    def run(
        self,
        user_input: str,
        context: AgentContext | None = None,
    ) -> AgentResult:
        """Route to appropriate agent and run."""
        # First, determine which agent to use
        routing_result = super().run(f"Which agent should handle: {user_input}", context)

        if not routing_result.success:
            return routing_result

        # Parse agent name from response
        agent_name = routing_result.output.strip().lower()

        # Find matching agent
        selected_agent = None
        for name, agent in self.agents.items():
            if name.lower() == agent_name or name.lower() in agent_name:
                selected_agent = agent
                break

        if selected_agent is None:
            # Default to first agent or return error
            if self.agents:
                selected_agent = list(self.agents.values())[0]
            else:
                return AgentResult(
                    success=False,
                    output="",
                    error=f"No agent found for: {agent_name}",
                    context=context,
                )

        # Run selected agent
        self.reset()  # Clear router memory
        return selected_agent.run(user_input, context)
