Blame for processor/app/schemas.py

"""Pydantic v2 schemas."""

from datetime import datetime
from typing import Any, Optional
from uuid import UUID
from pydantic import BaseModel, Field, Field


class EpisodeStart(BaseModel):
    task_description: Optional[str] = None
    name: Optional[str] = None
    fps: int = 30
    meta: dict[str, Any] = Field(default_factory=dict)


class EpisodeStartResponse(BaseModel):
    episode_id: UUID
    status: str = "collecting"


class FrameUpload(BaseModel):
    frame_index: int
    timestamp: Optional[float] = None
    observation: dict[str, Any] = Field(default_factory=dict)
    action: dict[str, Any] = Field(default_factory=dict)
    reward: Optional[float] = None
    is_terminal: bool = False
    is_truncated: bool = False


class FrameUploadResponse(BaseModel):
    frame_id: int
    image_paths: dict[str, str] = Field(default_factory=dict)


class EpisodeEndResponse(BaseModel):
    episode_id: UUID
    frame_count: int
    status: str = "completed"


class EpisodeOut(BaseModel):
    id: UUID
    session_id: Optional[UUID] = None
    name: Optional[str] = None
    task_description: Optional[str] = None
    task_id: Optional[int] = None
    status: str
    frame_count: int
    fps: int
    camera_names: Optional[list[str]] = None
    meta: Optional[dict[str, Any]] = None
    cleaning_report: Optional[dict[str, Any]] = None
    received_at: Optional[datetime] = None
    approved_at: Optional[datetime] = None
    created_at: datetime
    updated_at: datetime
    deleted_at: Optional[datetime] = None


class EpisodeListOut(BaseModel):
    episodes: list[EpisodeOut]
    total: int
    limit: int
    offset: int


class ExportStart(BaseModel):
    dataset_name: str
    episode_ids: Optional[list[UUID]] = None
    split_ratio: float = 0.9
    export_format: str = "lerobot_v3"


class BatchDownloadRequest(BaseModel):
    episode_ids: list[UUID] = Field(..., min_length=1)


class ExportJobOut(BaseModel):
    id: UUID
    dataset_name: str
    status: str
    episode_ids: Optional[list[UUID]] = None
    split_ratio: float
    progress: float
    output_dir: Optional[str] = None
    export_format: str
    error: Optional[str] = None
    created_at: datetime
    finished_at: Optional[datetime] = None


class HealthOut(BaseModel):
    status: str = "ok"
    db: str = "ok"
    storage: str = "ok"
    version: str = "0.1.0"


class MessageOut(BaseModel):
    message: str
    detail: Optional[str] = None


# ── Annotation Schemas ────────────────────────────────────────

class AnnotationKeyframeCreate(BaseModel):
    frame_index: int = Field(..., ge=0)
    event: Optional[str] = None


class AnnotationKeyframeOut(BaseModel):
    id: UUID
    segment_id: Optional[UUID] = None
    episode_id: UUID
    frame_index: int
    event: Optional[str] = None
    created_at: datetime


class AnnotationSegmentCreate(BaseModel):
    label: str = Field(..., min_length=1, max_length=256)
    start_frame_index: int = Field(..., ge=0)
    end_frame_index: int = Field(..., ge=0)
    color: Optional[str] = "#3B82F6"
    sort_order: int = 0
    notes: Optional[str] = None
    keyframes: list[AnnotationKeyframeCreate] = Field(default_factory=list)


class AnnotationSegmentUpdate(BaseModel):
    label: Optional[str] = None
    start_frame_index: Optional[int] = None
    end_frame_index: Optional[int] = None
    color: Optional[str] = None
    sort_order: Optional[int] = None
    notes: Optional[str] = None


class AnnotationSegmentOut(BaseModel):
    id: UUID
    episode_id: UUID
    label: str
    start_frame_index: int
    end_frame_index: int
    color: Optional[str] = None
    sort_order: int
    notes: Optional[str] = None
    keyframes: list[AnnotationKeyframeOut] = Field(default_factory=list)
    created_at: datetime
    updated_at: datetime


class AnnotationListOut(BaseModel):
    annotations: list[AnnotationSegmentOut]
    total: int
    episode_id: UUID


