# SPDX-FileCopyrightText: 2024-2026 Patryk Ciechański
# SPDX-License-Identifier: Apache-2.0

"""Configuration models for the pluggable vision-assist runtime."""

from __future__ import annotations

from typing import Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator

VisionBackendKind = Literal["transformers_local", "mlx_local", "openai_compatible_external"]
VisionExternalProviderName = Literal["generic", "openrouter", "google_ai_studio"]
VisionContractProfile = Literal["generic_full", "google_family_compare"]
VisionSegmentationProviderName = Literal["generic_sidecar"]
VisionModelCapabilitySource = Literal["fallback_registry", "openrouter_api", "env_override", "unknown"]


class VisionModelCapabilities(BaseModel):
    """Bounded model capability metadata used for request policy decisions."""

    model_config = ConfigDict(extra="forbid")

    model_id: str
    capability_source: VisionModelCapabilitySource = "unknown"
    context_length: int | None = None
    max_completion_tokens: int | None = None
    input_modalities: list[str] = []
    output_modalities: list[str] = []
    supported_parameters: list[str] = []


class VisionTransformersLocalConfig(BaseModel):
    """Configuration for local Hugging Face/Transformers vision runtimes."""

    model_config = ConfigDict(extra="forbid")

    model_id: str | None = None
    model_path: str | None = None
    device: str = "cpu"
    dtype: str = "auto"

    @model_validator(mode="after")
    def validate_source(self) -> "VisionTransformersLocalConfig":
        """Require one explicit model source for local runtimes."""

        if not self.model_id and not self.model_path:
            raise ValueError("transformers_local backend requires model_id or model_path")
        return self


class VisionOpenAICompatibleConfig(BaseModel):
    """Configuration for external OpenAI-compatible vision endpoints."""

    model_config = ConfigDict(extra="forbid")

    provider_name: VisionExternalProviderName = "generic"
    vision_contract_profile: VisionContractProfile = "generic_full"
    base_url: str | None = None
    model: str | None = None
    api_key: str | None = None
    api_key_env: str | None = None
    site_url: str | None = None
    site_name: str | None = None
    require_parameters: bool = False
    enable_response_healing: bool = True
    prefer_json_object_for_qwen: bool = True
    model_capabilities: VisionModelCapabilities | None = None

    @model_validator(mode="after")
    def validate_endpoint(self) -> "VisionOpenAICompatibleConfig":
        """Require one explicit endpoint target for external runtimes."""

        if not self.base_url:
            raise ValueError("openai_compatible_external backend requires base_url")
        if not self.model:
            raise ValueError("openai_compatible_external backend requires model")
        return self


class VisionMLXLocalConfig(BaseModel):
    """Configuration for local Apple Silicon MLX vision runtimes."""

    model_config = ConfigDict(extra="forbid")

    model_id: str | None = None
    model_path: str | None = None

    @model_validator(mode="after")
    def validate_source(self) -> "VisionMLXLocalConfig":
        """Require one explicit model source for MLX runtimes."""

        if not self.model_id and not self.model_path:
            raise ValueError("mlx_local backend requires model_id or model_path")
        return self


class VisionRuntimeConfig(BaseModel):
    """Top-level runtime configuration for bounded vision assistance."""

    model_config = ConfigDict(extra="forbid")

    enabled: bool = False
    provider: VisionBackendKind = "transformers_local"
    allow_on_guided: bool = True
    max_images: int = Field(default=8, ge=1, le=12)
    max_tokens: int = Field(default=400, ge=1)
    timeout_seconds: float = Field(default=20.0, gt=0)
    transformers_local: VisionTransformersLocalConfig | None = None
    mlx_local: VisionMLXLocalConfig | None = None
    openai_compatible_external: VisionOpenAICompatibleConfig | None = None
    segmentation_sidecar: "VisionSegmentationSidecarConfig | None" = None

    @property
    def effective_max_tokens(self) -> int:
        """Return the output-token cap after model-capability fallback policy."""

        model_capabilities = (
            self.openai_compatible_external.model_capabilities if self.openai_compatible_external is not None else None
        )
        model_cap = model_capabilities.max_completion_tokens if model_capabilities is not None else None
        if model_cap is None:
            return self.max_tokens
        profile = self.active_vision_contract_profile
        profile_floor = 4096 if profile == "google_family_compare" else 2048
        return min(max(self.max_tokens, profile_floor), model_cap)

    @model_validator(mode="after")
    def validate_provider_config(self) -> "VisionRuntimeConfig":
        """Require configuration for the selected provider when enabled."""

        if not self.enabled:
            return self

        if self.provider == "transformers_local" and self.transformers_local is None:
            raise ValueError(
                "enabled vision runtime with provider=transformers_local requires transformers_local config"
            )

        if self.provider == "mlx_local" and self.mlx_local is None:
            raise ValueError("enabled vision runtime with provider=mlx_local requires mlx_local config")

        if self.provider == "openai_compatible_external" and self.openai_compatible_external is None:
            raise ValueError(
                "enabled vision runtime with provider=openai_compatible_external requires openai_compatible_external config"
            )
        return self

    @property
    def active_backend_config(
        self,
    ) -> VisionTransformersLocalConfig | VisionMLXLocalConfig | VisionOpenAICompatibleConfig | None:
        """Return the config block for the selected backend."""

        if self.provider == "transformers_local":
            return self.transformers_local
        if self.provider == "mlx_local":
            return self.mlx_local
        return self.openai_compatible_external

    @property
    def active_model_name(self) -> str | None:
        """Return a human-readable model name for diagnostics."""

        active = self.active_backend_config
        if active is None:
            return None
        if isinstance(active, VisionTransformersLocalConfig):
            return active.model_id or active.model_path
        if isinstance(active, VisionMLXLocalConfig):
            return active.model_id or active.model_path
        return active.model

    @property
    def active_vision_contract_profile(self) -> VisionContractProfile | None:
        """Return the resolved external vision contract profile for diagnostics."""

        if self.provider != "openai_compatible_external" or self.openai_compatible_external is None:
            return None
        return self.openai_compatible_external.vision_contract_profile

    @property
    def active_segmentation_sidecar(self) -> "VisionSegmentationSidecarConfig | None":
        """Return the optional segmentation sidecar config when enabled."""

        return self.segmentation_sidecar


class VisionSegmentationSidecarConfig(BaseModel):
    """Configuration for the optional part-segmentation sidecar."""

    model_config = ConfigDict(extra="forbid")

    enabled: bool = False
    provider_name: VisionSegmentationProviderName = "generic_sidecar"
    endpoint: str | None = None
    model: str | None = None
    api_key: str | None = None
    api_key_env: str | None = None
    timeout_seconds: float = Field(default=15.0, gt=0)
    max_parts: int = Field(default=16, ge=1, le=64)

    @model_validator(mode="after")
    def validate_endpoint(self) -> "VisionSegmentationSidecarConfig":
        """Require an explicit endpoint only when the sidecar is enabled."""

        if self.enabled and not self.endpoint:
            raise ValueError("enabled segmentation_sidecar requires endpoint")
        return self
