"""Unit tests for autonomy_paused event + timeout/failure tracking (IQ-Boost U3, R10).

Verifies:
- Timeout triggers autonomy_paused with reason="timeout"
- Consecutive failures (3x) trigger autonomy_paused with reason="consecutive_failures"
- resume_handler returns True → counters reset, execution continues
- resume_handler returns False → execution stops (cancel)
- resume_handler=None → auto-resume (non-blocking, for tests/REST)
- Non-autonomy mode → no pause (gate is a no-op)
- Disabled thresholds (0) → no pause
- _track_tool_result_for_autonomy increments on error, resets on success
- U6: resume Future pre-registration race fix
"""

from __future__ import annotations

import asyncio
import time
from unittest.mock import AsyncMock, MagicMock

import pytest

from agentkit.core.react import ReActEngine, ReActEvent
from agentkit.server.config import DangerousToolsConfig


# ---------------------------------------------------------------------------
# Helper: simulates the caller pattern (detect → yield event → await resume)
# ---------------------------------------------------------------------------


async def _detect_and_await_pause(
    engine: ReActEngine,
    step: int,
    progress: dict,
    resume_handler=None,
) -> tuple[bool, list[ReActEvent]]:
    """Simulate the react.py caller pattern for autonomy pause.

    Mirrors the production code: detect (non-blocking) → build event →
    yield event → await resume (blocking). Returns (should_continue, events).
    """
    pause_info = engine._detect_autonomy_pause(step, progress)
    if pause_info is None:
        return True, []
    reason, token, event_data = pause_info
    event = ReActEvent(event_type="autonomy_paused", step=step, data=event_data)
    should_continue = await engine._await_autonomy_resume(token, reason, resume_handler)
    return should_continue, [event]


# ---------------------------------------------------------------------------
# _track_tool_result_for_autonomy
# ---------------------------------------------------------------------------


class TestTrackToolResult:
    """Failure tracking increments on error, resets on success."""

    def test_success_resets_counter(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(),
            autonomy_mode=True,
        )
        engine._consecutive_failures = 2
        engine._track_tool_result_for_autonomy({"output": "ok"})
        assert engine._consecutive_failures == 0

    def test_error_increments_counter(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(),
            autonomy_mode=True,
        )
        engine._consecutive_failures = 1
        engine._track_tool_result_for_autonomy({"error": "boom"})
        assert engine._consecutive_failures == 2

    def test_is_error_flag_increments(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(),
            autonomy_mode=True,
        )
        engine._track_tool_result_for_autonomy({"is_error": True, "output": ""})
        assert engine._consecutive_failures == 1

    def test_error_type_increments(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(),
            autonomy_mode=True,
        )
        engine._track_tool_result_for_autonomy({"error_type": "permission_denied"})
        assert engine._consecutive_failures == 1

    def test_non_dict_result_resets(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(),
            autonomy_mode=True,
        )
        engine._consecutive_failures = 2
        engine._track_tool_result_for_autonomy("plain string result")
        assert engine._consecutive_failures == 0

    def test_no_tracking_when_not_autonomy(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(),
            autonomy_mode=False,
        )
        engine._consecutive_failures = 0
        engine._track_tool_result_for_autonomy({"error": "boom"})
        assert engine._consecutive_failures == 0  # Unchanged


# ---------------------------------------------------------------------------
# _detect_autonomy_pause + _await_autonomy_resume — no pause conditions
# ---------------------------------------------------------------------------


class TestAutonomyPauseNoTrigger:
    """When conditions are not met, no pause is triggered (should_continue=True)."""

    @pytest.mark.asyncio
    async def test_no_autonomy_mode_no_pause(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(),
            autonomy_mode=False,
        )
        should_continue, events = await _detect_and_await_pause(
            engine, step=1, progress={}, resume_handler=None
        )
        assert should_continue is True
        assert events == []

    @pytest.mark.asyncio
    async def test_no_config_no_pause(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=None,
            autonomy_mode=True,
        )
        should_continue, events = await _detect_and_await_pause(
            engine, step=1, progress={}, resume_handler=None
        )
        assert should_continue is True
        assert events == []

    @pytest.mark.asyncio
    async def test_below_thresholds_no_pause(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(
                autonomy_timeout_minutes=30,
                max_consecutive_failures=3,
            ),
            autonomy_mode=True,
        )
        engine._autonomy_started_at = time.time()  # Just started
        engine._consecutive_failures = 1  # Below threshold
        should_continue, events = await _detect_and_await_pause(
            engine, step=1, progress={}, resume_handler=None
        )
        assert should_continue is True
        assert events == []

    @pytest.mark.asyncio
    async def test_disabled_thresholds_no_pause(self):
        """0 = disabled for both thresholds."""
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(
                autonomy_timeout_minutes=0,
                max_consecutive_failures=0,
            ),
            autonomy_mode=True,
        )
        # Even with expired time + many failures, disabled = no pause.
        engine._autonomy_started_at = time.time() - 999999
        engine._consecutive_failures = 99
        should_continue, events = await _detect_and_await_pause(
            engine, step=1, progress={}, resume_handler=None
        )
        assert should_continue is True
        assert events == []


