"""The `Agent` class for LLM agents."""

from abc import ABC
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Generic

from ..context import Context, DepsT
from ..formatting import FormattableT
from ..messages import UserContent
from ..models import Model
from ..responses import AsyncResponse, AsyncStreamResponse, Response, StreamResponse
from ..tools import AsyncContextTool, AsyncTool, ContextTool, Tool


@dataclass
class BaseAgent(Generic[DepsT, FormattableT], ABC):
    """Agent class for generating responses using LLMs with tools."""

    ctx: Context[DepsT]
    """The context for the agent, such as the history of messages."""

    format: type[FormattableT] | None
    """The response format for the generated response."""

    model: Model
    """The default model the agent will use if not specified through context."""


@dataclass
class Agent(BaseAgent[DepsT, FormattableT]):
    """Agent class for generating responses using LLMs with tools."""

    tools: Sequence[Tool | ContextTool[DepsT]] | None
    """The tools available to the agent, if any."""

    def __call__(
        self,
        content: UserContent,
        *,
        ctx: Context[DepsT] | None = None,
    ) -> Response[FormattableT]:
        """Generates a response by running the agent loop."""
        raise NotImplementedError()

    def call(
        self,
        content: UserContent,
        *,
        ctx: Context[DepsT] | None = None,
    ) -> Response[FormattableT]:
        """Generates a response by running the agent loop."""
        raise NotImplementedError()

    def stream(
        self,
        content: UserContent,
        *,
        ctx: Context[DepsT] | None = None,
    ) -> StreamResponse[FormattableT]:
        """Streams the response generated by running the agent loop."""
        raise NotImplementedError()


@dataclass
class AsyncAgent(BaseAgent[DepsT, FormattableT]):
    """Asynchronous agent class for generating responses using LLMs with tools."""

    tools: Sequence[AsyncTool | AsyncContextTool[DepsT]] | None
    """The tools available to the agent, if any."""

    async def __call__(
        self,
        content: UserContent,
        *,
        ctx: Context[DepsT] | None = None,
    ) -> AsyncResponse[FormattableT]:
        """Generates a response by running the agent loop asynchronously."""
        raise NotImplementedError()

    async def call(
        self,
        content: UserContent,
        *,
        ctx: Context[DepsT] | None = None,
    ) -> AsyncResponse[FormattableT]:
        """Generates a response by running the agent loop asynchronously."""
        raise NotImplementedError()

    async def stream(
        self,
        content: UserContent,
        *,
        ctx: Context[DepsT] | None = None,
    ) -> AsyncStreamResponse[FormattableT]:
        """Streams the response generated by running the agent loop asynchronously."""
        raise NotImplementedError()
