"""Integration tests for Stream C: Retry logic + Dead Letter Queue."""

from __future__ import annotations

from datetime import datetime, timedelta, timezone

import pytest

from loom.bus import channels
from loom.config import OrchestrationConfig
from loom.graph import cache, store
from loom.graph.project import create_project
from loom.graph.task import Priority, 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,
)


@pytest.fixture
async def project(pool):
    """Create a test project and return its ID as a string."""
    record = await create_project(pool, "retry-test-project", "Retry DLQ tests")
    return str(record["id"])


# ---------------------------------------------------------------------------
# Test 1: Exponential backoff
# ---------------------------------------------------------------------------


def test_compute_retry_delay_exponential():
    """Delay doubles with each retry, capped at max."""
    # Use base_delay=100, max_delay=10000, no jitter variance check here
    # Just check the general trend: each retry roughly doubles
    delays = []
    for i in range(6):
        # Run many samples and take the median to smooth jitter
        samples = [compute_retry_delay(i, base_delay=100, max_delay=10000) for _ in range(100)]
        avg = sum(samples) / len(samples)
        delays.append(avg)

    # delay[0] ~ 100, delay[1] ~ 200, delay[2] ~ 400, delay[3] ~ 800, etc.
    # Each should roughly double (within 30% due to jitter averaging)
    for i in range(1, len(delays)):
        expected = min(100 * (2 ** i), 10000)
        # Average should be close to the base (jitter is symmetric around 0)
        assert abs(delays[i] - expected) < expected * 0.3, (
            f"retry {i}: expected ~{expected}, got {delays[i]}"
        )

    # Verify cap: retry 7 would be 100 * 128 = 12800, capped at 10000
    capped_samples = [compute_retry_delay(7, base_delay=100, max_delay=10000) for _ in range(50)]
    for s in capped_samples:
        assert s <= 12000, f"Delay {s} should be capped near 10000"
        assert s >= 8000, f"Delay {s} should be near 10000 (with jitter)"


# ---------------------------------------------------------------------------
# Test 2: Jitter range
# ---------------------------------------------------------------------------


def test_compute_retry_delay_jitter():
    """Delay has +/-20% jitter: run multiple times, verify range."""
    base = 1000
    max_d = 100000  # won't hit cap
    retry = 0  # base * 2^0 = 1000

    samples = [compute_retry_delay(retry, base, max_d) for _ in range(500)]

    # Expected range: 1000 +/- 20% = [800, 1200]
    min_sample = min(samples)
    max_sample = max(samples)

    # With 500 samples, we should see values across the range
    assert min_sample >= 800 - 1, f"Min {min_sample} should be >= ~800"
    assert max_sample <= 1200 + 1, f"Max {max_sample} should be <= ~1200"
    # Verify we actually get spread (not all the same)
    assert max_sample - min_sample > 100, "Jitter should produce spread"


# ---------------------------------------------------------------------------
# Test 3: process_failed_task schedules retry
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_process_failed_task_schedules_retry(pool, redis_conn, project):
    """Failed task with retries remaining gets retry_after set and retry_count incremented."""
    task = Task(id=gen_task_id(), project_id=project, title="Retry candidate")
    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)

    # Should have incremented retry_count
    assert updated.retry_count == 1
    # Should have set retry_after to a future time
    assert updated.retry_after is not None
    assert updated.retry_after > datetime.now(timezone.utc)
    # Should NOT be dead lettered
    assert updated.dead_letter is False
    # last_failed_at should be set
    assert updated.last_failed_at is not None

    # Verify 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 4: process_failed_task dead letters
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_process_failed_task_dead_letter(pool, redis_conn, project):
    """Failed task exceeding max_retries goes to dead letter."""
    task = Task(id=gen_task_id(), project_id=project, title="Dead letter candidate")
    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")

    # Set retry_count to max_retries so next failure triggers dead letter
    async with pool.acquire() as conn:
        await conn.execute(
            "UPDATE tasks SET retry_count = $1 WHERE id = $2",
            3, task.id,
        )
    # Re-fetch to get updated retry_count
    task = await store.get_task(pool, task.id)
    assert task.retry_count == 3

    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)

    # Should be dead lettered
    assert updated.dead_letter is True
    assert updated.dead_letter_reason == "Max retries exceeded"
    # retry_after should be cleared
    assert updated.retry_after is None

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


