from typing import Any, Callable, Generic, Literal, TypeVar

from verse.content.image import Image, ImageParam
from verse.core import DataModel
from verse.core.data_model import DataModelField


class ChatCompletionMessageContentPart(DataModel):
    type: Literal["text", "image"]
    """Content part type."""

    text: str | None = None
    """Text for text content type."""

    image: Image | ImageParam | dict | None = None
    """Image for image content type."""


class ToolCallFunction(DataModel):
    name: str
    """The name of the function to call."""

    arguments: str | None = None
    """The arguments to call the function in JSON format."""


class ToolCall(DataModel):
    id: str | None = None
    """The ID of the tool call."""

    function: ToolCallFunction
    """The function that the model called."""

    type: Literal["function"] = "function"
    """The type of the tool. Currently, only `function` is supported."""


class ChatCompletionMessage(DataModel):
    role: Literal["system", "user", "assistant", "tool"] | None = None
    """The role of the author."""

    content: str | list[ChatCompletionMessageContentPart] | None = None
    """The contents of the message."""

    name: str | None = None
    """An optional name for the participant."""

    tool_calls: list[ToolCall] | None = None
    """The tool calls generated by the model."""

    tool_call_id: str | None = None
    """Tool call that this message is responding to."""

    refusal: str | None = None
    """The refusal message generated by the model."""


class FunctionDefinition(DataModel):
    name: str
    """Name of the function."""

    description: str | None = None
    """Description of the function."""

    parameters: dict[str, Any] | None = None
    """Parameters of the function."""

    strict: bool | None = None
    """Whether to enable strict schema adherence
    when generating the function call.
    """


class Tool(DataModel):
    function: FunctionDefinition
    """Function definition."""

    type: Literal["function"] = "function"


class ToolChoiceFunction(DataModel):
    name: str
    """Name of the function."""


class ToolChoice(DataModel):
    function: ToolChoiceFunction

    type: Literal["function"] = "function"


class JSONSchema(DataModel):
    name: str
    """The name of the response format."""

    description: str
    """A description of what the response format is for."""

    schema_: dict[str, object] = DataModelField(alias="schema")
    """The schema for the response format, as a JSON Schema object."""

    strict: bool | None = None
    """Whether to enable strict schema adherence when generating the output."""


class ResponseFormat(DataModel):
    type: Literal["text", "json_object", "json_schema"]
    """Response format type."""

    json_schema: str


class StreamOptions(DataModel):
    include_usage: bool
    """A value that indicates if usage info should
    be sent as part of streaming.
    """


class TopLogprob(DataModel):
    token: str
    """The token."""

    bytes: list[int] | None = None
    """A list of integers representing the UTF-8
    bytes representation of the token.
    """

    logprob: float
    """The log probability of this token.
    """


class ChatCompletionTokenLogprob(DataModel):
    token: str
    """The token."""

    bytes: list[int] | None = None
    """A list of integers representing the UTF-8
    bytes representation of the token.
    """

    logprob: float
    """The log probability of this token.
    """

    top_logprobs: list[TopLogprob]
    """List of the most likely tokens and their log probability, at this token
    position.
    """


class ChoiceLogProbs(DataModel):
    content: list[ChatCompletionTokenLogprob] | None = None
    """A list of message content tokens with
    log probability information.
    """


class Choice(DataModel):
    finish_reason: str | None = None
    """The reason the model stopped generating tokens.
    Values are ["stop", "length", "content_filter", "tool_calls",
    "error", "cancel"]
    """

    index: int | None = None
    """The index of the choice in the list of choices."""

    message: ChatCompletionMessage | None = None
    """Chat completion response generated by the model."""

    delta: ChatCompletionMessage | None = None
    """Delta message when streaming."""

    logprobs: ChoiceLogProbs | None = None
    """Log probability information for the choice."""


class Usage(DataModel):
    completion_tokens: int
    """Number of tokens in the generated completion."""

    prompt_tokens: int
    """Number of tokens in the prompt."""

    total_tokens: int
    """Total number of tokens used in the request (prompt + completion)."""


class ChatCompletionResult(DataModel):
    id: str | None = None
    """A unique identifier for the chat completion."""

    choices: list[Choice]
    """A list of chat completion choices."""

    created: int | None = None
    """The Unix timestampof when the chat completion was created."""

    model: str | None = None
    """The model used for the chat completion."""

    usage: Usage | None = None
    """Usage statistics for the completion request."""


T = TypeVar("T")


class Stream(Generic[T]):
    stream: Any
    convert_chunk: Callable[[Any], T]

    def __init__(self, stream: Any, convert_chunk: Callable[[Any], T]):
        self.stream = stream
        self.convert_chunk = convert_chunk

    def __iter__(self):
        return self

    def __next__(self) -> T:
        try:
            chunk = self.stream.__next__()
            return self.convert_chunk(chunk)
        except StopAsyncIteration:
            raise

    def __aiter__(self):
        return self

    async def __anext__(self) -> T:
        try:
            chunk = await self.stream.__anext__()
            return self.convert_chunk(chunk)
        except StopAsyncIteration:
            raise
