"""
Module containing logic for persisting basic information about Fargate task execution.
This allows for creating task records, updating their status and payload, and returning
metadata about the task.

This is currently backed by Redis, with expectations that the data is ephemeral
and only lasts for a period of 24 hours.
"""

import time
import uuid
from typing import Any, Literal

from pydantic import UUID4, BaseModel

from product_staging.api import datasources

TASK_TTL = 3600 * 24  # 24 hours


# Define a type alias for task status
TaskStatus = Literal["in_progress", "success", "failure"]


class Task(BaseModel):
    """
    Represents a task with a unique key, status, and optional payload.
    """

    token: UUID4
    status: TaskStatus
    created_by: UUID4 | None = None
    created_at: float
    payload: dict[str, Any] | None = None

    @property
    def key(self) -> str:
        return token_key(self.token)


def token_key(token: UUID4) -> str:
    """
    Generate a Redis key to persist a Task, given a provided token UUID.

    Args:
        token (UUID4): the unique identifier for the Task
    Returns:
        str: the Redis key for the Task
    """
    return f"task_{str(token)}"


async def create_task(*, identity_uuid: UUID4 | None = None) -> Task:
    """
    Persist a new task marker to Redis, returning a unique task token.
    """
    redis = datasources.get_redis_client()
    task = Task(
        token=uuid.uuid4(),
        status="in_progress",
        created_by=identity_uuid,
        created_at=time.time(),
    )
    if redis:
        await redis.set(
            key=task.key,
            item=task.model_dump(mode="json"),
            ttl=TASK_TTL,
        )
    return task


async def get_task(*, token: UUID4, status: TaskStatus | None = None) -> Task | None:
    """
    Retrieve a task by its token from Redis.

    Args:
        token (UUID4): the unique identifier for the Task
        status (TaskStatus | None): optional status to filter the Task by

    Returns:
        Task | None: The task object if found and matches the
            status (if provided), else None
    """
    redis = datasources.get_redis_client()
    item = await redis.get(key=token_key(token)) if redis else None

    if item is None:
        return None

    task = Task.model_validate(item)

    if status is not None and task.status != status:
        return None

    return task


async def update_task(
    token: UUID4, status: TaskStatus, payload: dict[str, Any] | None = None
) -> Task:
    """
    Update an existing task.

    Args:
        token (UUID4): the unique identifier for the Task
        status (TaskStatus): the new status for the Task
        payload (dict[str, Any] | None): optional payload to associate with the task.

    Returns:
        Task: The updated task object.

    Raises:
        ValueError: If the task with the given token does not exist.
    """
    redis = datasources.get_redis_client()

    task = await get_task(token=token)

    if not task:
        raise ValueError(f"Task with token {str(token)} does not exist.")

    task.status = status
    task.payload = payload

    if redis:
        await redis.set(
            key=task.key,
            item=task.model_dump(mode="json"),
            ttl=TASK_TTL,
        )

    return task


async def complete_task(token: UUID4, payload: dict[str, Any]) -> Task:
    """
    Mark a task as completed and store the result in Redis.

    Args:
        token (UUID4): the unique identifier for the Task
        payload (dict[str, Any]): payload to associate with the completed Task

    Returns:
        Task: the updated Task object
    """
    return await update_task(token=token, status="success", payload=payload)


async def fail_task(token: UUID4, payload: dict[str, Any]) -> Task:
    """
    Mark a task as failed and store the result in Redis.

    Args:
        token (UUID4): the unique identifier for the Task
        payload (dict[str, Any]): payload to associate with the failed Task

    Returns:
        Task: the updated Task object
    """
    return await update_task(token=token, status="failure", payload=payload)