"""Pydantic models for convert.py - subset that avoids discriminator issues.

This module contains only the models needed for the LangChain convert.py module,
extracted from models.py to avoid the problematic discriminated unions that
cause runtime import errors.
"""

from __future__ import annotations

from enum import Enum
from typing import Annotated, Any, Literal

from pydantic import BaseModel, Field


class MessageRole(Enum):
    user = "user"
    assistant = "assistant"
    system = "system"
    developer = "developer"


class MessageStatus(Enum):
    in_progress = "in_progress"
    completed = "completed"
    incomplete = "incomplete"


class FunctionCallStatus(Enum):
    """Status for function calls and function call outputs."""
    in_progress = "in_progress"
    completed = "completed"
    incomplete = "incomplete"


# Alias for backwards compatibility - both function calls and outputs use the same status values
FunctionCallOutputStatusEnum = FunctionCallStatus


class InputTextContent(BaseModel):
    type: Annotated[
        Literal["input_text"],
        Field(description="The type of the input item. Always `input_text`."),
    ]
    text: Annotated[str, Field(description="The text input to the model.")]


class UrlCitationBody(BaseModel):
    type: Annotated[
        Literal["url_citation"],
        Field(description="The type of the URL citation. Always `url_citation`."),
    ]
    url: Annotated[str, Field(description="The URL of the web resource.")]
    start_index: Annotated[
        int,
        Field(
            description="The index of the first character of the URL citation in the message."
        ),
    ]
    end_index: Annotated[
        int,
        Field(
            description="The index of the last character of the URL citation in the message."
        ),
    ]
    title: Annotated[str, Field(description="The title of the web resource.")]


class TopLogProb(BaseModel):
    token: str
    logprob: float
    bytes: list[int]


class LogProb(BaseModel):
    token: str
    logprob: float
    bytes: list[int]
    top_logprobs: list[TopLogProb]


class OutputTextContent(BaseModel):
    type: Annotated[
        Literal["output_text"],
        Field(description="The type of the output text. Always `output_text`."),
    ]
    text: Annotated[str, Field(description="The text output from the model.")]
    annotations: Annotated[
        list[UrlCitationBody], Field(description="The annotations of the text output.")
    ]
    logprobs: list[LogProb]


class FunctionCall(BaseModel):
    type: Annotated[
        Literal["function_call"],
        Field(description="The type of the item. Always `function_call`."),
    ]
    id: Annotated[str, Field(description="The unique ID of the function call item.")]
    call_id: Annotated[
        str,
        Field(
            description="The unique ID of the function tool call that was generated."
        ),
    ]
    name: Annotated[str, Field(description="The name of the function that was called.")]
    arguments: Annotated[
        str, Field(description="The arguments JSON string that was generated.")
    ]
    status: FunctionCallStatus


class FunctionCallOutput(BaseModel):
    type: Annotated[
        Literal["function_call_output"],
        Field(
            description="The type of the function tool call output. Always `function_call_output`."
        ),
    ]
    id: Annotated[
        str,
        Field(
            description="The unique ID of the function tool call output. Populated when this item is returned via API."
        ),
    ]
    call_id: Annotated[
        str,
        Field(
            description="The unique ID of the function tool call generated by the model."
        ),
    ]
    output: str
    status: FunctionCallOutputStatusEnum


class Message(BaseModel):
    type: Annotated[
        Literal["message"],
        Field(description="The type of the message. Always set to `message`."),
    ]
    id: Annotated[str, Field(description="The unique ID of the message.")]
    status: MessageStatus
    role: MessageRole
    content: Annotated[
        list[InputTextContent | OutputTextContent],
        Field(description="The content of the message"),
    ]


class FunctionTool(BaseModel):
    type: Annotated[
        Literal["function"],
        Field(description="The type of the function tool. Always `function`."),
    ]
    name: Annotated[str, Field(description="The name of the function to call.")]
    description: str | None
    parameters: dict[str, Any] | None
    strict: bool | None


class InputTokensDetails(BaseModel):
    cached_tokens: Annotated[
        int,
        Field(description="The number of input tokens that were served from cache."),
    ]


class OutputTokensDetails(BaseModel):
    reasoning_tokens: Annotated[
        int,
        Field(
            description="The number of output tokens that were attributed to reasoning."
        ),
    ]


class Usage(BaseModel):
    input_tokens: Annotated[
        int,
        Field(
            description="The number of input tokens that were used to generate the response."
        ),
    ]
    output_tokens: Annotated[
        int,
        Field(
            description="The number of output tokens that were generated by the model."
        ),
    ]
    total_tokens: Annotated[
        int, Field(description="The total number of tokens that were used.")
    ]
    input_tokens_details: InputTokensDetails
    output_tokens_details: OutputTokensDetails


class IncompleteDetails(BaseModel):
    reason: Annotated[
        str, Field(description="The reason the response could not be completed.")
    ]
