"""Tests for auto-heartbeat background claim TTL extension."""

from __future__ import annotations

import asyncio

import pytest

from loom.graph import cache, store
from loom.graph.task import Task, TaskStatus
from loom.ids import task_id as gen_task_id
from loom.orchestration.heartbeat import AutoHeartbeat


@pytest.fixture
async def heartbeat(pool, redis_conn):
    """Create an AutoHeartbeat with a short interval for testing."""
    hb = AutoHeartbeat(pool, redis_conn, interval=1)
    yield hb
    await hb.shutdown()


@pytest.fixture
async def claimed_task(pool, redis_conn, project):
    """Create and claim a task, returning (task, ttl_seconds)."""
    ttl = 300
    task = Task(
        id=gen_task_id(), project_id=project, title="test task",
        status=TaskStatus.PENDING, priority="p1",
    )
    task = await store.create_task(pool, task)
    task = await store.claim_task(pool, task.id, "agent-1", ttl_seconds=ttl)
    await cache.sync_task(redis_conn, task)
    return task, ttl


@pytest.mark.asyncio
async def test_register_starts_loop(heartbeat):
    """Registering the first task starts the background loop."""
    assert heartbeat._task is None
    heartbeat.register("task-1", 300)
    assert heartbeat._task is not None
    assert not heartbeat._task.done()


@pytest.mark.asyncio
async def test_unregister_last_stops_loop(heartbeat):
    """Unregistering the last task cancels the background loop."""
    heartbeat.register("task-1", 300)
    loop_task = heartbeat._task
    heartbeat.unregister("task-1")
    assert heartbeat._task is None
    # Give the cancelled task a moment to finish
    await asyncio.sleep(0.05)
    assert loop_task.done()


@pytest.mark.asyncio
async def test_unregister_unknown_task_is_noop(heartbeat):
    """Unregistering an unknown task_id does nothing."""
    heartbeat.unregister("nonexistent")
    assert heartbeat._task is None
    assert len(heartbeat._claims) == 0


@pytest.mark.asyncio
async def test_loop_extends_claims(pool, redis_conn, project, claimed_task):
    """After one interval, claim_expires_at is extended."""
    task, ttl = claimed_task
    hb = AutoHeartbeat(pool, redis_conn, interval=1)
    try:
        # Record original expiry
        original = await store.get_task(pool, task.id)
        original_expires = original.claim_expires_at

        hb.register(task.id, ttl)
        # Wait for one loop cycle + buffer
        await asyncio.sleep(1.5)

        updated = await store.get_task(pool, task.id)
        assert updated.claim_expires_at > original_expires
    finally:
        await hb.shutdown()


@pytest.mark.asyncio
async def test_loop_survives_db_error(pool, redis_conn, project, claimed_task):
    """A transient DB error doesn't kill the loop."""
    task, ttl = claimed_task
    hb = AutoHeartbeat(pool, redis_conn, interval=1)
    try:
        # Register a non-existent task (will cause TaskNotFoundError)
        # alongside a valid task
        hb.register("nonexistent-task", 300)
        hb.register(task.id, ttl)
        # Wait for loop to run
        await asyncio.sleep(1.5)
        # Loop should still be alive
        assert hb._task is not None
        assert not hb._task.done()
        # Valid task should still have been extended
        updated = await store.get_task(pool, task.id)
        assert updated.claim_expires_at is not None
    finally:
        await hb.shutdown()


@pytest.mark.asyncio
async def test_shutdown_cancels_loop(heartbeat):
    """Explicit shutdown cleanly cancels the loop and clears state."""
    heartbeat.register("task-1", 300)
    heartbeat.register("task-2", 300)
    assert heartbeat._task is not None

    await heartbeat.shutdown()

    assert heartbeat._task is None
    assert len(heartbeat._claims) == 0


@pytest.mark.asyncio
async def test_register_multiple_tasks(heartbeat):
    """Multiple tasks can be tracked simultaneously."""
    heartbeat.register("task-1", 300)
    heartbeat.register("task-2", 600)
    assert len(heartbeat._claims) == 2
    assert heartbeat._claims["task-1"] == 300
    assert heartbeat._claims["task-2"] == 600


@pytest.mark.asyncio
async def test_unregister_one_keeps_loop(heartbeat):
    """Unregistering one of multiple tasks keeps the loop running."""
    heartbeat.register("task-1", 300)
    heartbeat.register("task-2", 300)
    heartbeat.unregister("task-1")
    assert heartbeat._task is not None
    assert not heartbeat._task.done()
    assert "task-1" not in heartbeat._claims
    assert "task-2" in heartbeat._claims
