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

from pydantic import BeforeValidator, field_validator

# We don't want to make opentelemetry a dependency of the wasm module
if typing.TYPE_CHECKING:
    from opentelemetry.util.types import AttributeValue


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

    def as_gen_ai_otel_attributes(self) -> dict[str, int]:
        """The attributes specified by the GenAI Otel Semantic convention.

        See <https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/#genai-attributes>
        for more details.
        """
        return {
            "gen_ai.usage.input_tokens": self.prompt,
            "gen_ai.usage.output_tokens": self.completion,
        }


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"

    def as_gen_ai_otel_attributes(self) -> dict[str, "AttributeValue"]:
        """How to format the finish reason as a GenAI attribute.

        The OTel spec specifies two possibilities: Either including it in the message
        (`gen_ai.output.messages.0.finish_reason`) or as a separate field on the
        response (`gen_ai.response.finish_reason`). Langfuse get's a bit confused by
        the first option, so we use the second.
        """
        return {"gen_ai.response.finish_reasons": [self.value]}


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

    def as_gen_ai_otel_attributes(self) -> dict[str, Any]:
        """Format as specified by OpenTelemetry GenAI semantic conventions.

        See <https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/#genai-attributes>
        for more details.
        """
        return {
            "role": self.role.value,
            # While `parts` is required by the OTel specification, Langfuse only renders
            # text content properly if passed as `content`.
            "content": self.content,
        }

    @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)