# ---------------------------------------------------------------------------
# Test 5: sweep_retryable_tasks
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_sweep_retryable_tasks(pool, redis_conn, project):
    """Retryable tasks (retry_after in the past) get reset to pending."""
    # Create and fail a task, set retry_after to the past
    task = Task(id=gen_task_id(), project_id=project, title="Sweepable")
    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")

    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,
        )

    # Sweep should find and reset this task
    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


# ---------------------------------------------------------------------------
# Test 6: loom_fail integrates retry
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_loom_fail_integrates_retry(pool, redis_conn, project):
    """Calling fail_task followed by process_failed_task schedules a retry."""
    task = Task(id=gen_task_id(), project_id=project, title="Fail and retry")
    task = await store.create_task(pool, task)
    task = await store.claim_task(pool, task.id, "agent-1")

    # Simulate what loom_fail does: fail_task then process_failed_task
    task = await store.fail_task(pool, task.id, "timeout")
    await cache.sync_task(redis_conn, task)

    config = OrchestrationConfig(max_retries=3, retry_base_delay_seconds=5, retry_max_delay_seconds=300)
    updated = await process_failed_task(pool, redis_conn, task, config, project)

    # First failure: retry_count should be 1
    assert updated.retry_count == 1
    assert updated.retry_after is not None
    assert updated.dead_letter is False

    # Simulate second failure (need to re-claim first by resetting to pending)
    async with pool.acquire() as conn:
        await conn.execute(
            "UPDATE tasks SET status = 'pending', assignee = NULL WHERE id = $1",
            updated.id,
        )
    task2 = await store.claim_task(pool, updated.id, "agent-2")
    task2 = await store.fail_task(pool, task2.id, "timeout again")
    updated2 = await process_failed_task(pool, redis_conn, task2, config, project)

    assert updated2.retry_count == 2
    assert updated2.dead_letter is False

    # Third failure
    async with pool.acquire() as conn:
        await conn.execute(
            "UPDATE tasks SET status = 'pending', assignee = NULL WHERE id = $1",
            updated2.id,
        )
    task3 = await store.claim_task(pool, updated2.id, "agent-3")
    task3 = await store.fail_task(pool, task3.id, "timeout yet again")
    updated3 = await process_failed_task(pool, redis_conn, task3, config, project)

    assert updated3.retry_count == 3
    assert updated3.dead_letter is False

    # Fourth failure: should dead letter now (retry_count=3 == max_retries=3)
    async with pool.acquire() as conn:
        await conn.execute(
            "UPDATE tasks SET status = 'pending', assignee = NULL WHERE id = $1",
            updated3.id,
        )
    task4 = await store.claim_task(pool, updated3.id, "agent-4")
    task4 = await store.fail_task(pool, task4.id, "final failure")
    # Re-fetch to get current retry_count from DB
    task4 = await store.get_task(pool, task4.id)
    updated4 = await process_failed_task(pool, redis_conn, task4, config, project)

    assert updated4.dead_letter is True
    assert updated4.dead_letter_reason == "Max retries exceeded"


# ---------------------------------------------------------------------------
# Test 7: Dead letter in Redis
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_dead_letter_in_redis(pool, redis_conn, project):
    """Dead letter task ID appears in the dead letter Redis set."""
    task = Task(id=gen_task_id(), project_id=project, title="DLQ Redis test")
    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_count at max so it dead letters
    async with pool.acquire() as conn:
        await conn.execute(
            "UPDATE tasks SET retry_count = $1 WHERE id = $2",
            3, task.id,
        )
    task = await store.get_task(pool, task.id)

    config = OrchestrationConfig(max_retries=3)
    await process_failed_task(pool, redis_conn, task, config, project)

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

    # Verify NOT in the dead letter set for a different project
    other_key = channels.dead_letter_set("other-project")
    other_members = await redis_conn.smembers(other_key)
    assert task.id not in other_members
