from dataclasses import dataclass
from enum import Enum
from typing import Annotated, Any, Self

from pydantic import BeforeValidator, field_validator


def none_to_nan(v: float | None) -> float:
    """Custom deserialization for Nan.

    This custom deserialization is necessary because JSON does not support NaN (Not a Number)
    values as valid data types, as specified by the JSON standard.
    """

    return float("nan") if v is None else v


NanFloat = Annotated[float, BeforeValidator(none_to_nan)]


@dataclass
class Logprob:
    """Logarithmic probability of the token returned in the completion."""

    token: bytes
    logprob: NanFloat

    @field_validator("token", mode="before")
    @classmethod
    def convert_token_to_bytes(cls, v: Any) -> Any:
        if isinstance(v, list):
            return bytes(v)
        return v

    def try_as_utf8(self) -> str | None:
        """Try to decode the token as utf-8.

        A token may also represent just a part of an utf-8 character, in which
        case it does not have a valid utf-8 encoding on its own.
        """
        try:
            return self.token.decode("utf-8")
        except UnicodeDecodeError:
            return None


@dataclass
class Distribution:
    sampled: Logprob
    top: list[Logprob]

    @classmethod
    def from_dict(cls, body: dict[str, Any]) -> "Distribution":
        raise NotImplementedError(
            "Not implemented, maybe look into pydantic dataclasses"
        )


@dataclass
class MessageAppend:
    """A chunk of a message generated by the model.

    Attributes:
        content (str, required): A chunk of the message content.
        logprobs (list[Distribution], required): Corresponding log probabilities for each token in the message content.
    """

    content: str
    logprobs: list[Distribution]

    @classmethod
    def from_dict(cls, body: dict[str, Any]) -> Self:
        return cls(
            content=body["content"],
            logprobs=body["logprobs"],
        )


class Role(str, Enum):
    """A role used for a message in a chat."""

    User = "user"
    Assistant = "assistant"
    System = "system"
    Tool = "tool"


@dataclass
class MessageBegin:
    role: str


@dataclass
class TokenUsage:
    """Usage statistics for the completion request."""

    prompt: int
    completion: int


class FinishReason(str, Enum):
    """The reason the model finished generating.

    Attributes:
        STOP: The model hit a natural stopping point or a provided stop sequence.

        LENGTH: The maximum number of tokens specified in the request was reached.

        CONTENT_FILTER: Content was omitted due to a flag from content filters.
    """

    STOP = "stop"
    LENGTH = "length"
    CONTENT_FILTER = "content_filter"


ChatEvent = MessageBegin | MessageAppend | FinishReason | TokenUsage


@dataclass
class Message:
    """A single turn in a conversation.

    Parameters:
        role (Role, required): The role of the message.
        content (str, required): The content of the message.
    """

    role: Role
    content: str

    @classmethod
    def user(cls, content: str) -> Self:
        return cls(role=Role.User, content=content)

    @classmethod
    def assistant(cls, content: str) -> Self:
        return cls(role=Role.Assistant, content=content)

    @classmethod
    def system(cls, content: str) -> Self:
        return cls(role=Role.System, content=content)

    @classmethod
    def tool(cls, content: str) -> Self:
        return cls(role=Role.Tool, content=content)

    @classmethod
    def from_dict(cls, body: dict[str, Any]) -> Self:
        # the shell csi does not serialize the roles in lowercase
        role = Role(body["role"].lower())
        content = body["content"]
        return cls(role=role, content=content)
