from __future__ import annotations

from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Union

from pydantic import BaseModel, SerializeAsAny

OCIRoles = Literal["SYSTEM", "USER", "ASSISTANT", "TOOL"]


class OCIVendors(Enum):
    """
    A class to hold the vendor names for OCI models.
    This is used to map model names to their respective vendors.
    """

    COHERE = "COHERE"
    GENERIC = "GENERIC"


# --- Base Models and Content Parts ---


class OCIContentPart(BaseModel):
    """Base model for content parts in an OCI message."""

    pass


class OCITextContentPart(OCIContentPart):
    """Text content part for the OCI API."""

    type: Literal["TEXT"] = "TEXT"
    text: str


class OCIImageUrl(BaseModel):
    """ImageUrl object for OCI API. See: https://docs.oracle.com/en-us/iaas/tools/python/latest/api/generative_ai_inference/models/oci.generative_ai_inference.models.ImageUrl.html"""

    url: str
    detail: Optional[Literal["AUTO", "HIGH", "LOW"]] = None


class OCIImageContentPart(OCIContentPart):
    """Image content part for the OCI API."""

    type: Literal["IMAGE"] = "IMAGE"
    imageUrl: OCIImageUrl


OCIContentPartUnion = Union[OCITextContentPart, OCIImageContentPart]

# --- Models for Tools and Tool Calls ---


class OCIToolCall(BaseModel):
    """Represents a tool call made by the model."""

    id: Optional[str] = None  # absent in some provider responses (e.g. Google via OCI)
    type: Literal["FUNCTION"] = "FUNCTION"
    name: str
    arguments: str  # Arguments should be a JSON-serialized string


class OCIToolDefinition(BaseModel):
    """Defines a tool that can be used by the model."""

    type: Literal["FUNCTION"] = "FUNCTION"
    name: Optional[str] = None
    description: Optional[str] = None
    parameters: Optional[dict] = None


# --- Message Models (Request and Response) ---


class OCIMessage(BaseModel):
    """Model for a single message in the request/response payload."""

    role: OCIRoles
    content: Optional[List[OCIContentPartUnion]] = None
    toolCalls: Optional[List[OCIToolCall]] = None
    toolCallId: Optional[str] = None


# --- Request Payload Models ---


class OCIChatRequestPayload(BaseModel):
    """Internal 'chatRequest' payload for the OCI API."""

    apiFormat: str
    messages: List[OCIMessage]
    tools: Optional[List[OCIToolDefinition]] = None
    isStream: bool = False
    numGenerations: Optional[int] = None
    maxTokens: Optional[int] = None
    # GPT-5+ on OCI rejects maxTokens and requires maxCompletionTokens.
    maxCompletionTokens: Optional[int] = None
    temperature: Optional[float] = None
    topP: Optional[float] = None
    stop: Optional[List[str]] = None
    seed: Optional[int] = None
    frequencyPenalty: Optional[float] = None
    presencePenalty: Optional[float] = None
    # Reasoning-token budget knob (OCI: NONE/MINIMAL/LOW/MEDIUM/HIGH).
    # Honoured by GPT-5 family, Gemini 2.5, Grok reasoning variants,
    # Cohere Command-A-Reasoning. Ignored by non-reasoning models.
    reasoningEffort: Optional[str] = None
    responseFormat: Optional[Dict[str, Any]] = None
    toolChoice: Optional[Union[str, Dict[str, Any]]] = None
    logitBias: Optional[Dict[str, Any]] = None
    logProbs: Optional[int] = None


class OCIServingMode(BaseModel):
    """Defines the serving mode and the model to be used."""

    servingType: str
    endpointId: Optional[str] = None
    modelId: Optional[str] = None


class OCICompletionPayload(BaseModel):
    """Pydantic model for the complete OCI chat request body."""

    compartmentId: str
    servingMode: OCIServingMode
    chatRequest: Union[OCIChatRequestPayload, CohereChatRequest]


# --- API Response Models (Non-streaming) ---


class OCICompletionTokenDetails(BaseModel):
    """Completion token details in the OCI response."""

    acceptedPredictionTokens: Optional[int] = None
    reasoningTokens: Optional[int] = None


class OCIPromptTokensDetails(BaseModel):
    """Prompt token details in the OCI response."""

    cachedTokens: Optional[int] = None


