"""Retry logic with exponential backoff and dead letter queue."""

from __future__ import annotations

import random
from datetime import datetime, timedelta, timezone

import asyncpg
from redis.asyncio import Redis

from loom.bus import channels
from loom.bus.events import EventType
from loom.bus.publisher import publish_event
from loom.config import OrchestrationConfig
from loom.graph import cache, store
from loom.graph.task import Task, TaskStatus


def compute_retry_delay(
    retry_count: int,
    base_delay: int = 60,
    max_delay: int = 3600,
) -> int:
    """Compute retry delay with exponential backoff and +/-20% jitter.

    Args:
        retry_count: Current retry attempt (0-based, so first retry uses 2^0)
        base_delay: Base delay in seconds
        max_delay: Maximum delay cap in seconds

    Returns:
        Delay in seconds
    """
    delay = min(base_delay * (2 ** retry_count), max_delay)
    jitter = delay * 0.2
    return int(delay + random.uniform(-jitter, jitter))


async def process_failed_task(
    pool: asyncpg.Pool,
    redis: Redis,
    task: Task,
    config: OrchestrationConfig,
    project_id: str,
) -> Task:
    """Handle a failed task: retry with backoff or move to dead letter.

    If retry_count < max_retries: schedule for retry
    Otherwise: move to dead letter queue

    Returns the updated task.
    """
    max_retries = config.max_retries

    if task.retry_count < max_retries:
        # Schedule retry
        delay = compute_retry_delay(
            task.retry_count,
            config.retry_base_delay_seconds,
            config.retry_max_delay_seconds,
        )
        retry_after = datetime.now(timezone.utc) + timedelta(seconds=delay)
        updated = await store.increment_retry(pool, task.id, retry_after)
        await cache.sync_task(redis, updated)
        await publish_event(redis, project_id, EventType.TASK_RETRIED, task.id)
        return updated
    else:
        # Move to dead letter
        updated = await store.mark_dead_letter(pool, task.id, "Max retries exceeded")
        await cache.sync_task(redis, updated)
        await redis.sadd(channels.dead_letter_set(project_id), task.id)
        await publish_event(redis, project_id, EventType.TASK_DEAD_LETTER, task.id)
        return updated


async def sweep_retryable_tasks(
    pool: asyncpg.Pool,
    redis: Redis,
    project_id: str,
) -> int:
    """Find failed tasks ready for retry and reset them to pending.

    Queries for tasks where:
    - status = 'failed'
    - dead_letter = FALSE
    - retry_after <= NOW()

    For each: reset to pending, clear assignee, sync cache, add to ready queue.

    Returns count of tasks reset.
    """
    retryable = await store.get_retryable_tasks(pool, project_id)
    count = 0
    for task in retryable:
        # Reset to pending via set_task_status
        updated = await store.set_task_status(pool, task.id, TaskStatus.PENDING)
        await cache.sync_task(redis, updated)
        await cache.add_to_ready_queue(redis, updated)
        count += 1
    return count