class AnnotationPerFrameOut(BaseModel):
    episode_id: UUID
    frame_count: int
    fps: int
    per_frame: list[dict[str, Any]]  # [{frame_index, annotation, annotation_index}, ...]


# ── Task Definitions ────────────────────────────────────────

class TaskDefinitionCreate(BaseModel):
    name: str = Field(..., min_length=1, max_length=128)
    description: Optional[str] = None
    claimer: Optional[str] = Field(None, max_length=128)
    target_episodes: int = Field(0, ge=0)
    params: Optional[dict[str, Any]] = None
    status: str = Field("active")


class TaskDefinitionUpdate(BaseModel):
    description: Optional[str] = None
    claimer: Optional[str] = Field(None, max_length=128)
    target_episodes: Optional[int] = Field(None, ge=0)
    params: Optional[dict[str, Any]] = None
    status: Optional[str] = None


class TaskDefinitionOut(BaseModel):
    id: UUID
    name: str
    description: Optional[str] = None
    claimer: Optional[str] = None
    target_episodes: int
    params: Optional[dict[str, Any]] = None
    status: str
    current_episodes: int = 0
    episode_progress_pct: float = 0.0
    is_complete: bool = False
    last_upload_at: Optional[datetime] = None
    created_at: datetime
    updated_at: datetime


class TaskDefinitionListOut(BaseModel):
    definitions: list[TaskDefinitionOut]
    total: int


# ── Devices ─────────────────────────────────────────────────

class DeviceHeartbeat(BaseModel):
    device_name: str = Field(..., min_length=1, max_length=128)
    meta: Optional[dict[str, Any]] = None


class DeviceOut(BaseModel):
    id: UUID
    name: str
    status: str
    meta: Optional[dict[str, Any]] = None
    first_seen_at: datetime
    last_seen_at: Optional[datetime] = None


class DeviceListOut(BaseModel):
    devices: list[DeviceOut]


class DeviceTaskOut(BaseModel):
    id: str  # task definition id as string
    name: str
    description: Optional[str] = None
    status: str
    total_required: int
    current_count: int
    assigned_at: Optional[datetime] = None
    params: Optional[dict[str, Any]] = None


# ── Dashboard Schemas ────────────────────────────────────

class StatCard(BaseModel):
    total: int
    today: int = 0
    label: str = ""


class DashboardOverview(BaseModel):
    total: StatCard
    reviewing: StatCard
    approved: StatCard
    failed: StatCard
    active_tasks: StatCard = StatCard(total=0, today=0, label="Active Tasks")
    datasets: StatCard = StatCard(total=0, today=0, label="Datasets")
    updated_at: str


class TaskProgressOut(BaseModel):
    task_name: str
    current_episodes: int
    target_episodes: int
    progress_pct: float
    status: str


class RecentEpisodeOut(BaseModel):
    id: UUID
    task_name: str
    status: str
    received_at: Optional[str] = None
    duration_sec: float = 0.0
    frame_count: int = 0
    cleaning_passed: Optional[bool] = None


class PipelineStage(BaseModel):
    name: str
    count: int


class PipelineStatus(BaseModel):
    stages: list[PipelineStage]


# ── Hand Keypoints ────────────────────────────────────

class HandKeypointData(BaseModel):
    """Per-hand keypoint + gesture data for a single frame."""
    keypoints: Optional[list[list[float]]] = None   # 21 × 2 pixel coords
    gesture: Optional[str] = None
    extended: Optional[list[str]] = None
    extended_count: Optional[int] = None
    fist: Optional[bool] = None
    pinch: Optional[bool] = None
    center_x: Optional[float] = None
    center_y: Optional[float] = None
    motion: Optional[float] = None


class HandKeypointsOut(BaseModel):
    """Per-frame hand keypoint response for skeleton overlay."""
    frame_index: int
    source: str = "none"                            # detected | interpolated | none
    hand_0: Optional[HandKeypointData] = None
    hand_1: Optional[HandKeypointData] = None
    two_hand_distance: Optional[float] = None
    contact: Optional[str] = None
    interpolated_from: Optional[int] = None          # source=interpolated
    distance_frames: Optional[int] = None


# ── Workflow Studio ─────────────────────────────────────

