models.py

from pydantic import BaseModel, Field
from typing import Any, Optional, Dict
from enum import Enum
import uuid
import time


class JobStatus(str, Enum):
    PENDING = "pending"
    CLAIMED = "claimed"
    COMPLETED = "completed"
    FAILED = "failed"
    TIMEOUT = "timeout"


class APIJob(BaseModel):
    job_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    created_at: float = Field(default_factory=time.time)
    claimed_at: Optional[float] = None
    completed_at: Optional[float] = None

    # Request fields
    method: str  # GET, POST, PUT, DELETE, etc.
    endpoint: str  # e.g. "/api/v1/inference"
    headers: Dict[str, str] = {}
    body: Optional[Any] = None
    query_params: Dict[str, str] = {}

    # Routing: which mirror should handle this (optional, None = any mirror that has it)
    target_mirror: Optional[str] = None

    # Response fields
    status: JobStatus = JobStatus.PENDING
    response_status_code: Optional[int] = None
    response_headers: Dict[str, str] = {}
    response_body: Optional[Any] = None
    error: Optional[str] = None

    # TTL: jobs older than this (seconds) are considered timed out
    ttl: float = 30.0


class ClaimRequest(BaseModel):
    mirror_id: str
    available_endpoints: list[str]  # list of endpoint prefixes this mirror can serve


class CompleteRequest(BaseModel):
    mirror_id: str
    job_id: str
    response_status_code: int
    response_headers: Dict[str, str] = {}
    response_body: Optional[Any] = None
    error: Optional[str] = None