"""Tests for recursion guard async/thread safety (#1230).

Verifies that concurrent calls to stamp, evaluate, clear, and record
telemetry don't produce races, leaked grants, or corrupted state.
"""

from __future__ import annotations

import asyncio
import threading

import pytest

from cortex.hippo.tools.recursion_guard import (
    EXPANSION_DELEGATION_DEPTH_CAP,
    clear_delegated_expansion_context,
    create_expansion_request_id,
    evaluate_expansion_recursion_guard,
    get_delegated_context_for_tests,
    get_telemetry_snapshot_for_tests,
    record_expansion_delegation_telemetry,
    reset_for_tests,
    stamp_delegated_expansion_context,
)


@pytest.fixture(autouse=True)
def clean_state():
    """Reset all in-memory state before each test."""
    reset_for_tests()
    yield
    reset_for_tests()


class TestConcurrentStampAndEvaluate:
    """Verify no races when multiple coroutines stamp/evaluate concurrently."""

    @pytest.mark.asyncio
    async def test_concurrent_stamp_different_sessions(self):
        """Stamping different sessions concurrently should not interfere."""
        async def stamp_and_check(session_key: str, depth: int):
            stamp_delegated_expansion_context(
                session_key=session_key,
                request_id=f"req-{session_key}",
                expansion_depth=depth,
                origin_session_key="main",
                stamped_by="test",
            )
            ctx = get_delegated_context_for_tests(session_key)
            assert ctx is not None
            assert ctx.request_id == f"req-{session_key}"
            assert ctx.expansion_depth == depth

        # Run 20 concurrent stamps
        tasks = [
            stamp_and_check(f"session-{i}", i % 3)
            for i in range(20)
        ]
        await asyncio.gather(*tasks)

        # Verify all 20 contexts exist
        for i in range(20):
            ctx = get_delegated_context_for_tests(f"session-{i}")
            assert ctx is not None
            assert ctx.request_id == f"req-session-{i}"

    @pytest.mark.asyncio
    async def test_concurrent_evaluate_same_session(self):
        """Multiple concurrent evaluations on the same session should be consistent."""
        stamp_delegated_expansion_context(
            session_key="shared-session",
            request_id="req-shared",
            expansion_depth=EXPANSION_DELEGATION_DEPTH_CAP,
            origin_session_key="main",
            stamped_by="test",
        )

        results = []

        async def evaluate():
            decision = evaluate_expansion_recursion_guard(
                session_key="shared-session",
                request_id=create_expansion_request_id(),
            )
            results.append(decision)

        # 10 concurrent evaluations — all should be blocked
        await asyncio.gather(*[evaluate() for _ in range(10)])

        assert len(results) == 10
        assert all(d.blocked is True for d in results)
        # First one gets depth_cap, rest may get depth_cap or idempotent_reentry
        # depending on request_id uniqueness
        assert all(d.reason in ("depth_cap", "idempotent_reentry") for d in results)

    @pytest.mark.asyncio
    async def test_concurrent_stamp_and_clear(self):
        """Stamping then clearing concurrently should not leak contexts."""

        async def stamp_use_clear(i: int):
            session = f"ephemeral-{i}"
            stamp_delegated_expansion_context(
                session_key=session,
                request_id=f"req-{i}",
                expansion_depth=0,
                origin_session_key="main",
                stamped_by="test",
            )

            # Evaluate (should be allowed at depth 0)
            decision = evaluate_expansion_recursion_guard(
                session_key=session,
                request_id=f"req-{i}",
            )
            assert decision.blocked is False

            # Clear
            clear_delegated_expansion_context(session)

        await asyncio.gather(*[stamp_use_clear(i) for i in range(20)])

        # All contexts should be cleared — no leaked grants
        for i in range(20):
            assert get_delegated_context_for_tests(f"ephemeral-{i}") is None

    @pytest.mark.asyncio
    async def test_concurrent_expand_and_expand_query_simulation(self):
        """Simulate concurrent expand + expand_query tool calls.

        expand_query stamps a child context, evaluates, then clears.
        expand just evaluates (no stamp needed for direct expansion).
        Both running concurrently should not interfere.
        """
        errors = []

        async def simulate_expand_query(i: int):
            """Simulates what lcm_expand_query does with the guard."""
            child_key = f"child-eq-{i}"
            try:
                stamp_delegated_expansion_context(
                    session_key=child_key,
                    request_id=f"req-eq-{i}",
                    expansion_depth=1,
                    origin_session_key="main",
                    stamped_by="expand_query",
                )

                # Child evaluates — should be blocked at depth 1 (== cap)
                decision = evaluate_expansion_recursion_guard(
                    session_key=child_key,
                    request_id=f"req-eq-{i}",
                )
                assert decision.blocked is True

                record_expansion_delegation_telemetry(
                    component="test",
                    event="block",
                    request_id=f"req-eq-{i}",
                    session_key=child_key,
                    expansion_depth=1,
                    origin_session_key="main",
                    reason="depth_cap",
                )
            except Exception as e:
                errors.append(e)
            finally:
                clear_delegated_expansion_context(child_key)

        async def simulate_expand(i: int):
            """Simulates what lcm_expand does — just evaluates from main context."""
            try:
                decision = evaluate_expansion_recursion_guard(
                    session_key=f"expand-session-{i}",
                    request_id=f"req-exp-{i}",
                )
                # No stamped context → should be allowed
                assert decision.blocked is False
            except Exception as e:
                errors.append(e)

        # Run 10 expand_query + 10 expand concurrently
        tasks = []
        for i in range(10):
            tasks.append(simulate_expand_query(i))
            tasks.append(simulate_expand(i))

        await asyncio.gather(*tasks)

        assert errors == [], f"Concurrent errors: {errors}"

        # Verify all child contexts were cleaned up
        for i in range(10):
            assert get_delegated_context_for_tests(f"child-eq-{i}") is None


class TestConcurrentTelemetry:
    """Verify telemetry counters are consistent under concurrent access."""

    @pytest.mark.asyncio
    async def test_concurrent_telemetry_increments(self):
        """100 concurrent telemetry records should produce exact counter totals."""

        async def record(i: int):
            event = ["start", "block", "timeout", "success"][i % 4]
            record_expansion_delegation_telemetry(
                component="test",
                event=event,
                request_id=f"req-{i}",
                session_key=f"s-{i}",
                expansion_depth=0,
                origin_session_key="main",
            )

        await asyncio.gather(*[record(i) for i in range(100)])

        snap = get_telemetry_snapshot_for_tests()
        assert snap["start"] == 25
        assert snap["block"] == 25
        assert snap["timeout"] == 25
        assert snap["success"] == 25


class TestThreadSafety:
    """Verify thread-safety of the lock-protected operations."""

    def test_concurrent_threads_stamp_and_clear(self):
        """Multiple threads stamping and clearing should not corrupt state."""
        errors = []

        def worker(i: int):
            try:
                session = f"thread-{i}"
                stamp_delegated_expansion_context(
                    session_key=session,
                    request_id=f"req-{i}",
                    expansion_depth=0,
                    origin_session_key="main",
                    stamped_by="thread-test",
                )

                decision = evaluate_expansion_recursion_guard(
                    session_key=session,
                    request_id=f"req-{i}",
                )
                assert decision.blocked is False

                clear_delegated_expansion_context(session)
                assert get_delegated_context_for_tests(session) is None
            except Exception as e:
                errors.append(e)

        threads = [threading.Thread(target=worker, args=(i,)) for i in range(20)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert errors == [], f"Thread errors: {errors}"
