"""
Base agent abstraction and shared helper behavior.
"""

import json
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator, Callable
from typing import Any, Dict, List, Optional, Type, Union

from pydantic import BaseModel

from agentbyte.cancellation_token import CancellationToken
from agentbyte.context import AgentContext
from agentbyte.llm.base import BaseChatCompletionClient
from agentbyte.llm.types import ChatCompletionChunk
from agentbyte.memory import BaseMemory
from agentbyte.middleware import BaseMiddleware, MiddlewareChain
from agentbyte.messages import Message, SystemMessage, UserMessage
from agentbyte.tools.base import BaseTool, FunctionTool

from .types import AgentEvent, AgentResponse

MEMORY_ACCESS_OPERATION = "memory_access"


class AgentError(Exception):
    """Base exception for agent failures."""


class AgentConfigurationError(AgentError):
    """Raised when agent configuration is invalid."""


class AgentExecutionError(AgentError):
    """Raised when agent execution fails."""


class AgentToolError(AgentExecutionError):
    """Raised when tool execution path fails."""


class BaseAgent(ABC):
    """Abstract base class for all Agentbyte agents."""

    def __init__(
        self,
        name: str,
        description: str,
        instructions: str,
        model_client: BaseChatCompletionClient,
        tools: Optional[List[Union[BaseTool, Callable]]] = None,
        memory: Optional[BaseMemory] = None,
        middlewares: Optional[List[BaseMiddleware]] = None,
        max_iterations: int = 10,
        output_format: Optional[Type[BaseModel]] = None,
        summarize_tool_result: bool = True,
        required_tools: Optional[List[str]] = None,
        example_tasks: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        self.name = name
        self.description = description
        self.instructions = instructions
        self.model_client = model_client
        self.tools: List[BaseTool] = self._process_tools(tools or [])
        self.memory = memory
        self.middleware_chain = MiddlewareChain(middlewares)
        self.max_iterations = max_iterations
        self.output_format = output_format
        self.summarize_tool_result = summarize_tool_result
        self.required_tools = required_tools or []
        self.example_tasks = example_tasks or []
        self.extra_config = kwargs

        self._validate_configuration()

    def _validate_configuration(self) -> None:
        """Validate required configuration fields."""
        if not self.name or not isinstance(self.name, str):
            raise AgentConfigurationError("Agent name must be a non-empty string")
        if not self.description:
            raise AgentConfigurationError("Agent description cannot be empty")
        if not self.instructions:
            raise AgentConfigurationError("Agent instructions cannot be empty")
        if self.model_client is None:
            raise AgentConfigurationError("Model client is required")
        if self.max_iterations <= 0:
            raise AgentConfigurationError("max_iterations must be > 0")

    def _process_tools(self, tools: List[Union[BaseTool, Callable]]) -> List[BaseTool]:
        """Normalize mixed tool list into BaseTool instances."""
        processed: List[BaseTool] = []

        for tool in tools:
            if isinstance(tool, BaseTool):
                processed.append(tool)
            elif callable(tool):
                processed.append(FunctionTool(tool))
            else:
                raise AgentConfigurationError(
                    f"Invalid tool type: {type(tool)}. Must be BaseTool or callable.",
                )

        return processed

    def _find_tool(self, name: str) -> Optional[BaseTool]:
        """Find a tool by normalized name (handles versioned names)."""
        requested_name = name.split("_v")[0]
        for tool in self.tools:
            if tool.name == requested_name:
                return tool
        return None

    def _get_tools_for_llm(self) -> List[Dict[str, Any]]:
        """Convert tools into LLM function-calling schema format."""
        return [tool.to_llm_format() for tool in self.tools]

    def _convert_task_to_messages(
        self,
        task: Union[str, UserMessage, List[Message]],
    ) -> List[Message]:
        """Normalize task input into message list."""
        if isinstance(task, str):
            return [UserMessage(content=task, source="user")]
        if isinstance(task, UserMessage):
            return [task]
        if isinstance(task, list):
            return task
        raise AgentExecutionError(
            f"Unsupported task type: {type(task)}. "
            "Agent.run() only accepts str, UserMessage, or List[Message]."
        )

    async def _prepare_llm_messages(
        self,
        task_messages: List[Message],
        ctx: AgentContext,
    ) -> List[Message]:
        """Build LLM input from instructions, memory context, history, and task."""
        system_content = self.instructions

        if self.required_tools:
            tool_list = ", ".join(self.required_tools)
            system_content += (
                "\n\nIMPORTANT: You MUST use these tools in your response: "
                f"{tool_list}."
            )

        if self.memory is not None:
            try:

                async def _memory_access(payload: Dict[str, Any]) -> Any:
                    max_items = payload.get("max_items", 5)
                    if not isinstance(max_items, int):
                        max_items = 5
                    return await self.memory.get_context(max_items=max_items)

                memory_result = await self.middleware_chain.execute(
                    operation=MEMORY_ACCESS_OPERATION,
                    agent_name=self.name,
                    agent_context=ctx,
                    data={"max_items": 5},
                    metadata={"has_memory": True},
                    func=_memory_access,
                )
                if hasattr(memory_result, "results"):
                    context_lines = []
                    for item in memory_result.results:
                        content = item.content
                        if isinstance(content, str):
                            context_lines.append(content)
                        else:
                            context_lines.append(json.dumps(content))
                elif isinstance(memory_result, list):
                    context_lines = [str(item) for item in memory_result]
                else:
                    context_lines = []

                if context_lines:
                    system_content += "\n\nRelevant context from memory:\n"
                    system_content += "\n".join(context_lines)
            except Exception:
                pass

        return [
            SystemMessage(content=system_content, source="system"),
            *ctx.messages,
            *task_messages,
        ]

    def get_info(self) -> Dict[str, Any]:
        """Return agent metadata for orchestration/discovery."""
        return {
            "name": self.name,
            "description": self.description,
            "type": self.__class__.__name__,
            "tools_count": len(self.tools),
            "middlewares_count": len(self.middleware_chain.middlewares),
            "max_iterations": self.max_iterations,
        }

    def as_tool(
        self,
        task_parameter_name: str = "task",
        result_strategy: Union[str, Callable[[List[Message]], str]] = "last",
    ) -> BaseTool:
        """Wrap this agent as a tool for hierarchical composition."""
        from .agent_as_tool import AgentAsTool

        return AgentAsTool(
            agent=self,
            task_parameter_name=task_parameter_name,
            result_strategy=result_strategy,
        )

    @abstractmethod
    async def run(
        self,
        task: Optional[Union[str, UserMessage, List[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[CancellationToken] = None,
    ) -> AgentResponse:
        """Execute task and return only final response.

        Args:
            task: The task to execute. May be omitted when resuming from a
                non-empty ``context`` (e.g. approval resume).
            context: Per-call AgentContext. When provided, the agent threads
                it as the working context for this run and returns it mutated
                via AgentResponse.context. When ``None`` a fresh AgentContext
                is created so the agent remains stateless.
            cancellation_token: Optional token to cancel execution.
        """

    @abstractmethod
    def run_stream(
        self,
        task: Optional[Union[str, UserMessage, List[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[CancellationToken] = None,
        verbose: bool = False,
        stream_tokens: bool = False,
    ) -> AsyncGenerator[
        Union[Message, AgentEvent, AgentResponse, ChatCompletionChunk], None
    ]:
        """Stream agent execution events, ending with a final AgentResponse.

        Args:
            task: The task to execute. May be omitted when resuming from a
                non-empty ``context``.
            context: Per-call AgentContext threaded through execution. ``None``
                creates a fresh context, keeping the agent fully stateless.
            cancellation_token: Optional token to cancel execution.
            verbose: Emit detailed agent events.
            stream_tokens: Stream individual LLM token chunks.
        """
