from __future__ import annotations

from datetime import datetime
from enum import Enum

from pydantic import BaseModel, Field

from ai_orchestrator.schemas.task_state import utc_now


class QueueJobStatus(str, Enum):
    queued = "queued"
    leased = "leased"
    completed = "completed"
    failed = "failed"
    cancelled = "cancelled"


class QueueJob(BaseModel):
    job_id: str
    task_id: str
    engine_name: str
    status: QueueJobStatus = QueueJobStatus.queued
    attempts: int = 0
    max_attempts: int = 5
    worker_id: str | None = None
    lease_until: datetime | None = None
    available_at: datetime = Field(default_factory=utc_now)
    last_heartbeat: datetime | None = None
    error_message: str | None = None
    created_at: datetime = Field(default_factory=utc_now)
    updated_at: datetime = Field(default_factory=utc_now)
