"""LLM response schema types for Fluxibly.

Defines the standardized output format returned by all LLM implementations:
LLMResponse, LLMOutput, ContentItem, LLMMetadata, TokenUsage.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from pydantic import BaseModel, Field

from fluxibly.schema.tools import Citation, ToolCall, ToolResult

if TYPE_CHECKING:
    from fluxibly.llm.base import PricingConfig


class ContentItem(BaseModel):
    """A single content item in the output.

    Captures all output types: text, reasoning, tool calls, multimodal, etc.

    The ``type`` field is an open string — not restricted to a fixed enum.
    This allows forward-compatibility with new content types from providers
    without requiring framework changes.

    Known types (non-exhaustive):
    - "text"         — text response
    - "reasoning"    — model reasoning/thinking (OpenAI, Anthropic, Gemini)
    - "tool_call"    — function/tool call request
    - "tool_result"  — result from tool execution
    - "image"        — generated image output (OpenAI, Gemini)
    - "audio"        — generated audio/speech output (OpenAI, Gemini)
    - "video"        — generated video output (future)
    - "refusal"      — content policy refusal
    - "code"         — code execution (Gemini executableCode)
    - "code_result"  — code execution result (Gemini codeExecutionResult)

    New types can be added by providers or custom implementations.
    Unknown types are preserved as-is with raw data in ``extra``.
    """

    type: str = Field(
        ..., description="Content type (open string — not restricted)"
    )

    # ── Text ──
    text: str | None = Field(default=None, description="Text content")
    annotations: list[Citation] | None = Field(
        default=None, description="Citations/annotations on text"
    )

    # ── Reasoning / Thinking ──
    reasoning_text: str | None = Field(
        default=None, description="Reasoning/thinking text"
    )

    # ── Tool Calls ──
    tool_call: ToolCall | None = Field(
        default=None, description="Tool call data"
    )
    tool_result: ToolResult | None = Field(
        default=None, description="Tool result data"
    )

    # ── Image (generated output) ──
    image_data: str | None = Field(
        default=None, description="Base64-encoded image data"
    )
    image_media_type: str | None = Field(
        default=None, description="Image MIME type, e.g. 'image/png'"
    )

    # ── Audio (generated output / TTS) ──
    audio_data: str | None = Field(
        default=None, description="Base64-encoded audio data"
    )
    audio_media_type: str | None = Field(
        default=None, description="Audio MIME type, e.g. 'audio/wav'"
    )
    audio_transcript: str | None = Field(
        default=None, description="Text transcript of generated audio"
    )

    # ── Video (generated output — future) ──
    video_data: str | None = Field(
        default=None, description="Base64-encoded video data"
    )
    video_media_type: str | None = Field(
        default=None, description="Video MIME type, e.g. 'video/mp4'"
    )
    video_url: str | None = Field(
        default=None, description="URL to generated video"
    )

    # ── Code Execution (Gemini) ──
    code: str | None = Field(default=None, description="Generated code")
    code_language: str | None = Field(
        default=None, description="Code language, e.g. 'python'"
    )
    code_output: str | None = Field(
        default=None, description="Code execution output"
    )

    # ── Refusal ──
    refusal: str | None = Field(default=None, description="Refusal message")

    # ── Extensibility ──
    extra: dict[str, Any] | None = Field(
        default=None,
        description="Provider-specific data for unknown/new types",
    )


class LLMOutput(BaseModel):
    """The actual output content from the LLM."""

    # --- Primary text output (guaranteed field) ---
    output_text: str = Field(
        default="",
        description="Final text response. Concatenation of all text content items.",
    )

    # --- Structured content items (preserves full detail) ---
    content: list[ContentItem] = Field(
        default_factory=list, description="All output items in order"
    )

    # --- Convenience accessors (extracted from content items) ---
    tool_calls: list[ToolCall] = Field(
        default_factory=list, description="Extracted tool/function calls"
    )
    reasoning: str | None = Field(
        default=None, description="Extracted reasoning/thinking text"
    )
    citations: list[Citation] = Field(
        default_factory=list, description="Extracted citations/annotations"
    )
    images: list[ContentItem] = Field(
        default_factory=list, description="Extracted image content items"
    )
    audio: list[ContentItem] = Field(
        default_factory=list, description="Extracted audio content items"
    )


class TokenUsage(BaseModel):
    """Standardized token usage across all providers."""

    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0
    reasoning_tokens: int | None = None
    cached_tokens: int | None = None

    # --- Cost (populated when PricingConfig is set in LLMConfig) ---
    input_cost: float | None = None
    output_cost: float | None = None
    reasoning_cost: float | None = None
    cached_cost: float | None = None
    total_cost: float | None = None

    def calculate_cost(self, pricing: PricingConfig) -> None:
        """Calculate costs from token counts and pricing config.

        Called automatically by LLM._parse_response() when pricing is configured.
        All prices are in USD per 1M tokens.
        """
        self.input_cost = (self.input_tokens / 1_000_000) * pricing.input_price
        self.output_cost = (
            self.output_tokens / 1_000_000
        ) * pricing.output_price

        if self.cached_tokens and pricing.cached_input_price is not None:
            self.cached_cost = (
                self.cached_tokens / 1_000_000
            ) * pricing.cached_input_price
            # Cached tokens are cheaper — adjust input cost to only count non-cached tokens
            self.input_cost = (
                (self.input_tokens - self.cached_tokens) / 1_000_000
            ) * pricing.input_price

        if self.reasoning_tokens:
            rate = (
                pricing.reasoning_price
                if pricing.reasoning_price is not None
                else pricing.output_price
            )
            self.reasoning_cost = (self.reasoning_tokens / 1_000_000) * rate

        self.total_cost = (
            (self.input_cost or 0)
            + (self.output_cost or 0)
            + (self.cached_cost or 0)
            + (self.reasoning_cost or 0)
        )


class LLMMetadata(BaseModel):
    """Metadata about the LLM response."""

    model: str = Field(..., description="Model ID used")
    id: str = Field(default="", description="Response/completion unique ID")
    created_at: str = Field(default="", description="ISO timestamp")
    status: str = Field(
        default="completed",
        description="'completed' | 'in_progress' | 'incomplete' | 'failed'",
    )

    # --- Token Usage ---
    usage: TokenUsage = Field(default_factory=TokenUsage)

    # --- Stop/Finish info ---
    stop_reason: str | None = Field(
        default=None,
        description="Normalized: 'stop' | 'tool_calls' | 'max_tokens' | 'content_filter' | 'error'",
    )

    # --- Provider-specific extras ---
    provider: str = Field(
        default="",
        description="'openai' | 'anthropic' | 'gemini' | 'langchain' | 'litellm'",
    )
    raw_response: dict[str, Any] | None = Field(
        default=None, description="Original provider response (for debugging)"
    )
    extra: dict[str, Any] = Field(
        default_factory=dict,
        description="Any other provider-specific metadata",
    )


class LLMResponse(BaseModel):
    """Standardized response from any LLM forward() call."""

    output: LLMOutput = Field(default_factory=LLMOutput)
    metadata: LLMMetadata
