from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import AsyncGenerator

from flexai.message import AIMessage, Message, MessageContent, SystemMessage
from flexai.tool import Tool


@dataclass(frozen=True)
class Client(ABC):
    """Abstract base class for language model clients.

    Defines the interface for interacting with various language models.
    Subclasses should implement the necessary methods for specific LLM providers.
    """

    @abstractmethod
    async def get_chat_response(
        self,
        messages: list[Message],
        system: str | SystemMessage = "",
        tools: list[Tool] | None = None,
    ) -> AIMessage:
        """Retrieve a response from the chat model.

        Args:
            messages: Conversation history to send to the model.
            system: Optional system message to set the behavior of the AI.
            tools: Optional list of tools available for the model to use.

        Returns:
            A list of AI-generated messages in response to the input.
        """

    @abstractmethod
    async def stream_chat_response(
        self,
        messages: list[Message],
        system: str | SystemMessage = "",
        tools: list[Tool] | None = None,
        allow_tool: bool = True,
    ) -> AsyncGenerator[MessageContent | AIMessage, None]:  # type: ignore
        """Stream the response from the chat model in real-time.

        Args:
            messages: Conversation history to send to the model.
            system: Optional system message to set the behavior of the AI.
            tools: Optional list of tools available for the model to use.
            allow_tool: Boolean flag for whether to fully interpret tool calls/results in the history.

        Yields:
            AI-generated messages as they are generated by the model.
        """
