from __future__ import annotations

from datetime import datetime
from enum import Enum

from pydantic import BaseModel


class TaskLifecycleStatus(str, Enum):
    queued = "queued"
    running = "running"
    completed = "completed"
    failed = "failed"
    cancelled = "cancelled"

    @classmethod
    def terminal(cls) -> set["TaskLifecycleStatus"]:
        return {cls.completed, cls.failed, cls.cancelled}


class TaskSubmittedResponse(BaseModel):
    task_id: str
    status: TaskLifecycleStatus


class TaskStatusResponse(BaseModel):
    task_id: str
    status: TaskLifecycleStatus
    current_step: str | None = None
    attempt: int = 0
    updated_at: datetime | None = None
