"""Integration tests for Claim TTL + Heartbeat system (Phase 3 Stream B)."""

from __future__ import annotations

from datetime import datetime, timedelta, timezone

import pytest

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
from loom.orchestration.sweeper import sweep_expired_claims


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


# ── Test 1: Claiming a task with TTL sets claim_expires_at ──────────────


async def test_claim_with_ttl(pool, redis_conn, project):
    """Claiming a task with ttl_seconds sets claim_expires_at in the future."""
    t = Task(
        id=task_id(), project_id=project, title="Task with TTL",
        priority=Priority.P1,
    )
    created = await store.create_task(pool, t)

    before_claim = datetime.now(timezone.utc)
    claimed = await store.claim_task(pool, created.id, "agent-1", ttl_seconds=600)
    after_claim = datetime.now(timezone.utc)

    assert claimed.status == TaskStatus.CLAIMED
    assert claimed.assignee == "agent-1"
    assert claimed.claim_expires_at is not None

    # claim_expires_at should be approximately NOW + 600 seconds
    expected_min = before_claim + timedelta(seconds=599)
    expected_max = after_claim + timedelta(seconds=601)
    assert expected_min <= claimed.claim_expires_at <= expected_max


# ── Test 2: Heartbeat extends TTL ──────────────────────────────────────


async def test_heartbeat_extends_ttl(pool, redis_conn, project):
    """set_claim_expiry() pushes claim_expires_at forward."""
    t = Task(
        id=task_id(), project_id=project, title="Heartbeat task",
        priority=Priority.P0,
    )
    created = await store.create_task(pool, t)
    claimed = await store.claim_task(pool, created.id, "agent-1", ttl_seconds=300)
    original_expiry = claimed.claim_expires_at
    assert original_expiry is not None

    # Heartbeat: extend the TTL by another 600 seconds from now
    new_expiry = datetime.now(timezone.utc) + timedelta(seconds=600)
    updated = await store.set_claim_expiry(pool, created.id, new_expiry)

    assert updated.claim_expires_at is not None
    assert updated.claim_expires_at > original_expiry

    # Verify it persists in Postgres
    from_pg = await store.get_task(pool, created.id)
    assert from_pg.claim_expires_at is not None
    # The stored value should match what we set (within a small margin)
    delta = abs((from_pg.claim_expires_at - new_expiry).total_seconds())
    assert delta < 2  # within 2 seconds tolerance


# ── Test 3: Expired claims are detected ────────────────────────────────


async def test_expired_claim_detected(pool, redis_conn, project):
    """get_expired_claims() finds tasks whose claim TTL has passed."""
    t = Task(
        id=task_id(), project_id=project, title="Expiring task",
        priority=Priority.P1,
    )
    created = await store.create_task(pool, t)
    claimed = await store.claim_task(pool, created.id, "agent-1", ttl_seconds=600)

    # Not expired yet
    expired = await store.get_expired_claims(pool, project)
    assert not any(e.id == created.id for e in expired)

    # Manually set claim_expires_at to the past
    async with pool.acquire() as conn:
        await conn.execute(
            "UPDATE tasks SET claim_expires_at = NOW() - interval '1 second' WHERE id = $1",
            created.id,
        )

    # Now it should be found
    expired = await store.get_expired_claims(pool, project)
    assert any(e.id == created.id for e in expired)


# ── Test 4: Releasing an expired claim ─────────────────────────────────


async def test_release_expired_claim(pool, redis_conn, project):
    """release_expired_claim() moves task back to pending, clears assignee/TTL."""
    t = Task(
        id=task_id(), project_id=project, title="Release me",
        priority=Priority.P0,
    )
    created = await store.create_task(pool, t)
    claimed = await store.claim_task(pool, created.id, "agent-stuck", ttl_seconds=600)
    assert claimed.status == TaskStatus.CLAIMED
    assert claimed.assignee == "agent-stuck"

    # Expire the claim
    async with pool.acquire() as conn:
        await conn.execute(
            "UPDATE tasks SET claim_expires_at = NOW() - interval '1 second' WHERE id = $1",
            created.id,
        )

    # Release it
    released = await store.release_expired_claim(pool, created.id)
    assert released.status == TaskStatus.PENDING
    assert released.assignee is None
    assert released.claim_expires_at is None
    assert released.claimed_at is None

    # Verify in Postgres
    from_pg = await store.get_task(pool, created.id)
    assert from_pg.status == TaskStatus.PENDING
    assert from_pg.assignee is None


# ── Test 5: Sweeper releases multiple expired claims ───────────────────


async def test_sweep_expired_claims(pool, redis_conn, project):
    """sweep_expired_claims() finds and releases all expired claims in one pass."""
    # Create and claim 3 tasks
    tasks = []
    for i in range(3):
        t = Task(
            id=task_id(), project_id=project, title=f"Sweep task {i}",
            priority=Priority.P1,
        )
        created = await store.create_task(pool, t)
        claimed = await store.claim_task(pool, created.id, f"agent-{i}", ttl_seconds=600)
        tasks.append(claimed)

    # Expire the first two, leave the third
    async with pool.acquire() as conn:
        await conn.execute(
            "UPDATE tasks SET claim_expires_at = NOW() - interval '1 second' WHERE id = $1",
            tasks[0].id,
        )
        await conn.execute(
            "UPDATE tasks SET claim_expires_at = NOW() - interval '1 second' WHERE id = $1",
            tasks[1].id,
        )

    # Sync to cache first
    for t in tasks:
        await cache.sync_task(redis_conn, t)

    # Run the sweeper
    released_count = await sweep_expired_claims(pool, redis_conn, project)
    assert released_count == 2

    # Verify the first two are pending again
    t0 = await store.get_task(pool, tasks[0].id)
    t1 = await store.get_task(pool, tasks[1].id)
    assert t0.status == TaskStatus.PENDING
    assert t1.status == TaskStatus.PENDING
    assert t0.assignee is None
    assert t1.assignee is None

    # The third should still be claimed
    t2 = await store.get_task(pool, tasks[2].id)
    assert t2.status == TaskStatus.CLAIMED
    assert t2.assignee == "agent-2"

    # The released tasks should be back in the ready queue
    ready = await cache.get_ready_tasks(redis_conn, pool, project)
    ready_ids = {r.id for r in ready}
    assert tasks[0].id in ready_ids
    assert tasks[1].id in ready_ids


# ── Test 6: Claim without TTL ──────────────────────────────────────────


async def test_claim_without_ttl(pool, redis_conn, project):
    """claim_task() with no ttl_seconds leaves claim_expires_at as NULL."""
    t = Task(
        id=task_id(), project_id=project, title="No TTL task",
        priority=Priority.P2,
    )
    created = await store.create_task(pool, t)
    claimed = await store.claim_task(pool, created.id, "agent-free")

    assert claimed.status == TaskStatus.CLAIMED
    assert claimed.assignee == "agent-free"
    assert claimed.claim_expires_at is None

    # Verify in Postgres too
    from_pg = await store.get_task(pool, created.id)
    assert from_pg.claim_expires_at is None
