# generated by datamodel-codegen:
#   filename:  config-schema.yaml
#   timestamp: 2025-12-30T14:26:42+00:00

from __future__ import annotations

from enum import Enum
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, RootModel


class Region(Enum):
    """
    Platform region selection.

    `cn` uses ModelScope. `other` is reserved for future Hugging Face routing
    and currently also falls back to ModelScope until those artifacts are
    uploaded.
    """

    cn = 'cn'
    other = 'other'


class Metadata(BaseModel):
    model_config = ConfigDict(
        populate_by_name=True,
    )
    version: str = Field(
        ..., examples=['1.0.0', '2.1.3'], pattern='^\\d+\\.\\d+\\.\\d+$'
    )
    """
    Configuration version (semantic versioning)
    """
    region: Region
    """
    Platform region selection. `cn` uses ModelScope. `other` is reserved for
    future Hugging Face routing and currently also falls back to ModelScope.
    """
    cache_dir: str = Field(..., examples=['~/.lumen/models', '/opt/lumen/models'])
    """
    Model cache directory path (supports ~ expansion)
    """


class Mode(Enum):
    """
    Deployment mode
    """

    single = 'single'
    hub = 'hub'


class Service(RootModel[str]):
    model_config = ConfigDict(
        populate_by_name=True,
    )
    root: str = Field(..., pattern='^[a-z][a-z0-9_]*$')


class Deployment(BaseModel):
    model_config = ConfigDict(
        populate_by_name=True,
    )
    mode: Literal['single']
    """
    Deployment mode
    """
    service: str = Field(..., pattern='^[a-z][a-z0-9_]*$')
    """
    Service name for single mode (required if mode=single)
    """
    services: list[Service] | None = Field(None, min_length=1)
    """
    Service names for hub mode (required if mode=hub)
    """


class Deployment1(BaseModel):
    model_config = ConfigDict(
        populate_by_name=True,
    )
    mode: Literal['hub']
    """
    Deployment mode
    """
    service: str | None = Field(None, pattern='^[a-z][a-z0-9_]*$')
    """
    Service name for single mode (required if mode=single)
    """
    services: list[Service] = Field(..., min_length=1)
    """
    Service names for hub mode (required if mode=hub)
    """


class Mdns(BaseModel):
    model_config = ConfigDict(
        populate_by_name=True,
    )
    enabled: bool | None = False
    """
    Enable mDNS service discovery
    """
    service_name: str | None = Field(
        None, examples=['lumen-clip', 'lumen-hub'], pattern='^[a-z][a-z0-9-]*$'
    )
    """
    mDNS service name (required if enabled=true)
    """


class Server(BaseModel):
    model_config = ConfigDict(
        populate_by_name=True,
    )
    port: int = Field(..., ge=1024, le=65535)
    """
    gRPC server port
    """
    host: str | None = Field('0.0.0.0', examples=['0.0.0.0', '127.0.0.1', '[::]'])
    """
    Server bind address
    """
    mdns: Mdns | None = None


class ImportInfo(BaseModel):
    model_config = ConfigDict(
        populate_by_name=True,
    )
    registry_class: str = Field(
        ...,
        examples=[
            'lumen_clip.service_registry.ClipService',
            'lumen_face.service_registry.FaceService',
        ],
        pattern='^[a-z_][a-z0-9_.]*\\.[A-Z][a-zA-Z0-9]*$',
    )
    """
    Full dotted path to service registry class
    """
    add_to_server: str = Field(
        ...,
        examples=[
            'lumen_clip.proto.ml_service_pb2_grpc.add_InferenceServicer_to_server',
            'lumen_face.proto.ml_service_pb2_grpc.add_FaceServicer_to_server',
        ],
        pattern='^[a-z_][a-z0-9_.]*\\.add_[A-Za-z0-9_]+_to_server$',
    )
    """
    Full dotted path to gRPC add_to_server function
    """


class BackendSettings(BaseModel):
    """
    Optional settings for inference backend configuration.
    """

    model_config = ConfigDict(
        extra='forbid',
        populate_by_name=True,
    )
    device: str | None = None
    """
    Preferred device ('cuda', 'mps', 'cpu'). If null, auto-detects best available.
    """
    batch_size: int | None = Field(8, ge=1)
    """
    Maximum batch size for inference.
    """
    onnx_providers: list[Any] | None = None
    """
    List of ONNX execution providers. Each item can be a string or a tuple of (name, config_dict).
    """


class Runtime(Enum):
    """
    Model runtime type
    """

    torch = 'torch'
    onnx = 'onnx'
    rknn = 'rknn'


class ModelConfig(BaseModel):
    model_config = ConfigDict(
        populate_by_name=True,
    )
    model: str = Field(..., examples=['ViT-B-32', 'CN-CLIP-ViT-B-16', 'MobileCLIP-S2'])
    """
    Model repository name
    """
    runtime: Runtime
    """
    Model runtime type
    """
    rknn_device: str | None = Field(
        None, examples=['rk3566', 'rk3588'], pattern='^rk\\d+$'
    )
    """
    RKNN device identifier (required if runtime=rknn)
    """
    dataset: str | None = Field(None, examples=['ImageNet_1k', 'TreeOfLife-10M'])
    """
    Dataset name for zero-shot classification (optional)
    """
    precision: str | None = Field(None, examples=['fp32', 'int8', 'q4fp16'])
    """
    Preferred precision for running the model, valid only when runtime is 'onnx' or 'rknn'. The download validator will check the precision field in model_info.json to verify if the preferred precision is available for the current model.
    """


class Services(BaseModel):
    model_config = ConfigDict(
        populate_by_name=True,
    )
    enabled: bool
    """
    Whether to load this service
    """
    package: str = Field(
        ..., examples=['lumen_clip', 'lumen_face'], pattern='^[a-z][a-z0-9_]*$'
    )
    """
    Python package name
    """
    import_info: ImportInfo
    backend_settings: BackendSettings | None = None
    models: dict[str, ModelConfig]
    """
    Model configurations (alias → config)
    """


class LumenConfig(BaseModel):
    """
    Unified configuration schema for all Lumen ML services
    """

    model_config = ConfigDict(
        extra='forbid',
        populate_by_name=True,
    )
    metadata: Metadata
    deployment: Deployment | Deployment1
    server: Server
    services: dict[str, Services]
    """
    Service definitions
    """
