"""Agent provider implementations."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable

from psygnal import Signal
from pydantic_ai.result import FinalResult
from pydantic_graph import BaseNode, End
from typing_extensions import TypeVar

from llmling_agent.log import get_logger
from llmling_agent.tools import ToolCallInfo


if TYPE_CHECKING:
    from collections.abc import AsyncGenerator, AsyncIterator
    from contextlib import AbstractAsyncContextManager

    from pydantic_ai.agent import AgentRunResult
    import tokonomics
    from tokonomics.pydanticai_cost import Usage

    from llmling_agent.agent.context import AgentContext
    from llmling_agent.common_types import ModelProtocol, ModelType
    from llmling_agent.messaging.messages import ChatMessage, TokenCost
    from llmling_agent.models.content import Content
    from llmling_agent.tools.base import Tool


logger = get_logger(__name__)
TResult_co = TypeVar("TResult_co", default=str, covariant=True)
type TNode[TResult_co] = BaseNode | End[FinalResult[TResult_co]]


@dataclass
class ProviderResponse:
    """Raw response data from provider."""

    content: Any
    tool_calls: list[ToolCallInfo] = field(default_factory=list)
    model_name: str = ""
    cost_and_usage: TokenCost | None = None
    provider_extra: dict[str, Any] | None = None


class AgentRunProtocol(Protocol[TResult_co]):
    """Protocol representing an ongoing, iterable agent execution run.

    Matches the interface provided by pydantic_ai.models.AgentRun.
    Allows iterating through execution graph nodes and accessing results/usage.
    """

    def __aiter__(self) -> AsyncIterator[TNode]:
        """Allows `async for node in agent_run:`."""
        ...

    async def __anext__(self) -> TNode:
        """Gets the next node in the iteration."""
        ...

    # --- Manual Iteration Interface ---
    @property
    def next_node(self) -> TNode:
        """The next node that will be run in the agent graph."""
        ...

    async def next(self, node: BaseNode) -> TNode:
        """Manually advances the iteration by running the specified `node`.

        Returns the node *after* the processed one, or an End node.
        """
        ...

    # --- Results and Usage ---
    @property
    def result(self) -> AgentRunResult[TResult_co] | None:
        """The final result (AgentRunResult) of the agent run.

        Available after iteration completes (i.e., an End node is reached).
        Returns None if the run is not finished or failed.
        """
        ...

    def usage(self) -> Usage:
        """Retrieves usage statistics (Usage) for the run so far."""
        ...


@runtime_checkable
class StreamingResponseProtocol[TResult](Protocol):
    """Protocol for streaming responses.

    This matches PydanticAI's StreamedRunResult interface to make transition easier,
    but lives in our core package to remove the dependency.
    """

    model_name: str | None
    is_complete: bool
    formatted_content: TResult = None  # type: ignore

    def stream(self) -> AsyncGenerator[TResult, None]:
        """Stream individual chunks as they arrive."""
        ...

    def stream_text(self, delta: bool = False) -> AsyncGenerator[str, None]:
        """Stream individual chunks as they arrive."""
        ...

    def usage(self) -> Usage:
        """Get token usage statistics if available."""
        ...


@dataclass
class UsageLimits:
    """Limits on model usage."""

    request_limit: int | None = 50
    """The maximum number of requests allowed to the model."""

    input_tokens_limit: int | None = None
    """The maximum number of tokens allowed in requests to the model."""

    output_tokens_limit: int | None = None
    """The maximum number of tokens allowed in responses from the model."""

    total_tokens_limit: int | None = None
    """The maximum number of tokens allowed in requests and responses combined."""


class AgentProvider[TDeps]:
    """Base class for agent providers."""

    tool_used = Signal(ToolCallInfo)
    chunk_streamed = Signal(str, str)
    model_changed = Signal(object)  # Model | None
    NAME: str

    def __init__(
        self,
        *,
        name: str = "agent",
        context: AgentContext[TDeps] | None = None,
        debug: bool = False,
    ):
        super().__init__()
        self._name = name
        self._model: str | ModelProtocol | None = None
        self._context = context
        self._debug = debug

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}()"

    def set_model(self, model: ModelType):
        """Default no-op implementation for setting model."""

    def iterate_run(
        self,
        *prompts: str | Content,
        message_id: str,
        message_history: list[ChatMessage],
        result_type: type[Any] | None = None,
        model: ModelType = None,
        tools: list[Tool] | None = None,
        usage_limits: UsageLimits | None = None,
        **kwargs: Any,
    ) -> AbstractAsyncContextManager[AgentRunProtocol]:
        """Stream a response. Must be implemented by providers."""
        raise NotImplementedError

    @property
    def context(self) -> AgentContext[TDeps]:
        """Get context."""
        if self._context is None:
            msg = "Context not set"
            raise RuntimeError(msg)
        return self._context

    @context.setter
    def context(self, value: AgentContext[TDeps]):
        """Set context."""
        self._context = value

    @property
    def model_name(self) -> str | None:
        """Get model name."""
        return None

    @property
    def name(self) -> str:
        """Get agent name."""
        return self._name or "agent"

    @name.setter
    def name(self, value: str | None):
        """Set agent name."""
        self._name = value or "agent"

    async def generate_response(
        self,
        *prompts: str | Content,
        message_id: str,
        message_history: list[ChatMessage],
        result_type: type[Any] | None = None,
        model: ModelType = None,
        tools: list[Tool] | None = None,
        usage_limits: UsageLimits | None = None,
        **kwargs: Any,
    ) -> ProviderResponse:
        """Generate a response. Must be implemented by providers."""
        raise NotImplementedError

    def stream_response(
        self,
        *prompts: str | Content,
        message_id: str,
        message_history: list[ChatMessage],
        result_type: type[Any] | None = None,
        model: ModelType = None,
        tools: list[Tool] | None = None,
        usage_limits: UsageLimits | None = None,
        **kwargs: Any,
    ) -> AbstractAsyncContextManager[StreamingResponseProtocol]:
        """Stream a response. Must be implemented by providers."""
        raise NotImplementedError

    async def supports_feature(self, capability: Literal["vision"]) -> bool:
        """Check if provider supports a specific capability."""
        import tokonomics

        match capability:
            case "vision":
                if not self.model_name:
                    return False
                caps = await tokonomics.get_model_capabilities(self.model_name)
                return (caps and caps.supports_vision) or False
            case _:
                return False

    async def get_token_limits(self) -> tokonomics.TokenLimits | None:
        """Get token limits for the current model."""
        import tokonomics

        if not self.model_name:
            return None

        try:
            return await tokonomics.get_model_limits(self.model_name)
        except ValueError:
            logger.debug("Could not get token limits for model: %s", self.model_name)
            return None


class AgentLLMProvider[TDeps](AgentProvider[TDeps]):
    """Provider using LLM backend."""

    def __init__(
        self,
        model_settings: dict[str, Any] | None = None,
        model: str | ModelProtocol | None = None,
        **kwargs: Any,
    ):
        self.model_settings = model_settings or {}
        self._model = model
        super().__init__(**kwargs)

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(model={self.model_name})"

    def set_model(self, model: ModelType):
        """Set the model for this agent.

        Args:
            model: New model to use (name or instance)

        Emits:
            model_changed signal with the new model
        """
        old_name = self.model_name
        self._model = model
        self.model_changed.emit(model)
        if old_name:
            logger.debug("Changed model from %s to %s", old_name, self.model_name)