# ---------------------------------------------------------------------------
# _detect_autonomy_pause + _await_autonomy_resume — timeout trigger
# ---------------------------------------------------------------------------


class TestAutonomyPauseTimeout:
    """Timeout triggers autonomy_paused with reason='timeout'."""

    @pytest.mark.asyncio
    async def test_timeout_triggers_pause_auto_resume(self):
        """resume_handler=None → auto-resume (non-blocking)."""
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(
                autonomy_timeout_minutes=1,  # 1 minute
                max_consecutive_failures=99,  # Disabled for this test
            ),
            autonomy_mode=True,
        )
        # Simulate started 2 minutes ago (exceeds 1-min timeout).
        engine._autonomy_started_at = time.time() - 120

        should_continue, events = await _detect_and_await_pause(
            engine, step=5, progress={"step": 5}, resume_handler=None
        )

        # Auto-resume → should_continue=True
        assert should_continue is True
        assert len(events) == 1
        assert events[0].event_type == "autonomy_paused"
        assert events[0].data["reason"] == "timeout"
        assert "resume_token" in events[0].data
        assert events[0].data["elapsed_seconds"] >= 120
        # Counters reset after auto-resume
        assert engine._consecutive_failures == 0
        assert engine._autonomy_started_at > time.time() - 1  # Fresh timer

    @pytest.mark.asyncio
    async def test_timeout_triggers_pause_user_resumes(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(
                autonomy_timeout_minutes=1,
                max_consecutive_failures=99,
            ),
            autonomy_mode=True,
        )
        engine._autonomy_started_at = time.time() - 120

        resume_handler = AsyncMock(return_value=True)
        should_continue, events = await _detect_and_await_pause(
            engine, step=3, progress={}, resume_handler=resume_handler
        )

        assert should_continue is True
        assert events[0].data["reason"] == "timeout"
        resume_handler.assert_awaited_once()
        # Counters reset
        assert engine._consecutive_failures == 0

    @pytest.mark.asyncio
    async def test_timeout_triggers_pause_user_cancels(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(
                autonomy_timeout_minutes=1,
                max_consecutive_failures=99,
            ),
            autonomy_mode=True,
        )
        engine._autonomy_started_at = time.time() - 120

        resume_handler = AsyncMock(return_value=False)
        should_continue, events = await _detect_and_await_pause(
            engine, step=3, progress={}, resume_handler=resume_handler
        )

        assert should_continue is False  # Cancelled
        assert events[0].data["reason"] == "timeout"


# ---------------------------------------------------------------------------
# _detect_autonomy_pause + _await_autonomy_resume — consecutive failures trigger
# ---------------------------------------------------------------------------


class TestAutonomyPauseConsecutiveFailures:
    """3 consecutive failures trigger autonomy_paused."""

    @pytest.mark.asyncio
    async def test_consecutive_failures_triggers_pause(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(
                autonomy_timeout_minutes=0,  # Disabled
                max_consecutive_failures=3,
            ),
            autonomy_mode=True,
        )
        engine._autonomy_started_at = time.time()
        engine._consecutive_failures = 3  # At threshold

        resume_handler = AsyncMock(return_value=True)
        should_continue, events = await _detect_and_await_pause(
            engine, step=5, progress={"step": 5}, resume_handler=resume_handler
        )

        assert should_continue is True  # Resumed
        assert len(events) == 1
        assert events[0].event_type == "autonomy_paused"
        assert events[0].data["reason"] == "consecutive_failures"
        assert events[0].data["consecutive_failures"] == 3
        # Reset after resume
        assert engine._consecutive_failures == 0

    @pytest.mark.asyncio
    async def test_below_failure_threshold_no_pause(self):
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(
                autonomy_timeout_minutes=0,
                max_consecutive_failures=3,
            ),
            autonomy_mode=True,
        )
        engine._autonomy_started_at = time.time()
        engine._consecutive_failures = 2  # Below threshold

        should_continue, events = await _detect_and_await_pause(
            engine, step=1, progress={}, resume_handler=None
        )

        assert should_continue is True
        assert events == []

    @pytest.mark.asyncio
    async def test_resume_handler_exception_cancels(self):
        """If resume_handler raises, treat as cancel (should_continue=False)."""
        engine = ReActEngine(
            llm_gateway=MagicMock(),
            dangerous_tools_config=DangerousToolsConfig(
                autonomy_timeout_minutes=0,
                max_consecutive_failures=3,
            ),
            autonomy_mode=True,
        )
        engine._consecutive_failures = 3

        async def failing_handler(token, reason):
            raise RuntimeError("handler crashed")

        should_continue, events = await _detect_and_await_pause(
            engine, step=1, progress={}, resume_handler=failing_handler
        )

        assert should_continue is False  # Cancelled due to exception
        assert events[0].data["reason"] == "consecutive_failures"


