lotusgen/backend/schemas/models.py

import re
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel, ConfigDict, Field, field_validator

SLUG_PATTERN = r"^[a-z0-9-]+$"


def validate_slug_format(slug: Optional[str]) -> Optional[str]:
    if slug is not None and slug.strip():
        cleaned = slug.strip().lower()
        if not re.match(SLUG_PATTERN, cleaned):
            raise ValueError("Slug can only contain lowercase alphanumerics and hyphens (a-z, 0-9, -).")
        return cleaned
    return None


class FieldSchema(BaseModel):
    name: str
    role: Literal["aggregated", "randomized", "function", "reference", "static"]
    type: Literal["string", "integer", "float", "datetime", "boolean", "uuid7", "choice", "collection"]
    active: bool = True
    config: Optional[Dict[str, Any]] = Field(default_factory=dict)


class TableCreate(BaseModel):
    name: str
    slug: Optional[str] = None
    archetype: Literal["record", "aggregate", "reference", "collection"]
    format: Optional[Literal["csv", "parquet"]] = "csv"
    row_limit: Optional[int] = 100000
    exclude_if: Optional[str] = None
    fields: Optional[List[FieldSchema]] = Field(default_factory=list)
    columns: Optional[List[str]] = Field(default_factory=list)
    data: Optional[List[List[Any]]] = Field(default_factory=list)

    @field_validator("slug")
    @classmethod
    def check_slug(cls, v: Optional[str]) -> Optional[str]:
        return validate_slug_format(v)


class TableUpdate(BaseModel):
    name: Optional[str] = None
    slug: Optional[str] = None
    format: Optional[Literal["csv", "parquet"]] = None
    row_limit: Optional[int] = None
    exclude_if: Optional[str] = None
    fields: Optional[List[FieldSchema]] = None
    columns: Optional[List[str]] = None
    data: Optional[List[List[Any]]] = None

    @field_validator("slug")
    @classmethod
    def check_slug(cls, v: Optional[str]) -> Optional[str]:
        return validate_slug_format(v)


class TableResponse(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: str
    project_id: str
    name: str
    slug: str = ""
    archetype: Literal["record", "aggregate", "reference", "collection"]
    status: Literal["draft", "valid"]
    format: Literal["csv", "parquet"]
    row_limit: Optional[int] = None
    exclude_if: Optional[str] = None
    fields: List[FieldSchema] = Field(default_factory=list)
    columns: List[str] = Field(default_factory=list)
    data: List[List[Any]] = Field(default_factory=list)
    referencing_tables: List[str] = Field(default_factory=list)
    created_at: datetime
    updated_at: datetime


class ProjectCreate(BaseModel):
    name: str
    slug: Optional[str] = None
    description: Optional[str] = ""
    seed: Optional[int] = 42

    @field_validator("slug")
    @classmethod
    def check_slug(cls, v: Optional[str]) -> Optional[str]:
        return validate_slug_format(v)


class ProjectUpdate(BaseModel):
    name: Optional[str] = None
    slug: Optional[str] = None
    description: Optional[str] = None
    seed: Optional[int] = None

    @field_validator("slug")
    @classmethod
    def check_slug(cls, v: Optional[str]) -> Optional[str]:
        return validate_slug_format(v)


class ProjectCardResponse(BaseModel):
    id: str
    name: str
    slug: str = ""
    description: str
    seed: int
    table_count: int
    record_count: int
    aggregate_count: int
    reference_count: int
    collection_count: int = 0
    valid_count: int
    created_at: datetime
    updated_at: datetime


class ProjectResponse(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: str
    name: str
    slug: str = ""
    description: str
    seed: int
    tables: List[TableResponse] = Field(default_factory=list)
    created_at: datetime
    updated_at: datetime


class ExportTableSchema(BaseModel):
    id: str
    name: str
    archetype: Literal["record", "aggregate", "reference", "collection"]
    status: Literal["draft", "valid"]
    format: Optional[Literal["csv", "parquet"]] = "csv"
    row_limit: Optional[int] = None
    exclude_if: Optional[str] = None
    fields: Optional[List[FieldSchema]] = None
    columns: Optional[List[str]] = None
    data: Optional[List[List[Any]]] = None


class ProjectExportSchema(BaseModel):
    version: str = "1.0"
    id: str
    name: str
    description: str
    seed: int
    tables: List[ExportTableSchema]


class DryRunResponse(BaseModel):
    success: bool
    sample_rows: List[Dict[str, Any]] = Field(default_factory=list)
    error: Optional[str] = None


class JobResponse(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: str
    project_id: str
    table_id: str
    table_name: str
    status: Literal["pending", "running", "completed", "failed", "cancelled"]
    format: Literal["csv", "parquet"]
    total_rows: int
    rows_generated: int
    rows_per_sec: float
    artifact_path: Optional[str] = None
    file_size_bytes: Optional[int] = None
    error_message: Optional[str] = None
    created_at: datetime
    completed_at: Optional[datetime] = None