from dataclasses import dataclass
from typing import Protocol, Optional, Union
from .configuration import AIConfig
from .tags import Tag
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.
    """

    ID: str
    name: str
    config: AIConfig

    def request(self,
                question: Message,
                chat: Chat,
                num_answers: int = 1,
                otags: Optional[set[Tag]] = None) -> AIResponse:
        """
        Make an AI request. Parameters:
        * question: the question to ask
        * chat: the chat history to be added as context
        * num_answers: nr. of requested answers (corresponds
          to the nr. of messages in the 'AIResponse')
        * otags: the output tags, i. e. the tags that all
          returned messages should contain
        """
        raise NotImplementedError

    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

    def print(self) -> None:
        """
        Print some info about the current AI, like system message.
        """
        pass