from pathlib import Path
from typing import Annotated, Any, Literal
from pydantic import BaseModel, Field, JsonValue
from uuid import uuid4
import base64


type DataContent = bytes | str

type FinishReason = Literal[
    "stop",
    "length",
    "content-filter",
    "tool-calls",
    "interrupted",
    "error",
    "other",
    "unknown",
]


type ToolChoice = Literal["auto", "none", "required"] | str


class Usage(BaseModel):
    input_tokens: int | None = None
    output_tokens: int | None = None
    total_tokens: int | None = None
    reasoning_tokens: int | None = None
    cached_input_tokens: int | None = None

    def __add__(self, y: "Usage") -> "Usage":
        x = Usage(**self.model_dump())
        if y.input_tokens:
            x.input_tokens = (x.input_tokens or 0) + y.input_tokens
        if y.output_tokens:
            x.output_tokens = (x.output_tokens or 0) + y.output_tokens
        if y.total_tokens:
            x.total_tokens = (x.total_tokens or 0) + y.total_tokens
        if y.reasoning_tokens:
            x.reasoning_tokens = (x.reasoning_tokens or 0) + y.reasoning_tokens
        if y.cached_input_tokens:
            x.cached_input_tokens = (x.cached_input_tokens or 0) + y.cached_input_tokens
        return x


class FunctionTool(BaseModel):
    type: Literal["function"] = "function"
    name: str
    input_schema: JsonValue
    description: str | None = None

    def to_openai_schema(self) -> Any:
        return {
            "type": self.type,
            "name": self.name,
            "description": self.description,
            "parameters": self.input_schema,
            "strict": True,
        }


class ProviderDefinedTool(BaseModel):
    type: Literal["provider-defined"] = "provider-defined"
    id: str
    name: str
    args: dict[str, JsonValue] = Field(default_factory=dict)

    def to_openai_schema(self) -> Any:
        return {"type": self.id, **self.args}


type ToolSchema = Annotated[
    FunctionTool | ProviderDefinedTool, Field(discriminator="type")
]


class ToolCall(BaseModel):
    type: Literal["tool-call"] = "tool-call"

    tool_call_id: str
    """The identifier of the tool call. It must be unique across all tool calls."""

    tool_name: str
    """The name of the tool that should be called."""

    input: dict[str, JsonValue]
    """JSON object with the tool call arguments. Must match the parameters schema of the tool."""

    provider_executed: bool | None = None
    """
    Whether the tool call will be executed by the provider.
    If this flag is not set or is false, the tool call will be executed by the client.
    """


class File(BaseModel):
    media_type: str
    data: DataContent
    id: str = Field(default_factory=lambda: str(uuid4()))

    def to_url(self) -> str:
        if isinstance(self.data, bytes):
            return (
                f"data:{self.media_type};base64,{base64.b64encode(self.data).decode()}"
            )
        elif isinstance(self.data, str):
            if self.data.startswith(("http://", "https://", "data:")):
                return self.data
            else:
                path = Path(self.data)
                if not path.is_file():
                    raise ValueError(f"File path '{self.data}' does not exist")
                return f"data:{self.media_type};base64,{base64.b64encode(path.read_bytes()).decode()}"
        else:
            return str(self.data)

    def to_bytes(self) -> bytes:
        if isinstance(self.data, bytes):
            return self.data
        elif isinstance(self.data, str):
            if self.data.startswith("data:"):
                header, b64data = self.data.split(",", 1)
                return base64.b64decode(b64data)
            elif self.data.startswith(("http://", "https://")):
                raise ValueError("Cannot convert URL data to bytes")
            else:
                path = Path(self.data)
                if not path.is_file():
                    raise ValueError(f"File path '{self.data}' does not exist")
                return path.read_bytes()
        else:
            raise ValueError("Invalid data type for File")


class ToolCallResponse(BaseModel):
    type: Literal["tool-result"] = "tool-result"

    tool_call_id: str
    """The ID of the tool call that this result is associated with."""

    tool_name: str
    """Name of the tool that generated this result."""

    input: dict[str, JsonValue]

    output: JsonValue
    """Result of the tool call. This is a JSON-serializable object."""

    output_files: list[File] = Field(default_factory=list)
    """Optional list of files generated by the tool call."""

    is_error: bool | None = None
    """Optional flag if the result is an error or an error message."""

    provider_executed: bool | None = None
    """
    Whether the tool result was generated by the provider.
    If this flag is set to true, the tool result was generated by the provider.
    If this flag is not set or is false, the tool result was generated by the client.
    """

    preliminary: bool | None = None
    """
    Whether the tool result is preliminary.

    Preliminary tool results replace each other, e.g. image previews.
    There always has to be a final, non-preliminary tool result.

    If this flag is set to true, the tool result is preliminary.
    If this flag is not set or is false, the tool result is not preliminary.
    """


class Annotation(BaseModel):
    type: Literal["annotation"] = "annotation"

    end: int
    """The index of the last character of the URL citation in the message."""

    start: int
    """The index of the first character of the URL citation in the message."""

    title: str
    """The title of the web resource."""

    url: str
    """The URL of the web resource."""


class UserConsentRequest(BaseModel):
    type: Literal["user-consent-request"] = "user-consent-request"
    id: str = Field(default_factory=lambda: str(uuid4()))
    message: str
    details: str | None = None

    def __init__(self, message: str, details: str | None = None):
        super().__init__(message=message, details=details)


__all__ = [
    "DataContent",
    "FinishReason",
    "ToolChoice",
    "Usage",
    "FunctionTool",
    "ProviderDefinedTool",
    "ToolSchema",
    "ToolCall",
    "ToolCallResponse",
    "File",
    "Annotation",
    "UserConsentRequest",
]
