"""Inference gateway domain models.

Core types for the capability-aware inference gateway.
Requirements: 2.1, 2.8, 2.9
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Literal
from uuid import UUID

from pydantic import BaseModel, Field


@dataclass(frozen=True)
class ProviderCapabilities:
    """Declared capabilities for an inference endpoint/deployment."""

    chat_completions: bool = False
    responses_api: bool = False
    json_schema: bool = False
    json_object: bool = False
    seed: bool = False
    usage: bool = False
    max_completion_tokens: bool = False
    reasoning_toggle: bool = False
    model_listing: bool = False


@dataclass(frozen=True)
class InferenceTarget:
    """Resolved target for an inference request."""

    endpoint_id: UUID
    deployment_id: UUID
    protocol: Literal["ollama_native", "openai_chat", "specialist_http"]
    base_url: str
    model: str
    capabilities: ProviderCapabilities
    auth_secret_ref: str | None = None
    auth_scheme: str = "bearer"
    extra_headers: dict[str, str] = field(default_factory=dict)
    extra_body: dict[str, Any] = field(default_factory=dict)
    max_retries: int = 3
    timeout_seconds: float = 120.0
    context_window: int = 0
    max_output_tokens: int | None = None


class ChatMessage(BaseModel):
    """A single chat message."""

    role: Literal["system", "user", "assistant"]
    content: str


class StructuredGenerationRequest(BaseModel):
    """Request for structured generation via the inference gateway."""

    messages: list[ChatMessage]
    json_schema: dict[str, Any] | None = None
    max_output_tokens: int = 4096
    temperature: float = 0.0
    seed: int | None = 0
    timeout_seconds: float = 120.0
    trace_id: str = ""


class TokenUsage(BaseModel):
    """Token usage metadata from an inference response."""

    input_tokens: int | None = None
    output_tokens: int | None = None
    total_tokens: int | None = None


class InferenceResult(BaseModel):
    """Result of an inference request.

    Contains the generated content plus metadata for lineage,
    observability, and audit.
    """

    content: str
    parsed: dict[str, Any] | None = None
    endpoint_id: UUID | None = None
    deployment_id: UUID | None = None
    model: str = ""
    protocol: Literal["ollama_native", "openai_chat", "specialist_http"] = "openai_chat"
    structured_mode: Literal["json_schema", "json_object", "prompt_only", "none"] = "none"
    latency_ms: int = 0
    usage: TokenUsage = Field(default_factory=TokenUsage)
    request_id: str | None = None
    finish_reason: str | None = None
    repaired: bool = False
    retries: int = 0
    error: str | None = None
    error_category: str | None = None
    schema_valid: bool | None = None


class ModelLineage(BaseModel):
    """Lineage record capturing which endpoint, model, and route served a request.

    Used for persistence so actual endpoint, model, and route lineage are
    recorded (fixes hardcoded model_provider = 'ollama').
    """

    endpoint_id: UUID | None = None
    deployment_id: UUID | None = None
    model: str = ""
    protocol: Literal["ollama_native", "openai_chat", "specialist_http"] = "openai_chat"
    structured_mode: str = "none"
    request_id: str | None = None
    latency_ms: int = 0
    retries: int = 0
    trace_id: str = ""


# Error categories for provider failures
class ErrorCategory:
    """Normalized error categories for inference failures."""

    TIMEOUT = "timeout"
    AUTHENTICATION = "authentication"
    RATE_LIMIT = "rate_limit"
    SERVER_ERROR = "server_error"
    INVALID_RESPONSE = "invalid_response"
    SCHEMA_VIOLATION = "schema_violation"
    CAPABILITY_ERROR = "capability_error"
    POLICY_ERROR = "policy_error"
    CONNECTION_ERROR = "connection_error"