"""Chat API data models"""

from typing import Optional, List, Dict, Any, Union
from pydantic import Field, field_validator
from .base import OctenBaseModel
from ..utils.constants import (
    REASONING_EFFORTS,
    RESPONSE_FORMAT_TYPES,
    SAFESEARCH_OPTIONS,
    WEB_SEARCH_OPTIONS,
    TOOL_CHOICE_KEYWORDS,
)


# ============================================================
# Request Models
# ============================================================


class ChatMessage(OctenBaseModel):
    """A single message in the conversation.

    Attributes:
        role: Message role. One of ``"system"``, ``"user"``, ``"assistant"``,
              ``"tool"``.
        content: Message content. Either a plain string or a list of
                 multimodal content blocks (e.g. text + image).
        tool_call_id: ID of the tool call this message responds to
                      (only for ``role="tool"``).
        tool_calls: List of tool calls made by the assistant
                    (only for ``role="assistant"``).
    """

    role: str = Field(..., description="Message role: system/user/assistant/tool")
    content: Union[str, List[Any]] = Field(..., description="Message content")
    tool_call_id: Optional[str] = Field(None, description="Tool call ID for tool messages")
    tool_calls: Optional[List[Any]] = Field(None, description="Tool calls made by assistant")


class ChatHighlightOptions(OctenBaseModel):
    """Highlight snippet configuration for chat search results.

    Attributes:
        enable: Whether to return highlighted content (default True)
        max_tokens: Maximum token count for highlighted snippet (default 300)
    """

    enable: bool = Field(True, description="Whether to return highlights")
    max_tokens: int = Field(300, ge=100, le=2048, description="Max tokens for highlight")


class ChatFullContentOptions(OctenBaseModel):
    """Full page content configuration for chat search results.

    Attributes:
        enable: Whether to fetch full page content (default False)
        max_tokens: Maximum token count for full content (default 1000)
    """

    enable: bool = Field(False, description="Whether to fetch full page content")
    max_tokens: int = Field(1000, ge=1, le=100000, description="Max tokens for full content")


class WebSearchOptions(OctenBaseModel):
    """Web search configuration.

    Attributes:
        safesearch: Safe search level. ``"off"`` or ``"strict"`` (default ``"off"``).
        count: Number of search results to retrieve (default 5).
        country: Country code for localised results (e.g. ``"CN"``).
        include_domains: Restrict results to these domains.
        exclude_domains: Exclude these domains from results.
        include_text: Only return pages containing these texts.
        exclude_text: Exclude pages containing these texts.
        time_basis: Time field to use for filtering (``"auto"``, ``"published"``,
                    ``"crawled"``).
        start_time: Earliest result time (ISO 8601).
        end_time: Latest result time (ISO 8601).
        format: Content format (``"text"`` or ``"markdown"``).
        full_content: Full page content options.
        highlight: Highlight snippet options.
    """

    safesearch: Optional[str] = Field("off", description="Safe search: off/strict")
    count: Optional[int] = Field(None, ge=1, le=100, description="Number of results")
    country: Optional[str] = Field(None, description="Country code (e.g. CN)")
    include_domains: Optional[List[str]] = Field(None, description="Restrict to domains")
    exclude_domains: Optional[List[str]] = Field(None, description="Exclude domains")
    include_text: Optional[List[str]] = Field(None, description="Must-include text filters")
    exclude_text: Optional[List[str]] = Field(None, description="Must-exclude text filters")
    time_basis: Optional[str] = Field(None, description="Time basis: auto/published/crawled")
    start_time: Optional[str] = Field(None, description="Start time (ISO 8601)")
    end_time: Optional[str] = Field(None, description="End time (ISO 8601)")
    format: Optional[str] = Field(None, description="Format: text/markdown")
    full_content: Optional[ChatFullContentOptions] = Field(None, description="Full content options")
    highlight: Optional[ChatHighlightOptions] = Field(None, description="Highlight options")

    @field_validator("safesearch")
    @classmethod
    def validate_safesearch(cls, v: Optional[str]) -> Optional[str]:
        if v is not None and v not in SAFESEARCH_OPTIONS:
            raise ValueError(f"safesearch must be one of {SAFESEARCH_OPTIONS}")
        return v


class JSONSchema(OctenBaseModel):
    """JSON Schema definition for structured output.

    Attributes:
        name: Schema name (required).
        description: Optional schema description.
        strict: Whether to enforce strict schema adherence.
        schema: The JSON Schema object.
    """

    name: str = Field(..., description="Schema name")
    description: Optional[str] = Field(None, description="Schema description")
    strict: bool = Field(False, description="Strict schema enforcement")
    schema_: Optional[Dict[str, Any]] = Field(None, alias="schema", description="JSON Schema")


