"""Shared activity timeout policy for Temporal workflows."""

from datetime import timedelta

from app.infra.temporal.task_queues import (
    IMAGE_PIPELINE_IMAGE_GEN_TASK_QUEUE,
    IMAGE_PIPELINE_POST_PROCESS_TASK_QUEUE,
)

DEFAULT_ACTIVITY_TIMEOUT = timedelta(seconds=30)
LONG_RUNNING_ACTIVITY_TIMEOUT = timedelta(minutes=5)

LONG_RUNNING_TASK_QUEUES = {
    IMAGE_PIPELINE_IMAGE_GEN_TASK_QUEUE,
    IMAGE_PIPELINE_POST_PROCESS_TASK_QUEUE,
}


def activity_timeout_for_task_queue(task_queue: str) -> timedelta:
    """Return the timeout that matches the workload behind the given task queue."""

    if task_queue in LONG_RUNNING_TASK_QUEUES:
        return LONG_RUNNING_ACTIVITY_TIMEOUT
    return DEFAULT_ACTIVITY_TIMEOUT