"""Unit and integration tests for loom.orchestration.retry."""

from __future__ import annotations

from datetime import datetime, timezone

import pytest

from loom.bus import channels
from loom.config import OrchestrationConfig
from loom.graph import store
from loom.graph.task import Task, TaskStatus
from loom.ids import task_id as gen_task_id
from loom.orchestration.retry import (
    compute_retry_delay,
    process_failed_task,
    sweep_retryable_tasks,
)


# ---------------------------------------------------------------------------
# Test 1: compute_retry_delay returns value in expected range
# ---------------------------------------------------------------------------


def test_compute_retry_delay_in_expected_range():
    """base * 2^n +/- 20% jitter for several retry counts."""
    for retry_count in range(5):
        base = 60
        expected_center = base * (2 ** retry_count)
        low = int(expected_center * 0.8) - 1
        high = int(expected_center * 1.2) + 1

        for _ in range(100):
            delay = compute_retry_delay(retry_count, base_delay=base, max_delay=100000)
            assert low <= delay <= high, (
                f"retry_count={retry_count}: delay {delay} outside [{low}, {high}]"
            )


# ---------------------------------------------------------------------------
# Test 2: compute_retry_delay respects max_delay cap
# ---------------------------------------------------------------------------


def test_compute_retry_delay_respects_max_delay():
    """When base * 2^n exceeds max_delay, delay is capped around max_delay +/- 20%."""
    max_delay = 500
    # retry_count=10 => 60 * 1024 = 61440, well above cap
    low = int(max_delay * 0.8) - 1
    high = int(max_delay * 1.2) + 1

    for _ in range(100):
        delay = compute_retry_delay(10, base_delay=60, max_delay=max_delay)
        assert low <= delay <= high, f"Capped delay {delay} outside [{low}, {high}]"


# ---------------------------------------------------------------------------
# Test 3: compute_retry_delay at retry_count=0 returns ~base_delay
# ---------------------------------------------------------------------------


def test_compute_retry_delay_at_zero():
    """At retry_count=0, delay centers on base_delay (2^0 = 1)."""
    base = 100
    samples = [compute_retry_delay(0, base_delay=base, max_delay=10000) for _ in range(200)]
    avg = sum(samples) / len(samples)

    # Average should be close to base (jitter is symmetric)
    assert abs(avg - base) < base * 0.1, f"Average {avg} too far from base {base}"

    # All samples within 20% jitter
    for s in samples:
        assert 80 - 1 <= s <= 120 + 1, f"Sample {s} outside [79, 121]"


# ---------------------------------------------------------------------------
# Test 4: process_failed_task retries when under max_retries
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_process_failed_task_retries_under_max(pool, redis_conn, project):
    """Task with retry_count < max_retries gets retry scheduled."""
    task = Task(id=gen_task_id(), project_id=project, title="Retry me")
    task = await store.create_task(pool, task)
    task = await store.claim_task(pool, task.id, "agent-1")
    task = await store.fail_task(pool, task.id, "transient error")

    config = OrchestrationConfig(
        max_retries=3,
        retry_base_delay_seconds=10,
        retry_max_delay_seconds=600,
    )

    updated = await process_failed_task(pool, redis_conn, task, config, project)

    # retry_count incremented from 0 to 1
    assert updated.retry_count == 1
    # retry_after set to a future time
    assert updated.retry_after is not None
    assert updated.retry_after > datetime.now(timezone.utc)
    # Not dead-lettered
    assert updated.dead_letter is False

    # Verify persisted in Postgres
    pg_task = await store.get_task(pool, task.id)
    assert pg_task.retry_count == 1
    assert pg_task.retry_after is not None


# ---------------------------------------------------------------------------
# Test 5: process_failed_task dead-letters when at max_retries
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_process_failed_task_dead_letters_at_max(pool, redis_conn, project):
    """Task with retry_count >= max_retries goes to dead letter."""
    task = Task(id=gen_task_id(), project_id=project, title="Dead letter me")
    task = await store.create_task(pool, task)
    task = await store.claim_task(pool, task.id, "agent-1")
    task = await store.fail_task(pool, task.id, "persistent error")

    # Bump retry_count to max_retries via SQL
    max_retries = 3
    async with pool.acquire() as conn:
        await conn.execute(
            "UPDATE tasks SET retry_count = $1 WHERE id = $2",
            max_retries, task.id,
        )
    task = await store.get_task(pool, task.id)
    assert task.retry_count == max_retries

    config = OrchestrationConfig(
        max_retries=max_retries,
        retry_base_delay_seconds=10,
        retry_max_delay_seconds=600,
    )

    updated = await process_failed_task(pool, redis_conn, task, config, project)

    assert updated.dead_letter is True
    assert updated.dead_letter_reason == "Max retries exceeded"
    assert updated.retry_after is None

    # Verify in dead letter Redis set
    dlq_key = channels.dead_letter_set(project)
    members = await redis_conn.smembers(dlq_key)
    assert task.id in members

    # Verify in Postgres
    pg_task = await store.get_task(pool, task.id)
    assert pg_task.dead_letter is True


# ---------------------------------------------------------------------------
# Test 6: sweep_retryable_tasks resets eligible tasks to pending
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_sweep_retryable_tasks_resets_to_pending(pool, redis_conn, project):
    """Tasks with retry_after in the past get swept back to pending."""
    task = Task(id=gen_task_id(), project_id=project, title="Sweepable task")
    task = await store.create_task(pool, task)
    task = await store.claim_task(pool, task.id, "agent-1")
    task = await store.fail_task(pool, task.id, "error")

    # Set retry_after to the past so it's eligible for sweep
    async with pool.acquire() as conn:
        await conn.execute(
            "UPDATE tasks SET retry_after = NOW() - interval '1 second', "
            "status = 'failed', dead_letter = FALSE WHERE id = $1",
            task.id,
        )

    count = await sweep_retryable_tasks(pool, redis_conn, project)
    assert count == 1

    # Verify task is now pending
    pg_task = await store.get_task(pool, task.id)
    assert pg_task.status == TaskStatus.PENDING

    # Verify task is in the ready queue in Redis
    queue_key = channels.ready_queue(project)
    members = await redis_conn.zrange(queue_key, 0, -1)
    assert task.id in members