class ResponseFormat(OctenBaseModel):
    """Output format specification.

    Attributes:
        type: Format type. One of ``"text"``, ``"json_object"``, ``"json_schema"``.
        json_schema: Schema definition (required when ``type="json_schema"``).
    """

    type: str = Field(..., description="Format type: text/json_object/json_schema")
    json_schema: Optional[JSONSchema] = Field(None, description="JSON Schema (for json_schema type)")

    @field_validator("type")
    @classmethod
    def validate_type(cls, v: str) -> str:
        if v not in RESPONSE_FORMAT_TYPES:
            raise ValueError(f"type must be one of {RESPONSE_FORMAT_TYPES}")
        return v


class ToolFunction(OctenBaseModel):
    """Function definition for a tool.

    Attributes:
        name: Function name.
        description: Brief description of what the function does.
        parameters: JSON Schema describing the function's parameters.
    """

    name: str = Field(..., description="Function name")
    description: Optional[str] = Field(None, description="Function description")
    parameters: Optional[Dict[str, Any]] = Field(None, description="JSON Schema for parameters")


class Tool(OctenBaseModel):
    """A tool available to the model.

    Attributes:
        type: Tool type (always ``"function"``).
        function: Function definition.
    """

    type: str = Field("function", description="Tool type (currently only 'function')")
    function: ToolFunction = Field(..., description="Function definition")


class Reasoning(OctenBaseModel):
    """Reasoning configuration.

    Attributes:
        effort: Reasoning effort level. One of ``"low"``, ``"medium"``, ``"high"``.
    """

    effort: str = Field(..., description="Reasoning effort: low/medium/high")

    @field_validator("effort")
    @classmethod
    def validate_effort(cls, v: str) -> str:
        if v not in REASONING_EFFORTS:
            raise ValueError(f"effort must be one of {REASONING_EFFORTS}")
        return v


class ChatRequest(OctenBaseModel):
    """Chat API request model.

    Attributes:
        model: Model ID (e.g. ``"openai/gpt-5.4"``).
        messages: Conversation history.
        stream: If ``True``, stream the response as SSE events.
        web_search: ``"on"`` to enable web search augmentation, ``"off"`` to disable.
        web_search_options: Fine-grained web search configuration.
        max_tokens: Maximum number of output tokens.
        max_completion_tokens: Alternative max-token parameter.
        temperature: Sampling temperature [0, 2].
        top_p: Nucleus sampling probability (0, 1].
        frequency_penalty: Frequency penalty [-2, 2].
        presence_penalty: Presence penalty [-2, 2].
        response_format: Output format specification.
        stop: Stop sequences.
        seed: Reproducibility seed.
        reasoning: Reasoning configuration.
        logprobs: Whether to return log probabilities.
        top_logprobs: Number of top log-prob tokens to return [0, 20].
        logit_bias: Token ID to bias mapping.
        tools: List of tools available to the model.
        tool_choice: Tool selection strategy or specific tool.
        user: Opaque user identifier for monitoring.
    """

    model: str = Field(..., description="Model ID")
    messages: List[ChatMessage] = Field(..., min_length=1, description="Conversation history")
    stream: Optional[bool] = Field(None, description="Enable SSE streaming")
    web_search: Optional[str] = Field(None, description="Web search: on/off")
    web_search_options: Optional[WebSearchOptions] = Field(None, description="Web search options")
    max_tokens: Optional[int] = Field(None, ge=1, description="Max output tokens")
    max_completion_tokens: Optional[int] = Field(None, ge=1, description="Max completion tokens")
    temperature: Optional[float] = Field(None, ge=0, le=2, description="Temperature [0,2]")
    top_p: Optional[float] = Field(None, gt=0, le=1, description="Top-p (0,1]")
    frequency_penalty: Optional[float] = Field(None, ge=-2, le=2, description="Frequency penalty")
    presence_penalty: Optional[float] = Field(None, ge=-2, le=2, description="Presence penalty")
    response_format: Optional[ResponseFormat] = Field(None, description="Response format")
    stop: Optional[List[str]] = Field(None, description="Stop sequences")
    seed: Optional[int] = Field(None, description="Reproducibility seed")
    reasoning: Optional[Reasoning] = Field(None, description="Reasoning configuration")
    logprobs: Optional[bool] = Field(None, description="Return log probabilities")
    top_logprobs: Optional[int] = Field(None, ge=0, le=20, description="Top log-prob tokens [0,20]")
    logit_bias: Optional[Dict[str, float]] = Field(None, description="Token bias map")
    tools: Optional[List[Tool]] = Field(None, description="Available tools")
    tool_choice: Optional[Union[str, Dict[str, Any]]] = Field(None, description="Tool choice")
    user: Optional[str] = Field(None, description="Opaque user identifier")

    @field_validator("web_search")
    @classmethod
    def validate_web_search(cls, v: Optional[str]) -> Optional[str]:
        if v is not None and v not in WEB_SEARCH_OPTIONS:
            raise ValueError(f"web_search must be one of {WEB_SEARCH_OPTIONS}")
        return v

    @field_validator("tool_choice")
    @classmethod
    def validate_tool_choice(cls, v: Optional[Union[str, Dict[str, Any]]]) -> Optional[Union[str, Dict[str, Any]]]:
        if isinstance(v, str) and v not in TOOL_CHOICE_KEYWORDS:
            raise ValueError(f"tool_choice string must be one of {TOOL_CHOICE_KEYWORDS} or a dict")
        return v