# ---------------------------------------------------------------------------
# Config parsing
# ---------------------------------------------------------------------------


class TestAutonomyConfigParsing:
    """DangerousToolsConfig parses autonomy thresholds from dict."""

    def test_defaults(self):
        cfg = DangerousToolsConfig()
        assert cfg.autonomy_timeout_minutes == 30
        assert cfg.max_consecutive_failures == 3

    def test_custom_values_from_dict(self):
        cfg = DangerousToolsConfig.from_dict(
            {
                "enabled": True,
                "autonomy_timeout_minutes": 60,
                "max_consecutive_failures": 5,
            }
        )
        assert cfg.autonomy_timeout_minutes == 60
        assert cfg.max_consecutive_failures == 5

    def test_disabled_values_from_dict(self):
        cfg = DangerousToolsConfig.from_dict(
            {
                "autonomy_timeout_minutes": 0,
                "max_consecutive_failures": 0,
            }
        )
        assert cfg.autonomy_timeout_minutes == 0
        assert cfg.max_consecutive_failures == 0

    def test_partial_config_uses_defaults(self):
        cfg = DangerousToolsConfig.from_dict({"enabled": True})
        assert cfg.autonomy_timeout_minutes == 30
        assert cfg.max_consecutive_failures == 3

    def test_empty_dict_uses_defaults(self):
        cfg = DangerousToolsConfig.from_dict({})
        assert cfg.autonomy_timeout_minutes == 30
        assert cfg.max_consecutive_failures == 3


# ---------------------------------------------------------------------------
# U6: resume Future pre-registration race fix
# ---------------------------------------------------------------------------


