from enum import StrEnum

from typing import Annotated, Optional
from pydantic import BaseModel, Field


class TaskStatus(StrEnum):
    QUEUED = "queued"
    RUNNING = "running"
    DONE = "done"
    CANCELLED = "cancelled"

    @classmethod
    def _missing_(cls, value):
        value = value.lower()  # type: ignore
        for member in cls:
            if member.value == value:
                return member
        return None


class AsyncTaskMetadata(BaseModel):
    """Metadata associated with an async task."""

    action: Annotated[str, Field(description="The path of the endpoint used to invoke the task.")]
    task_status: Annotated[TaskStatus, Field(description="The current status of the task.")]
    userdata: Annotated[Optional[str], Field(description="Additional metadata associated with the task.")] = None


class AsyncTaskMap(BaseModel):
    """Metadata describing a list of async tasks."""

    tasks: Annotated[dict[str, AsyncTaskMetadata], Field(description="A map of task ids to metadata for async tasks.")]