# ============================================================
# Response Models
# ============================================================


class SearchItem(OctenBaseModel):
    """A single search result item.

    Attributes:
        title: Page title.
        url: Page URL.
        highlight: Highlighted snippet.
        full_content: Full page content (when requested).
        authors: Author information.
        time_published: Publication time (ISO 8601).
        time_last_crawled: Last crawl time (ISO 8601).
    """

    title: str = Field(..., description="Page title")
    url: str = Field(..., description="Page URL")
    highlight: Optional[str] = Field(None, description="Highlighted snippet")
    full_content: Optional[str] = Field(None, description="Full page content")
    authors: Optional[str] = Field(None, description="Authors")
    time_published: Optional[str] = Field(None, description="Published time (ISO 8601)")
    time_last_crawled: Optional[str] = Field(None, description="Last crawled time (ISO 8601)")


class ChatSearchResult(OctenBaseModel):
    """Search results for a single query (within a chat response).

    Attributes:
        query: The search query that produced these results.
        results: List of result items.
    """

    query: str = Field(..., description="Search query")
    results: List[SearchItem] = Field(default_factory=list, description="Result items")


class Usage(OctenBaseModel):
    """Token usage statistics.

    Attributes:
        prompt_tokens: Number of tokens in the prompt.
        completion_tokens: Number of tokens in the completion.
        total_tokens: Total tokens used.
        num_search_queries: Number of web search queries executed.
        full_content_tokens: Tokens consumed by full page content.
        reasoning_tokens: Tokens consumed by reasoning (chain-of-thought).
    """

    prompt_tokens: int = Field(0, description="Prompt token count")
    completion_tokens: int = Field(0, description="Completion token count")
    total_tokens: int = Field(0, description="Total token count")
    num_search_queries: Optional[int] = Field(None, description="Number of search queries")
    full_content_tokens: Optional[int] = Field(None, description="Full content token count")
    reasoning_tokens: Optional[int] = Field(None, description="Reasoning token count")


class ToolCallFunction(OctenBaseModel):
    """Function call details within a tool call.

    Attributes:
        name: Function name.
        arguments: JSON-encoded function arguments.
    """

    name: Optional[str] = Field(None, description="Function name")
    arguments: Optional[str] = Field(None, description="JSON-encoded arguments")


class ToolCallItem(OctenBaseModel):
    """A tool call made by the model.

    Attributes:
        id: Unique tool call ID.
        type: Tool type (always ``"function"``).
        function: Function call details.
    """

    id: Optional[str] = Field(None, description="Tool call ID")
    type: str = Field("function", description="Tool type")
    function: ToolCallFunction = Field(default_factory=ToolCallFunction, description="Function call")


class AssistantMessage(OctenBaseModel):
    """The assistant's response message.

    Attributes:
        role: Always ``"assistant"``.
        content: Text content of the response.
        reasoning_content: Chain-of-thought reasoning (when available).
        refusal: Refusal message (when the model declines to respond).
        tool_calls: Tool calls made by the model.
    """

    role: str = Field("assistant", description="Message role")
    content: Optional[str] = Field(None, description="Response content")
    reasoning_content: Optional[str] = Field(None, description="Chain-of-thought reasoning")
    refusal: Optional[str] = Field(None, description="Refusal message")
    tool_calls: Optional[List[ToolCallItem]] = Field(None, description="Tool calls")


class Choice(OctenBaseModel):
    """A completion choice.

    Attributes:
        index: Zero-based choice index.
        finish_reason: Why generation stopped (``"stop"``, ``"length"``,
                       ``"tool_calls"``, etc.).
        message: The assistant's response message.
    """

    index: int = Field(0, description="Choice index")
    finish_reason: Optional[str] = Field(None, description="Finish reason")
    message: AssistantMessage = Field(..., description="Response message")


