from collections.abc import Sequence
from typing import Any, Dict, List, Literal, TypedDict, Union

from typing_extensions import Required, TypeAlias

from openinference.semconv.trace import (
    OpenInferenceLLMProviderValues,
    OpenInferenceLLMSystemValues,
    OpenInferenceMimeTypeValues,
    OpenInferenceSpanKindValues,
)

OpenInferenceSpanKind = Union[
    Literal[
        "agent",
        "chain",
        "embedding",
        "evaluator",
        "guardrail",
        "llm",
        "prompt",
        "reranker",
        "retriever",
        "tool",
        "unknown",
    ],
    OpenInferenceSpanKindValues,
]
OpenInferenceMimeType = Union[
    Literal["application/json", "text/plain"],
    OpenInferenceMimeTypeValues,
]
OpenInferenceLLMProvider: TypeAlias = Union[str, OpenInferenceLLMProviderValues]
OpenInferenceLLMSystem: TypeAlias = Union[str, OpenInferenceLLMSystemValues]
AnnotationScope: TypeAlias = Literal["span", "trace", "session"]


class Annotation(TypedDict, total=False):
    """A single annotation or evaluation result."""

    name: Required[str]
    score: Union[int, float]
    label: str
    explanation: str
    annotator_kind: str
    identifier: str
    metadata: Union[str, Dict[str, Any]]


class Image(TypedDict, total=False):
    url: str


class TextMessageContent(TypedDict, total=False):
    type: Required[Literal["text"]]
    text: Required[str]
    signature: str


class ImageMessageContent(TypedDict):
    type: Literal["image"]
    image: Image


class ReasoningMessageContent(TypedDict, total=False):
    type: Required[Literal["reasoning"]]
    text: str
    signature: str
    data: str
    encrypted_content: str


MessageContent: TypeAlias = Union[TextMessageContent, ImageMessageContent, ReasoningMessageContent]


class ToolCallFunction(TypedDict, total=False):
    name: str
    arguments: Union[str, Dict[str, Any]]


class ToolCall(TypedDict, total=False):
    id: str
    function: ToolCallFunction
    reasoning_signature: str


class Message(TypedDict, total=False):
    role: str
    content: str
    contents: "Sequence[MessageContent]"
    tool_call_id: str
    tool_calls: "Sequence[ToolCall]"


class PromptDetails(TypedDict, total=False):
    audio: int
    cache_read: int
    cache_write: int


class TokenCount(TypedDict, total=False):
    prompt: int
    completion: int
    total: int
    prompt_details: PromptDetails


class Tool(TypedDict, total=False):
    json_schema: Required[Union[str, Dict[str, Any]]]


class Embedding(TypedDict, total=False):
    text: str
    vector: List[float]


class Document(TypedDict, total=False):
    content: str
    id: Union[str, int]
    metadata: Union[str, Dict[str, Any]]
    score: float
