from __future__ import annotations

import enum
from typing import Any
import pydantic


class AiProviderType(str, enum.Enum):
    """Canonical enterprise AI providers."""

    ANTIGRAVITY = "antigravity"
    OPENROUTER = "openrouter"
    LMSTUDIO = "lmstudio"
    GEMINI_DIRECT = "gemini"

    @classmethod
    def from_str(cls, value: str | None) -> "AiProviderType":
        """Normalize various string forms (e.g. 'antigravity_sdk') into canonical Enum."""
        if not value:
            return cls.ANTIGRAVITY
        v = str(value).lower().strip()
        if v in ("antigravity", "antigravity_sdk", "google_antigravity"):
            return cls.ANTIGRAVITY
        elif v in ("openrouter", "open_router"):
            return cls.OPENROUTER
        elif v in ("lmstudio", "lm_studio", "local"):
            return cls.LMSTUDIO
        elif v in ("gemini", "google", "gemini_direct"):
            return cls.GEMINI_DIRECT
        return cls.ANTIGRAVITY


class AiCapability(str, enum.Enum):
    """Capabilities supported by models in DSCons ERP."""

    TEXT_GENERATION = "text_generation"
    MULTIMODAL_VISION = "multimodal_vision"
    STRUCTURED_JSON = "structured_json"
    CODE_EXECUTION = "code_execution"


class AiModelSpec(pydantic.BaseModel):
    """Specification of an AI model registered in the Enterprise Registry."""

    canonical_id: str
    provider_type: AiProviderType
    provider_model_id: str
    display_name: str
    capabilities: list[AiCapability] = pydantic.Field(default_factory=list)
    context_window: int = 1048576
    max_output_tokens: int = 4096
    fallback_canonical_id: str | None = None
    is_active: bool = True

    model_config = pydantic.ConfigDict(frozen=True)


class AiTextRequest(pydantic.BaseModel):
    """Input contract for standard text and chat completion."""

    prompt: str = ""
    system_prompt: str | None = None
    messages: list[dict[str, Any]] | None = None
    canonical_model_id: str | None = None
    preferred_provider: AiProviderType | None = None
    agent_code: str | None = None
    temperature: float = 0.2
    max_tokens: int | None = None
    extra_body: dict[str, Any] | None = None


class AiVisionRequest(pydantic.BaseModel):
    """Input contract for multimodal vision / document analysis."""

    prompt: str
    image_bytes: bytes | None = None
    image_mime: str = "image/jpeg"
    data_uri: str | None = None
    file_path: str | None = None
    canonical_model_id: str | None = None
    preferred_provider: AiProviderType | None = None
    agent_code: str | None = None
    temperature: float = 0.1
    max_tokens: int = 2000


class AiStructuredRequest(pydantic.BaseModel):
    """Input contract for strictly typed JSON extraction."""

    prompt: str
    system_prompt: str | None = None
    schema_dict: dict[str, Any] | None = None
    canonical_model_id: str | None = None
    preferred_provider: AiProviderType | None = None
    agent_code: str | None = None
    temperature: float = 0.1


class AiResponse(pydantic.BaseModel):
    """Normalized response from any AI provider."""

    content: str
    structured_data: dict[str, Any] | None = None
    model_used: str
    provider_used: AiProviderType
    latency_ms: float = 0.0
    prompt_tokens: int = 0
    completion_tokens: int = 0
    cost_usd: float = 0.0
    is_fallback: bool = False