class ChatCompletion(OctenBaseModel):
    """Non-streaming chat completion response.

    Attributes:
        id: Unique completion ID.
        model: Model used for generation.
        created: Unix timestamp of creation.
        object: Object type (``"chat.completion"``).
        choices: List of generated choices.
        search_results: Web search results (present when ``web_search="on"``).
        meta: Metadata including usage statistics.
        warning: Optional warning message.
        citations: Citation string referencing search results.
    """

    id: str = Field(..., description="Completion ID")
    model: str = Field(..., description="Model used")
    created: float = Field(..., description="Unix timestamp")
    object: Optional[str] = Field(None, description="Object type")
    choices: List[Choice] = Field(..., description="Generated choices")
    search_results: Optional[List[ChatSearchResult]] = Field(None, description="Web search results")
    meta: Optional[Dict[str, Any]] = Field(None, description="Metadata")
    warning: Optional[str] = Field(None, description="Warning message")
    citations: Optional[str] = Field(None, description="Citation string")

    @property
    def usage(self) -> Optional[Usage]:
        """Token usage statistics extracted from ``meta``."""
        if self.meta and "usage" in self.meta:
            return Usage(**self.meta["usage"])
        return None

    @property
    def text(self) -> str:
        """Convenience accessor for the first choice's content."""
        if self.choices and self.choices[0].message.content:
            return self.choices[0].message.content
        return ""


# ============================================================
# Streaming Models
# ============================================================


class StreamDelta(OctenBaseModel):
    """Incremental content delta in a streaming chunk.

    Attributes:
        role: Role of the message being streamed (usually only in first chunk).
        content: Incremental text content.
        reasoning_content: Incremental reasoning content.
        refusal: Incremental refusal message.
        tool_calls: Incremental tool call data.
    """

    role: Optional[str] = Field(None, description="Message role (first chunk only)")
    content: Optional[str] = Field(None, description="Incremental content")
    reasoning_content: Optional[str] = Field(None, description="Incremental reasoning content")
    refusal: Optional[str] = Field(None, description="Incremental refusal message")
    tool_calls: Optional[List[Dict[str, Any]]] = Field(None, description="Incremental tool calls")


class StreamChoice(OctenBaseModel):
    """A streaming choice.

    Attributes:
        index: Zero-based choice index.
        delta: Incremental content for this chunk.
        finish_reason: Set on the final chunk (``"stop"``, ``"length"``, etc.).
    """

    index: int = Field(0, description="Choice index")
    delta: StreamDelta = Field(default_factory=StreamDelta, description="Incremental content")
    finish_reason: Optional[str] = Field(None, description="Finish reason (last chunk only)")


class StreamError(OctenBaseModel):
    """Error payload within a streaming event.

    Attributes:
        message: Human-readable error description.
        code: Numeric error code.
    """

    message: str = Field(..., description="Error message")
    code: int = Field(0, description="Error code")


class StreamEvent(OctenBaseModel):
    """A single SSE event from a streaming chat response.

    **Event types**:

    - ``"search_done"`` -- web search completed; inspect :attr:`search_results`
      and :attr:`citations`.
    - ``"content"`` -- incremental text chunk; inspect
      ``choices[0].delta.content``.
    - ``"finish"`` -- generation complete; inspect
      ``choices[0].finish_reason``.
    - ``"usage"`` -- token usage summary; inspect :attr:`usage`.
    - ``"error"`` -- in-stream error; :class:`~octen.core.exceptions.OctenStreamError`
      is raised after this event is yielded.

    Attributes:
        type: Event type string (see above).
        id: Completion ID.
        object: Object type string.
        created: Unix timestamp.
        model: Model used for generation.
        choices: List of streaming choices (present for content/finish events).
        search_results: Web search results (present in ``search_done`` event).
        meta: Metadata (present in ``usage`` event).
        citations: Citation string (present in ``search_done`` event).
        error: Error payload (present in ``error`` event).
    """

    type: Optional[str] = Field(None, description="Event type")
    id: Optional[str] = Field(None, description="Completion ID")
    object: Optional[str] = Field(None, description="Object type")
    created: Optional[int] = Field(None, description="Unix timestamp")
    model: Optional[str] = Field(None, description="Model used")
    choices: Optional[List[StreamChoice]] = Field(None, description="Streaming choices")
    search_results: Optional[List[ChatSearchResult]] = Field(None, description="Web search results")
    meta: Optional[Dict[str, Any]] = Field(None, description="Metadata")
    citations: Optional[str] = Field(None, description="Citation string")
    error: Optional[StreamError] = Field(None, description="Error payload")

    @property
    def usage(self) -> Optional[Usage]:
        """Token usage statistics (available on ``usage`` events)."""
        if self.meta and "usage" in self.meta:
            return Usage(**self.meta["usage"])
        return None
