"""
Wrap agents as tools for hierarchical agent composition.
"""

from collections.abc import AsyncGenerator, Callable
from typing import Any, Dict, List, Optional, Protocol, Union, runtime_checkable

from agentbyte.cancellation_token import CancellationToken
from agentbyte.context import AgentContext
from agentbyte.llm.types import ChatCompletionChunk
from agentbyte.messages import Message, UserMessage
from agentbyte.tools.base import BaseTool
from agentbyte.types import ToolResult

from .types import AgentEvent, AgentResponse

ResultStrategy = Union[str, Callable[[List[Message]], str]]


@runtime_checkable
class AgentToolCompatible(Protocol):
    name: str
    description: str

    async def run(
        self,
        task: Optional[Union[str, UserMessage, List[Message]]] = None,
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[CancellationToken] = None,
    ) -> AgentResponse: ...

    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
    ]: ...


class AgentAsTool(BaseTool):
    """Expose any BaseAgent as a BaseTool."""

    def __init__(
        self,
        agent: AgentToolCompatible,
        task_parameter_name: str = "task",
        result_strategy: ResultStrategy = "last",
    ) -> None:
        if not isinstance(agent, AgentToolCompatible):
            raise TypeError("agent must implement the agent tool interface")

        super().__init__(name=agent.name, description=agent.description)

        self.agent = agent
        self.task_parameter_name = task_parameter_name
        self.result_strategy = result_strategy
        self._validate_result_strategy()

    def _validate_result_strategy(self) -> None:
        if callable(self.result_strategy):
            return

        if not isinstance(self.result_strategy, str):
            raise TypeError("result_strategy must be a string or callable")

        if self.result_strategy in {"last", "all"}:
            return

        if self.result_strategy.startswith("last:"):
            try:
                count = int(self.result_strategy.split(":", maxsplit=1)[1])
            except ValueError as exc:
                raise ValueError(
                    "Invalid result_strategy format, expected 'last:N'",
                ) from exc

            if count <= 0:
                raise ValueError("N must be positive in 'last:N' strategy")
            return

        raise ValueError(
            "Unknown result_strategy. Use 'last', 'last:N', 'all', or callable.",
        )

    def _extract_result(self, messages: List[Message]) -> str:
        if not messages:
            return ""

        if callable(self.result_strategy):
            return self.result_strategy(messages)

        if self.result_strategy == "last":
            return messages[-1].content

        if self.result_strategy == "all":
            return "\n".join(message.content for message in messages)

        count = int(self.result_strategy.split(":", maxsplit=1)[1])
        selected = messages[-count:]
        return "\n".join(message.content for message in selected)

    @property
    def parameters(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                self.task_parameter_name: {
                    "type": "string",
                    "description": f"Task for {self.agent.name} to complete",
                },
            },
            "required": [self.task_parameter_name],
        }

    async def execute(
        self,
        parameters: Dict[str, Any],
        context: Optional[AgentContext] = None,
    ) -> ToolResult:
        task_str = parameters.get(self.task_parameter_name, "")
        task = UserMessage(content=task_str, source="user")

        try:
            response = await self.agent.run(task=task, context=context)
            final_content = self._extract_result(response.messages)
            return ToolResult(
                success=True,
                result=final_content,
                metadata={
                    "agent_name": self.agent.name,
                    "message_count": len(response.messages),
                    "usage": response.usage.model_dump() if response.usage else None,
                    "finish_reason": response.finish_reason,
                },
            )
        except Exception as exc:
            return ToolResult(
                success=False,
                result="",
                error=f"Agent execution failed: {exc}",
                metadata={"agent_name": self.agent.name},
            )

    async def execute_stream(
        self,
        parameters: Dict[str, Any],
        context: Optional[AgentContext] = None,
        cancellation_token: Optional[CancellationToken] = None,
    ) -> AsyncGenerator[Union[Message, AgentEvent, ToolResult], None]:
        task_str = parameters.get(self.task_parameter_name, "")
        task = UserMessage(content=task_str, source="user")
        final_response: Optional[AgentResponse] = None

        try:
            async for item in self.agent.run_stream(
                task=task,
                context=context,
                cancellation_token=cancellation_token,
                verbose=False,
                stream_tokens=False,
            ):
                if isinstance(item, AgentResponse):
                    final_response = item
                else:
                    yield item

            result_content = ""
            metadata: Dict[str, Any] = {"agent_name": self.agent.name}

            if final_response is not None:
                result_content = self._extract_result(final_response.messages)
                metadata["message_count"] = len(final_response.messages)
                metadata["usage"] = (
                    final_response.usage.model_dump() if final_response.usage else None
                )
                metadata["finish_reason"] = final_response.finish_reason

            yield ToolResult(success=True, result=result_content, metadata=metadata)
        except Exception as exc:
            yield ToolResult(
                success=False,
                result="",
                error=f"Agent execution failed: {exc}",
                metadata={"agent_name": self.agent.name},
            )


__all__ = ["AgentAsTool", "ResultStrategy", "AgentToolCompatible"]