class OCIResponseUsage(BaseModel):
    """Token usage in the OCI response."""

    promptTokens: int
    # completionTokens may be absent for reasoning models when all the output
    # budget is consumed by reasoning tokens before any visible content is produced.
    completionTokens: Optional[int] = None
    totalTokens: int
    completionTokensDetails: Optional[OCICompletionTokenDetails] = None
    promptTokensDetails: Optional[OCIPromptTokensDetails] = None


class OCIResponseChoice(BaseModel):
    """A completion choice in the OCI response."""

    index: int
    # message is absent when a reasoning model exhausts max_tokens in the
    # reasoning phase without producing any visible content.
    message: Optional[OCIMessage] = None
    finishReason: Optional[str] = None
    logprobs: Optional[Dict[str, Any]] = None


class OCIChatResponse(BaseModel):
    """The 'chatResponse' object in the OCI response."""

    apiFormat: str
    timeCreated: str
    choices: List[OCIResponseChoice]
    usage: OCIResponseUsage


class OCICompletionResponse(BaseModel):
    """Model for the complete non-streaming OCI response body."""

    modelId: str
    modelVersion: str
    chatResponse: OCIChatResponse


# --- API Response Models (Streaming) ---


class OCIStreamDelta(BaseModel):
    """The content delta in a streaming chunk."""

    content: Optional[List[OCIContentPartUnion]] = None
    role: Optional[str] = None
    toolCalls: Optional[List[OCIToolCall]] = None


class OCIStreamChunk(BaseModel):
    """Model for a single SSE event chunk from OCI."""

    finishReason: Optional[str] = None
    message: Optional[OCIStreamDelta] = None
    pad: Optional[str] = None
    index: Optional[int] = None


# --- Cohere-Specific Models ---


class CohereStreamChunk(BaseModel):
    """Model for a single SSE event chunk from OCI Cohere API."""

    apiFormat: str
    text: Optional[str] = None
    chatHistory: Optional[List[CohereMessage]] = None
    finishReason: Optional[str] = None
    toolCalls: Optional[List[CohereToolCall]] = None
    pad: Optional[str] = None
    index: Optional[int] = None


class CohereMessage(BaseModel):
    """Base model for Cohere messages."""

    role: str
    message: Optional[str] = None
    toolCalls: Optional[List[CohereToolCall]] = None


class CohereUserMessage(CohereMessage):
    """User message in Cohere chat."""

    role: Literal["USER"] = "USER"


class CohereChatBotMessage(CohereMessage):
    """Chatbot message in Cohere chat."""

    role: Literal["CHATBOT"] = "CHATBOT"


class CohereSystemMessage(CohereMessage):
    """System message in Cohere chat."""

    role: Literal["SYSTEM"] = "SYSTEM"


class CohereToolMessage(CohereMessage):
    """Tool message in Cohere chat.

    The OCI Cohere API represents tool results via a ``toolResults`` list on the
    TOOL-role history entry — not via a ``toolCallId`` string.
    """

    role: Literal["TOOL"] = "TOOL"
    toolResults: List[CohereToolResult]


class CohereParameterDefinition(BaseModel):
    """Parameter definition for Cohere tools."""

    description: str
    type: str
    isRequired: bool = False


class CohereTool(BaseModel):
    """Tool definition for Cohere."""

    name: str
    description: str
    parameterDefinitions: Dict[str, CohereParameterDefinition]


class CohereToolCall(BaseModel):
    """Tool call made by Cohere model."""

    name: str
    parameters: Dict[str, Any]


class CohereToolResult(BaseModel):
    """Result of a tool call.

    Matches the OCI SDK's CohereToolResult: each result carries the originating
    tool call (name + parameters) and a list of output objects.
    """

    call: CohereToolCall
    outputs: List[Dict[str, Any]]


