"""Tests for performance optimizations — batched queries and circuit breaker."""

import asyncio
import time

import pytest

from loom.graph import cache, store
from loom.graph.cache import _CircuitBreaker, _redis_cb
from loom.graph.task import Task, TaskStatus


@pytest.fixture
async def tasks_with_deps(pool, project):
    """Create multiple tasks with dependencies for batch testing."""
    t1 = Task(id="perf-001", project_id=project, title="Base 1", status=TaskStatus.DONE)
    t2 = Task(id="perf-002", project_id=project, title="Base 2", status=TaskStatus.DONE)
    t3 = Task(
        id="perf-003", project_id=project, title="Dep on 1",
        status=TaskStatus.PENDING, depends_on=["perf-001"],
    )
    t4 = Task(
        id="perf-004", project_id=project, title="Dep on 1,2",
        status=TaskStatus.PENDING, depends_on=["perf-001", "perf-002"],
    )
    t5 = Task(id="perf-005", project_id=project, title="No deps", status=TaskStatus.PENDING)
    for t in [t1, t2, t3, t4, t5]:
        await store.create_task(pool, t)
    return [t1, t2, t3, t4, t5]


async def test_batch_load_deps(pool, tasks_with_deps):
    result = await store._batch_load_deps(pool, ["perf-003", "perf-004", "perf-005"])
    assert set(result["perf-003"]) == {"perf-001"}
    assert set(result["perf-004"]) == {"perf-001", "perf-002"}
    assert result["perf-005"] == []


async def test_batch_load_deps_empty(pool):
    result = await store._batch_load_deps(pool, [])
    assert result == {}


async def test_get_ready_tasks_no_n_plus_one(pool, project, tasks_with_deps):
    """Ready tasks should use batched query (all done deps means tasks are ready)."""
    ready = await store.get_ready_tasks(pool, project)
    ready_ids = {t.id for t in ready}
    assert "perf-003" in ready_ids  # its dep (perf-001) is done
    assert "perf-004" in ready_ids  # its deps (perf-001, perf-002) are done
    assert "perf-005" in ready_ids  # no deps


async def test_get_expired_claims_no_n_plus_one(pool, project):
    """Expired claims should use batched dep loading."""
    t = Task(id="exp-001", project_id=project, title="Expired", status=TaskStatus.PENDING)
    await store.create_task(pool, t)
    await store.claim_task(pool, "exp-001", "agent-1", ttl_seconds=1)
    # Manually set expiry to past
    async with pool.acquire() as conn:
        await conn.execute(
            "UPDATE tasks SET claim_expires_at = NOW() - interval '1 hour' WHERE id = $1",
            "exp-001",
        )
    expired = await store.get_expired_claims(pool, project)
    assert len(expired) == 1
    assert expired[0].id == "exp-001"


async def test_circuit_breaker_opens_on_failures():
    cb = _CircuitBreaker(failure_threshold=3, cooldown_seconds=0.1)
    assert cb.state == "closed"
    for _ in range(3):
        cb.record_failure()
    assert cb.state == "open"
    assert not cb.should_allow_request


async def test_circuit_breaker_half_open_recovery():
    cb = _CircuitBreaker(failure_threshold=2, cooldown_seconds=0.05)
    cb.record_failure()
    cb.record_failure()
    assert cb.state == "open"
    await asyncio.sleep(0.1)  # wait for cooldown
    assert cb.state == "half-open"
    assert cb.should_allow_request


async def test_circuit_breaker_closes_on_success():
    cb = _CircuitBreaker(failure_threshold=2, cooldown_seconds=0.05)
    cb.record_failure()
    cb.record_failure()
    assert cb.state == "open"
    await asyncio.sleep(0.1)
    assert cb.state == "half-open"
    cb.record_success()
    assert cb.state == "closed"


async def test_cache_fallback_on_circuit_open(pool, redis_conn, project):
    """When circuit breaker is open, cache should fall back to Postgres."""
    t = Task(id="cb-001", project_id=project, title="CB Test", status=TaskStatus.PENDING)
    await store.create_task(pool, t)
    # Force circuit breaker open
    old_state = _redis_cb._state
    old_failures = _redis_cb._consecutive_failures
    try:
        _redis_cb._state = "open"
        _redis_cb._last_failure_time = time.monotonic()
        # Should still get the task via Postgres fallback
        task = await cache.get_task(redis_conn, pool, project, "cb-001")
        assert task.id == "cb-001"
    finally:
        _redis_cb._state = old_state
        _redis_cb._consecutive_failures = old_failures
