"""Core agent implementation - the heart of the system."""

from __future__ import annotations

import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

from agenai.agent.config import AgentConfig
from agenai.agent.context import AgentContext
from agenai.agent.hooks import HookData, HookEvent, HookRegistry
from agenai.agent.llm import AnthropicProvider, LLMProvider, LLMResponse
from agenai.agent.memory import ConversationMemory
from agenai.agent.retry import RetryConfig, RetryError, with_retry
from agenai.agent.tools import ToolRegistry, ToolResult


class StepType(str, Enum):
    """Type of agent step."""

    THINKING = "thinking"
    TOOL_CALL = "tool_call"
    FINAL_ANSWER = "final_answer"
    ERROR = "error"


@dataclass
class StepResult:
    """Result of a single agent step."""

    step_type: StepType
    content: str
    tool_name: str | None = None
    tool_args: dict[str, Any] | None = None
    tool_result: ToolResult | None = None
    # NEW: Include any thinking/text that came with tool calls
    thinking: str | None = None


@dataclass
class AgentResult:
    """Final result from running the agent."""

    success: bool
    output: str
    steps: list[StepResult] = field(default_factory=list)
    error: str | None = None
    total_tokens: int = 0
    # NEW: Include context for inspection
    context: AgentContext | None = None


class Agent:
    """The main agent orchestrator.

    Manages the agent loop lifecycle following the perceive -> reason -> act pattern.

    This class is designed to be subclassed for specialized agents. Override
    these methods to customize behavior:
    - _build_system_prompt(): Customize the system prompt
    - _pre_step(): Run logic before each step
    - _post_step(): Run logic after each step
    - _handle_tool_calls(): Customize tool execution
    - _should_continue(): Control loop termination
    """

    def __init__(
        self,
        config: AgentConfig | None = None,
        llm: LLMProvider | None = None,
        tools: ToolRegistry | None = None,
        memory: ConversationMemory | None = None,
        hooks: HookRegistry | None = None,
        retry_config: RetryConfig | None = None,
    ):
        """Initialize the agent.

        Args:
            config: Agent configuration. Uses defaults if not provided.
            llm: LLM provider. Creates AnthropicProvider if not provided.
            tools: Tool registry. Creates empty registry if not provided.
            memory: Conversation memory. Creates new memory if not provided.
            hooks: Hook registry for observability. Creates empty registry if not provided.
            retry_config: Configuration for API retry logic.
        """
        self.config = config or AgentConfig()
        self.tools = tools or ToolRegistry()
        self.memory = memory or ConversationMemory()
        self.hooks = hooks or HookRegistry()
        self.retry_config = retry_config or RetryConfig(
            max_retries=3,
            retry_on=(Exception,),  # Retry on API errors
            never_retry_on=(KeyboardInterrupt, SystemExit),
        )

        # Initialize LLM provider
        if llm:
            self.llm = llm
        else:
            self.llm = AnthropicProvider(model=self.config.model)

        # Current context (set during run)
        self._context: AgentContext | None = None

    @property
    def context(self) -> AgentContext | None:
        """Get the current execution context."""
        return self._context

    def run(
        self,
        user_input: str,
        context: AgentContext | None = None,
    ) -> AgentResult:
        """Run the agent with user input.

        This is the main entry point. It executes the agent loop until
        a final answer is produced or max_steps is reached.

        Args:
            user_input: The user's message/request.
            context: Optional context to use. Creates new if not provided.

        Returns:
            AgentResult with the final output and step history.
        """
        # Initialize context
        self._context = context or AgentContext(
            run_id=str(uuid.uuid4()),
            user_input=user_input,
            metadata={
                "model": self.config.model,
                "max_steps": self.config.max_steps,
            },
        )

        # Trigger agent start hook
        self._trigger_hook(HookEvent.AGENT_START)

        # Add user input to memory
        self.memory.add_user(user_input)

        steps: list[StepResult] = []
        total_tokens = 0

        try:
            for step_num in range(self.config.max_steps):
                self._context.step = step_num + 1

                # Pre-step hook (can be overridden)
                self._pre_step(step_num)
                self._trigger_hook(HookEvent.STEP_START, step_number=step_num + 1)

                # Execute step
                step_result, response = self._step()
                steps.append(step_result)

                if response:
                    total_tokens += response.usage.get("input_tokens", 0)
                    total_tokens += response.usage.get("output_tokens", 0)

                # Post-step hook (can be overridden)
                self._post_step(step_num, step_result)
                self._trigger_hook(HookEvent.STEP_END, step_number=step_num + 1)

                # Check if we've reached a final answer
                if step_result.step_type == StepType.FINAL_ANSWER:
                    result = AgentResult(
                        success=True,
                        output=step_result.content,
                        steps=steps,
                        total_tokens=total_tokens,
                        context=self._context,
                    )
                    self._trigger_hook(HookEvent.AGENT_END)
                    return result

                # Check for errors
                if step_result.step_type == StepType.ERROR:
                    result = AgentResult(
                        success=False,
                        output="",
                        steps=steps,
                        error=step_result.content,
                        total_tokens=total_tokens,
                        context=self._context,
                    )
                    self._trigger_hook(HookEvent.AGENT_ERROR, error=Exception(step_result.content))
                    return result

                # Check if we should continue (can be overridden)
                if not self._should_continue(step_num, step_result):
                    break

            # Max steps reached
            result = AgentResult(
                success=False,
                output="",
                steps=steps,
                error=f"Max steps ({self.config.max_steps}) reached without final answer",
                total_tokens=total_tokens,
                context=self._context,
            )
            self._trigger_hook(HookEvent.AGENT_END)
            return result

        except Exception as e:
            self._trigger_hook(HookEvent.AGENT_ERROR, error=e)
            return AgentResult(
                success=False,
                output="",
                steps=steps,
                error=str(e),
                total_tokens=total_tokens,
                context=self._context,
            )

    def _step(self) -> tuple[StepResult, LLMResponse | None]:
        """Execute a single step of the agent loop.

        Returns:
            Tuple of (StepResult, LLMResponse or None on error)
        """
        try:
            # Build messages for LLM
            messages = self.memory.get_anthropic_messages()

            # Get tool schemas
            tools = self.tools.list() if len(self.tools) > 0 else None

            # Get system prompt (can be overridden)
            system_prompt = self._build_system_prompt()

            # Trigger LLM start hook
            self._trigger_hook(HookEvent.LLM_START, messages=messages)

            # Call LLM with retry
            def call_llm() -> LLMResponse:
                return self.llm.chat(
                    messages=messages,
                    system=system_prompt,
                    tools=tools,
                    max_tokens=self.config.max_tokens,
                    temperature=self.config.temperature,
                )

            try:
                response = with_retry(call_llm, self.retry_config)
            except RetryError as e:
                self._trigger_hook(HookEvent.LLM_ERROR, error=e.last_error)
                raise e.last_error

            # Trigger LLM end hook
            self._trigger_hook(HookEvent.LLM_END, response=response)

            # Handle response
            if response.is_final:
                # Final text answer
                text = response.text or ""
                self.memory.add_assistant(text)
                return StepResult(step_type=StepType.FINAL_ANSWER, content=text), response

            # Process tool calls (may include thinking text)
            return self._handle_tool_calls(response), response

        except Exception as e:
            return StepResult(step_type=StepType.ERROR, content=str(e)), None

    def _handle_tool_calls(self, response: LLMResponse) -> StepResult:
        """Handle tool calls from the LLM response.

        This method can be overridden to customize tool execution behavior.
        """
        # FIX: Capture any text that came with the tool calls (thinking)
        thinking_text = response.text

        # Store assistant message with tool calls
        tool_call_blocks = []

        # Include text block if present
        if thinking_text:
            tool_call_blocks.append({"type": "text", "text": thinking_text})

        # Add tool use blocks
        for tc in response.tool_calls:
            tool_call_blocks.append({
                "type": "tool_use",
                "id": tc.id,
                "name": tc.name,
                "input": tc.arguments,
            })

        self.memory.add_assistant_tool_calls(tool_call_blocks)

        # Execute each tool and collect results
        results: list[str] = []
        last_tool_name = None
        last_tool_args = None
        last_tool_result = None

        for tool_call in response.tool_calls:
            last_tool_name = tool_call.name
            last_tool_args = tool_call.arguments

            # Trigger tool start hook
            self._trigger_hook(
                HookEvent.TOOL_START,
                tool_name=tool_call.name,
                tool_args=tool_call.arguments,
            )

            if tool_call.name not in self.tools:
                result = ToolResult.fail(f"Unknown tool: {tool_call.name}")
            else:
                try:
                    # Pass context to tool if it accepts it
                    result = self._execute_tool(tool_call.name, tool_call.arguments)
                except Exception as e:
                    result = ToolResult.fail(f"Tool error: {e}")
                    self._trigger_hook(
                        HookEvent.TOOL_ERROR,
                        tool_name=tool_call.name,
                        error=e,
                    )

            last_tool_result = result

            # Trigger tool end hook
            self._trigger_hook(
                HookEvent.TOOL_END,
                tool_name=tool_call.name,
                tool_args=tool_call.arguments,
                tool_result=result,
            )

            # Record in context
            if self._context:
                self._context.record_tool_call(
                    name=tool_call.name,
                    args=tool_call.arguments,
                    result=result.output if result.success else (result.error or ""),
                    success=result.success,
                )

            # Add tool result to memory
            result_text = result.output if result.success else f"Error: {result.error}"
            self.memory.add_tool_result(tool_call.id, tool_call.name, result_text)
            results.append(f"{tool_call.name}: {result_text}")

        return StepResult(
            step_type=StepType.TOOL_CALL,
            content="\n".join(results),
            tool_name=last_tool_name,
            tool_args=last_tool_args,
            tool_result=last_tool_result,
            thinking=thinking_text,
        )

    def _execute_tool(self, name: str, args: dict[str, Any]) -> ToolResult:
        """Execute a tool. Can be overridden to add custom behavior."""
        return self.tools.execute(name, **args)

    def _build_system_prompt(self) -> str:
        """Build the system prompt. Override to customize."""
        return self.config.system_prompt

    def _pre_step(self, step_num: int) -> None:
        """Called before each step. Override to add custom logic."""
        pass

    def _post_step(self, step_num: int, result: StepResult) -> None:
        """Called after each step. Override to add custom logic."""
        pass

    def _should_continue(self, step_num: int, result: StepResult) -> bool:
        """Determine if the agent loop should continue. Override to customize."""
        return True

    def _trigger_hook(
        self,
        event: HookEvent,
        **kwargs: Any,
    ) -> None:
        """Trigger a hook event."""
        data = HookData(
            event=event,
            context=self._context,
            **kwargs,
        )
        self.hooks.trigger(data)

    def reset(self) -> None:
        """Reset the agent state for a new conversation."""
        self.memory.clear()
        self._context = None

    def register_tool(self, tool: Any) -> None:
        """Register a tool with the agent."""
        self.tools.register(tool)

    def register_hook(self, hook: Any) -> None:
        """Register a hook with the agent."""
        self.hooks.register(hook)