class CohereChatRequest(BaseModel):
    """Cohere chat request model."""

    # Required fields
    message: str
    apiFormat: Literal["COHERE"] = "COHERE"

    # Optional fields
    # ``SerializeAsAny`` preserves subclass-specific fields (e.g. ``toolResults``
    # on ``CohereToolMessage``) when this request is serialized via ``model_dump``.
    # Without it, Pydantic v2 would serialize each element using the declared
    # ``CohereMessage`` schema and silently drop subclass fields.
    chatHistory: Optional[List[SerializeAsAny[CohereMessage]]] = None
    maxTokens: Optional[int] = None
    temperature: Optional[float] = None
    topP: Optional[float] = None
    topK: Optional[int] = None
    frequencyPenalty: Optional[float] = None
    presencePenalty: Optional[float] = None
    stopSequences: Optional[List[str]] = None
    seed: Optional[int] = None
    tools: Optional[List[CohereTool]] = None
    # NOTE: OCI's Cohere chat endpoint does not accept ``toolChoice`` — see
    # ``OCIChatConfig.openai_to_oci_cohere_param_map`` which marks
    # ``tool_choice`` as unsupported. The field is intentionally absent here
    # so it isn't silently dropped or surfaced as a supported feature.
    # OCI Cohere responseFormat is {"type": "TEXT" | "JSON_OBJECT", "schema"?: ...};
    # there is no JSON_SCHEMA type. The shape is built in
    # OCIChatConfig._normalize_response_format.
    responseFormat: Optional[Dict[str, Any]] = None
    preambleOverride: Optional[str] = None
    documents: Optional[List[Dict[str, Any]]] = None
    searchQueriesOnly: Optional[bool] = None
    searchEntryPoint: Optional[str] = None
    grounding: Optional[Dict[str, Any]] = None
    isEcho: Optional[bool] = None
    isSearchQueriesOnly: Optional[bool] = None
    isRawPrompting: Optional[bool] = None
    isForceSingleStep: Optional[bool] = None
    promptTruncation: Optional[str] = None
    safetyMode: Optional[str] = None
    citationQuality: Optional[str] = None
    maxInputTokens: Optional[int] = None
    isStream: Optional[bool] = None
    streamOptions: Optional[Dict[str, Any]] = None


class CohereUsage(BaseModel):
    """Usage information for Cohere response."""

    promptTokens: int
    completionTokens: int
    totalTokens: int
    promptTokensDetails: Optional[Dict[str, Any]] = None
    completionTokensDetails: Optional[Dict[str, Any]] = None


class CohereCitation(BaseModel):
    """Citation in Cohere response."""

    start: int
    end: int
    text: str
    document_ids: List[str]


class CohereSearchQuery(BaseModel):
    """Search query generated by Cohere."""

    text: str
    generation_id: str


class CohereChatResponse(BaseModel):
    """Cohere chat response model."""

    # Required fields
    text: str
    apiFormat: Literal["COHERE"] = "COHERE"
    # Accept any string (with ``None`` for absent) so unknown finish reasons
    # — e.g. a value OCI adds in a future API revision — degrade gracefully
    # via ``handle_cohere_response``'s ``elif oci_finish_reason is not None``
    # fallback instead of crashing Pydantic validation. Mirrors
    # ``CohereStreamChunk.finishReason`` which has always been ``Optional[str]``.
    finishReason: Optional[str] = None

    # Optional fields
    chatHistory: Optional[List[CohereMessage]] = None
    citations: Optional[List[CohereCitation]] = None
    documents: Optional[List[Dict[str, Any]]] = None
    errorMessage: Optional[str] = None
    isSearchRequired: Optional[bool] = None
    prompt: Optional[str] = None
    searchQueries: Optional[List[CohereSearchQuery]] = None
    toolCalls: Optional[List[CohereToolCall]] = None
    usage: Optional[CohereUsage] = None


class CohereChatDetails(BaseModel):
    """Chat details for Cohere request."""

    compartmentId: str
    servingMode: OCIServingMode
    chatRequest: CohereChatRequest


class CohereChatResult(BaseModel):
    """Complete Cohere chat result."""

    modelId: str
    modelVersion: str
    chatResponse: CohereChatResponse


# ---------------------------------------------------------------------------
# OCI Embed types
# ---------------------------------------------------------------------------


class OCIEmbedRequest(BaseModel):
    """Request body for POST /20231130/actions/embedText."""

    compartmentId: str
    servingMode: OCIServingMode
    inputs: List[str]
    inputType: Optional[str] = (
        None  # SEARCH_DOCUMENT | SEARCH_QUERY | CLASSIFICATION | CLUSTERING | IMAGE
    )
    truncate: Optional[str] = "END"  # NONE | START | END
    outputDimensions: Optional[int] = (
        None  # cohere.embed-v4.0+; valid: 256, 512, 1024, 1536
    )


class OCIEmbedUsage(BaseModel):
    promptTokens: int
    totalTokens: int


class OCIEmbedResponse(BaseModel):
    """Response body from POST /20231130/actions/embedText."""

    id: Optional[str] = None  # present in the official SDK response
    embeddings: List[List[float]]
    modelId: str
    modelVersion: str
    # OCI returns per-input token counts in inputTextTokenCounts (summed for total usage)
    inputTextTokenCounts: Optional[List[int]] = None
    # Some deployments may return a usage object instead
    usage: Optional[OCIEmbedUsage] = None