class TestResumeFutureRaceFix:
    """U6: Future is pre-registered before send_json so the WebSocket loop
    can resolve it even if the resume message arrives before _resume_handler
    is called by the engine.

    These tests mirror the chat.py ``autonomy_paused`` event handler +
    ``_resume_handler`` closure interaction without spinning up a full
    WebSocket server.
    """

    @pytest.mark.asyncio
    async def test_future_registered_before_send(self):
        """Future exists in pending_autonomy_resumes before send_json fires.

        Simulates the autonomy_paused event handler: the pre-registration
        step (U6) must insert the Future before the event is forwarded to
        the frontend.
        """
        pending: dict[str, asyncio.Future] = {}
        event_data = {
            "resume_token": "autonomy_pause:timeout:1",
            "reason": "timeout",
        }
        # Mirror chat.py: autonomy_paused event handler pre-registration.
        resume_token = event_data.get("resume_token")
        if resume_token and resume_token not in pending:
            loop = asyncio.get_running_loop()
            pending[resume_token] = loop.create_future()
        # At this point send_json has NOT run yet — Future is already
        # registered, so a concurrent resume message can resolve it.
        assert resume_token in pending
        assert not pending[resume_token].done()

    @pytest.mark.asyncio
    async def test_resume_message_not_lost_when_arriving_before_handler(self):
        """Resume message arriving before _resume_handler is awaited still
        resolves the Future (no lost message).

        This is the core race the fix addresses: pre-registration lets the
        WebSocket loop find the Future even if the engine has not yet
        called _resume_handler.
        """
        pending: dict[str, asyncio.Future] = {}
        resume_token = "autonomy_pause:timeout:1"

        # Step 1: autonomy_paused event handler pre-registers Future.
        loop = asyncio.get_running_loop()
        pre_future: asyncio.Future[bool] = loop.create_future()
        pending[resume_token] = pre_future

        # Step 2: WebSocket loop receives resume BEFORE engine calls
        # _resume_handler (the race window the fix closes).
        assert resume_token in pending
        fut = pending[resume_token]
        if not fut.done():
            fut.set_result(True)

        # Step 3: engine finally calls _resume_handler, which awaits the
        # already-resolved Future — no deadlock.
        assert pre_future.done()
        assert pre_future.result() is True

    @pytest.mark.asyncio
    async def test_resume_handler_uses_pre_registered_future(self):
        """_resume_handler finds and awaits the pre-registered Future
        instead of creating a new one (which would orphan the pre-registered
        Future and lose the resume message)."""
        pending: dict[str, asyncio.Future] = {}
        resume_token = "autonomy_pause:consecutive_failures:3"

        # Pre-register (autonomy_paused event handler).
        loop = asyncio.get_running_loop()
        pre_future: asyncio.Future[bool] = loop.create_future()
        pending[resume_token] = pre_future

        # Mirror _resume_handler: look up existing Future first.
        future = pending.get(resume_token)
        assert future is pre_future  # Same object — no orphaned Future.

        # Resolve and await.
        pre_future.set_result(True)
        result = await asyncio.wait_for(future, timeout=1.0)
        assert result is True

    @pytest.mark.asyncio
    async def test_resume_handler_fallback_creates_future(self):
        """When no pre-registered Future exists (non-WS caller or event not
        seen), _resume_handler creates one as a fallback."""
        pending: dict[str, asyncio.Future] = {}
        resume_token = "autonomy_pause:timeout:1"

        # Mirror _resume_handler fallback path.
        future = pending.get(resume_token)
        if future is None:
            loop = asyncio.get_running_loop()
            future = loop.create_future()
            pending[resume_token] = future

        assert resume_token in pending
        assert pending[resume_token] is future
        assert not future.done()

    @pytest.mark.asyncio
    async def test_double_resolve_guard(self):
        """Duplicate resume messages do not crash — the second is ignored
        because the Future is already done.

        Mirrors the WebSocket loop's double-resolve guard
        (``if not fut.done(): fut.set_result(...)``).
        """
        pending: dict[str, asyncio.Future] = {}
        resume_token = "autonomy_pause:timeout:1"
        loop = asyncio.get_running_loop()
        future: asyncio.Future[bool] = loop.create_future()
        pending[resume_token] = future

        # First resume resolves the Future.
        assert not future.done()
        future.set_result(True)

        # Second resume: guard prevents double-resolve (would raise
        # InvalidStateError without the done() check).
        if not future.done():
            future.set_result(True)  # Not reached — guard holds.

        assert future.done()
        assert future.result() is True

    @pytest.mark.asyncio
    async def test_pre_registration_does_not_overwrite_existing(self):
        """Pre-registration skips when a Future already exists (idempotent).

        Guards against the edge case where the engine emits autonomy_paused
        twice with the same token (e.g. retry) — the second pre-registration
        must not replace the first Future.
        """
        pending: dict[str, asyncio.Future] = {}
        resume_token = "autonomy_pause:timeout:1"
        loop = asyncio.get_running_loop()
        original: asyncio.Future[bool] = loop.create_future()
        pending[resume_token] = original

        # Mirror chat.py: pre-registration checks ``not in``.
        if resume_token and resume_token not in pending:
            pending[resume_token] = loop.create_future()

        assert pending[resume_token] is original  # Unchanged.

    @pytest.mark.asyncio
    async def test_full_flow_pre_register_then_handler_awaits(self):
        """End-to-end: pre-register → resolve from WS loop → handler awaits
        and returns True.

        Validates the complete U6 fix: Future is pre-registered before the
        event is sent, the WS loop resolves it, and _resume_handler returns
        the resolved value without creating a new Future.
        """
        pending: dict[str, asyncio.Future] = {}
        resume_token = "autonomy_pause:timeout:5"
        event_data = {"resume_token": resume_token, "reason": "timeout"}

        # 1. autonomy_paused event handler: pre-register BEFORE send_json.
        rt = event_data.get("resume_token")
        if rt and rt not in pending:
            loop = asyncio.get_running_loop()
            pending[rt] = loop.create_future()

        # 2. WS loop resolves the Future (resume message arrives).
        fut = pending[rt]
        if not fut.done():
            fut.set_result(True)

        # 3. _resume_handler: look up pre-registered Future, await it.
        handler_future = pending.get(rt)
        assert handler_future is not None
        result = await asyncio.wait_for(handler_future, timeout=1.0)

        assert result is True
        # 4. Cleanup (finally block in _resume_handler).
        pending.pop(rt, None)
        assert rt not in pending