class ProjectCreate(BaseModel):
    name: str = Field(..., min_length=1, max_length=256)
    description: Optional[str] = None
    workflow_id: Optional[UUID] = None
    device_type: Optional[str] = None
    params: dict = Field(default_factory=dict)


class ProjectUpdate(BaseModel):
    name: Optional[str] = Field(None, min_length=1, max_length=256)
    description: Optional[str] = None
    workflow_id: Optional[UUID] = None
    device_type: Optional[str] = None
    params: Optional[dict] = None
    status: Optional[str] = None


class ProjectOut(BaseModel):
    id: UUID
    name: str
    description: Optional[str] = None
    workflow_id: Optional[UUID] = None
    workflow_name: Optional[str] = None
    device_type: Optional[str] = None
    params: dict = Field(default_factory=dict)
    status: str
    task_count: int = 0
    episode_count: int = 0
    created_at: datetime
    updated_at: datetime

    model_config = {"from_attributes": True}


class ProjectListOut(BaseModel):
    projects: list[ProjectOut]
    total: int


class WorkflowCreate(BaseModel):
    name: str = Field(..., min_length=1, max_length=256)
    description: Optional[str] = None
    graph: dict = Field(default_factory=dict, description="React Flow JSON: {nodes, edges}")
    node_configs: dict = Field(default_factory=dict)


class WorkflowUpdate(BaseModel):
    name: Optional[str] = Field(None, min_length=1, max_length=256)
    description: Optional[str] = None
    graph: Optional[dict] = None
    node_configs: Optional[dict] = None
    status: Optional[str] = None


class WorkflowOut(BaseModel):
    id: UUID
    name: str
    description: Optional[str] = None
    graph: dict
    node_configs: dict
    status: str
    created_at: datetime
    updated_at: datetime

    model_config = {"from_attributes": True}


class WorkflowListItem(BaseModel):
    id: UUID
    name: str
    status: str
    updated_at: datetime

    model_config = {"from_attributes": True}


class WorkflowListOut(BaseModel):
    workflows: list[WorkflowListItem]
    total: int
    limit: int
    offset: int


class WorkflowRunOut(BaseModel):
    id: UUID
    workflow_id: UUID
    episode_id: Optional[UUID] = None
    status: str
    started_at: Optional[datetime] = None
    finished_at: Optional[datetime] = None
    node_states: dict
    error_log: Optional[str] = None
    worker_id: Optional[str] = None
    heartbeat_at: Optional[datetime] = None
    attempt: int = 0
    progress: float = 0.0
    outputs: dict = Field(default_factory=dict)
    created_at: datetime

    model_config = {"from_attributes": True}


class WorkflowModuleOut(BaseModel):
    type: str
    slug: str
    version: str
    category: str
    label: str
    icon: str
    color: str
    inputs: list[dict] = Field(default_factory=list)
    outputs: list[dict] = Field(default_factory=list)
    defaultConfig: dict = Field(default_factory=dict)
    configSchema: list[dict] = Field(default_factory=list)
    executionTarget: str = "worker"
    capabilities: list[str] = Field(default_factory=list)


class WorkerClaimRequest(BaseModel):
    worker_id: str = Field(..., min_length=1, max_length=128)
    capabilities: list[str] = Field(default_factory=list)
    device: str = "auto"


class WorkerHeartbeatRequest(BaseModel):
    worker_id: str
    progress: float = Field(0.0, ge=0.0, le=1.0)
    node_states: dict = Field(default_factory=dict)


class WorkerFailureRequest(BaseModel):
    worker_id: str
    error: str
    retry: bool = True


class WorkerJobOut(BaseModel):
    run_id: UUID
    workflow_id: UUID
    episode_id: UUID
    workflow_name: str
    graph: dict
    node_configs: dict = Field(default_factory=dict)
    attempt: int = 0
    video_path: Optional[str] = None
    video_paths: dict[str, str] = Field(default_factory=dict)
    cameras: list[str] = Field(default_factory=list)
    camera: Optional[str] = None
    fps: int = 30
    input_url: str


class WorkerCompletionOut(BaseModel):
    run_id: UUID
    status: str
    outputs: dict = Field(default_factory=dict)