from typing import Literal, TypeVar

from railtracks.llm.content import Content

from .message import Message, Role


class MessageInfo:
    def __init__(
        self,
        *,
        input_tokens: int | None = None,
        output_tokens: int | None = None,
        latency: float | None = None,
        model_name: str | None = None,
        total_cost: float | None = None,
        system_fingerprint: str | None = None,
    ):
        """
        Creates a new instance of a message info object.

        Args:
            input_tokens: The number of input tokens used in the message.
            output_tokens: The number of output tokens generated by the model.
            latency: The latency of the model response in seconds.
            model_name: The name of the model used to generate the response.
            total_cost: The total cost of the request, if applicable.
            system_fingerprint: A unique identifier for the system that processed the request.
        """
        self.input_tokens = input_tokens
        self.output_tokens = output_tokens
        self.latency = latency
        self.model_name = model_name
        self.total_cost = total_cost
        self.system_fingerprint = system_fingerprint

    @property
    def total_tokens(self):
        """
        Gets the total number of tokens used in the message.

        If either input_tokens or output_tokens is None, it will return None.
        """
        if self.output_tokens is None or self.input_tokens is None:
            return None

        return self.output_tokens + self.input_tokens

    def __repr__(self):
        return (
            f"MessageInfo(input_tokens={self.input_tokens}, "
            f"output_tokens={self.output_tokens}, "
            f"latency={self.latency}, "
            f"model_name={self.model_name}, "
            f"total_cost={self.total_cost}, "
            f"system_fingerprint={self.system_fingerprint})"
        )


_T = TypeVar("_T", bound=Content)


class Response:
    """
    A simple object that represents a response from a model. It includes specific detail about the returned message
    and any other additional information from the model.
    """

    def __init__(
        self,
        message: Message[_T, Literal[Role.assistant]],
        message_info: MessageInfo = MessageInfo(),
    ):
        """
        Creates a new instance of a response object.

        Args:
            message: The message that was returned as part of this.
            streamer: A generator that streams the response as a collection of chunked strings.
            message_info: Additional information about the message, such as input/output tokens and latency.
        """
        if message is not None and not isinstance(message, Message):
            raise TypeError(f"message must be of type Message, got {type(message)}")
        self._message = message
        self._message_info = message_info

    @property
    def message(self):
        """
        Gets the message that was returned as part of this response.

        If none exists, this will return None.
        """
        return self._message

    @property
    def message_info(self) -> MessageInfo:
        """
        Gets the message info that was returned as part of this response.

        This object contains additional information about the message, such as input/output tokens and latency.
        """
        return self._message_info

    def __str__(self):
        if self._message is not None:
            return str(self._message)
        else:
            return "Response(<no-data>)"

    def __repr__(self):
        return f"Response(message={self._message}, message_info={self._message_info})"
