backend/ml_module/api/schemas.py

from __future__ import annotations

from datetime import datetime
from typing import Any, Dict, List, Literal, Optional

from pydantic import BaseModel, Field, ConfigDict, field_validator


class ArtifactRegistryEntry(BaseModel):
    artifact_type: str
    version: int
    path: str
    checksum: str
    size_bytes: int
    content_type: Optional[str] = None
    created_at: datetime
    status: str = Field(default="active")
    metadata: Dict[str, Any] = Field(default_factory=dict)
    superseded_at: Optional[datetime] = None
    archived_at: Optional[datetime] = None

    model_config = ConfigDict(populate_by_name=True, extra="forbid")


class ArtifactResource(ArtifactRegistryEntry):
    download_url: Optional[str] = Field(default=None, description="Presigned download URL where applicable")
    inline_preview: Optional[Dict[str, Any]] = Field(default=None, description="Lightweight preview payload for UI rendering")
    correlation_id: Optional[str] = Field(default=None, description="Correlation id linking this artifact to run events")


class PaginationMeta(BaseModel):
    page: int
    page_size: int
    total_items: int
    total_pages: int
    has_next: bool
    has_previous: bool

    model_config = ConfigDict(extra="forbid")


class ResponseMeta(BaseModel):
    warnings: List[str] = Field(default_factory=list)
    correlation_ids: List[str] = Field(default_factory=list)
    pagination: Optional[PaginationMeta] = None

    model_config = ConfigDict(extra="forbid")


class ArtifactListResponse(BaseModel):
    data: List[ArtifactResource]
    meta: ResponseMeta = Field(default_factory=ResponseMeta)

    model_config = ConfigDict(populate_by_name=True, extra="forbid")


class ArtifactResponse(BaseModel):
    data: ArtifactResource
    meta: ResponseMeta = Field(default_factory=ResponseMeta)

    model_config = ConfigDict(populate_by_name=True, extra="forbid")


class ArtifactListQuery(BaseModel):
    artifact_type: Optional[str] = Field(
        default=None,
        pattern=r"^[A-Za-z0-9_.\-]+$",
        description="Optional artifact type filter (alphanumeric, dash, underscore, dot)",
    )
    status: str = Field(default="active", description="Artifact status filter (active|archived|superseded|all)")
    page: int = Field(default=1, ge=1, le=1000)
    page_size: int = Field(default=50, ge=1, le=200)
    sort: Literal["version", "-version", "created_at", "-created_at"] = Field(default="-created_at")

    model_config = ConfigDict(populate_by_name=True, extra="forbid")

    @field_validator("status")
    @classmethod
    def validate_status(cls, value: str) -> str:
        allowed = {"active", "archived", "superseded", "all"}
        normalised = value.lower()
        if normalised not in allowed:
            raise ValueError("status must be one of: active, archived, superseded, all")
        return normalised