from dataclasses import dataclass
from abc import abstractmethod
from typing import Protocol, Optional, Union
from .configuration import AIConfig
from .message import Message
from .chat import Chat


class AIError(Exception):
    pass


@dataclass
class Tokens:
    prompt: int = 0
    completion: int = 0
    total: int = 0


@dataclass
class AIResponse:
    """
    The response to an AI request. Consists of one or more messages
    (each containing the question and a single answer) and the nr.
    of used tokens.
    """
    messages: list[Message]
    tokens: Optional[Tokens] = None


class AI(Protocol):
    """
    The base class for AI clients.
    """

    name: str
    config: AIConfig

    @abstractmethod
    def request(self,
                question: Message,
                context: Chat,
                num_answers: int = 1) -> AIResponse:
        """
        Make an AI request, asking the given question with the given
        context (i. e. chat history). The nr. of requested answers
        corresponds to the nr. of messages in the 'AIResponse'.
        """
        raise NotImplementedError

    @abstractmethod
    def models(self) -> list[str]:
        """
        Return all models supported by this AI.
        """
        raise NotImplementedError

    def tokens(self, data: Union[Message, Chat]) -> int:
        """
        Computes the nr. of AI language tokens for the given message
        or chat. Note that the computation may not be 100% accurate
        and is not implemented for all AIs.
        """
        raise NotImplementedError