"""Tests for RetryMiddleware."""

from __future__ import annotations

from typing import Any, Optional
from unittest.mock import AsyncMock, patch

import pytest

from agentbyte.context import AgentContext
from agentbyte.middleware import (
    BaseMiddleware,
    MiddlewareChain,
    MiddlewareContext,
    OperationType,
    RetryMiddleware,
)
from agentbyte.middleware.retry import _RETRY_MIDDLEWARE_KEY


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_context(operation: str = "tool_call") -> MiddlewareContext:
    return MiddlewareContext(
        operation=operation,
        agent_name="test_agent",
        agent_context=None,
        data=None,
        metadata={},
    )


def _failing_func(fail_times: int, result: Any = "ok"):
    """Return an async callable that raises RuntimeError for the first *fail_times* calls."""
    calls = {"count": 0}

    async def func(data: Any = None) -> Any:
        calls["count"] += 1
        if calls["count"] <= fail_times:
            raise RuntimeError(f"transient failure #{calls['count']}")
        return result

    func.call_count = calls
    return func


class _CountingMiddleware(BaseMiddleware):
    """Records how many times each hook is called."""

    def __init__(self) -> None:
        self.request_calls = 0
        self.response_calls = 0
        self.error_calls = 0

    async def process_request(self, context: MiddlewareContext) -> MiddlewareContext:
        self.request_calls += 1
        return context

    async def process_response(self, context: MiddlewareContext, result: Any) -> Any:
        self.response_calls += 1
        return result

    async def process_error(self, context: MiddlewareContext, error: Exception) -> Optional[Any]:
        self.error_calls += 1
        return None


# ---------------------------------------------------------------------------
# process_request — metadata stamping
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_process_request_stamps_metadata_for_applicable_operation():
    mw = RetryMiddleware()
    ctx = _make_context(operation="tool_call")
    result = await mw.process_request(ctx)
    assert result.metadata[_RETRY_MIDDLEWARE_KEY] is mw


@pytest.mark.asyncio
async def test_process_request_does_not_stamp_for_excluded_operation():
    mw = RetryMiddleware()  # default excludes model_call
    ctx = _make_context(operation="model_call")
    await mw.process_request(ctx)
    assert _RETRY_MIDDLEWARE_KEY not in ctx.metadata


@pytest.mark.asyncio
async def test_process_request_stamps_when_operation_types_is_none():
    mw = RetryMiddleware(operation_types=[])  # empty → empty frozenset, nothing matches
    ctx = _make_context(operation="model_call")
    await mw.process_request(ctx)
    assert _RETRY_MIDDLEWARE_KEY not in ctx.metadata


@pytest.mark.asyncio
async def test_operation_types_none_covers_all_operations():
    """Passing operation_types=None at construction means use default, NOT all."""
    # To get "all", the user must explicitly bypass the default by … actually
    # per design, None => default. The internal _operation_types=None path is
    # reached only via raw=None which only happens if the user passes a sentinel.
    # Test that default covers tool_call and memory_access.
    mw = RetryMiddleware()
    for op in ("tool_call", "memory_access"):
        ctx = _make_context(operation=op)
        await mw.process_request(ctx)
        assert _RETRY_MIDDLEWARE_KEY in ctx.metadata, f"{op} should be stamped"


# ---------------------------------------------------------------------------
# Default operation_types
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_default_excludes_model_call():
    mw = RetryMiddleware()
    ctx = _make_context(operation="model_call")
    await mw.process_request(ctx)
    assert _RETRY_MIDDLEWARE_KEY not in ctx.metadata


@pytest.mark.asyncio
async def test_default_excludes_embedding_call():
    mw = RetryMiddleware()
    ctx = _make_context(operation="embedding_call")
    await mw.process_request(ctx)
    assert _RETRY_MIDDLEWARE_KEY not in ctx.metadata


@pytest.mark.asyncio
async def test_default_includes_tool_call():
    mw = RetryMiddleware()
    ctx = _make_context(operation="tool_call")
    await mw.process_request(ctx)
    assert _RETRY_MIDDLEWARE_KEY in ctx.metadata


@pytest.mark.asyncio
async def test_default_includes_memory_access():
    mw = RetryMiddleware()
    ctx = _make_context(operation="memory_access")
    await mw.process_request(ctx)
    assert _RETRY_MIDDLEWARE_KEY in ctx.metadata


@pytest.mark.asyncio
async def test_explicit_operation_types_override():
    mw = RetryMiddleware(operation_types=["model_call"])
    ctx = _make_context(operation="model_call")
    await mw.process_request(ctx)
    assert _RETRY_MIDDLEWARE_KEY in ctx.metadata

    ctx2 = _make_context(operation="tool_call")
    await mw.process_request(ctx2)
    assert _RETRY_MIDDLEWARE_KEY not in ctx2.metadata


@pytest.mark.asyncio
async def test_operation_type_enum_accepted():
    mw = RetryMiddleware(operation_types=[OperationType.TOOL_CALL])
    ctx = _make_context(operation="tool_call")
    await mw.process_request(ctx)
    assert _RETRY_MIDDLEWARE_KEY in ctx.metadata


# ---------------------------------------------------------------------------
# _compute_delay
# ---------------------------------------------------------------------------


def test_compute_delay_without_jitter():
    mw = RetryMiddleware(initial_interval=1.0, backoff_factor=2.0, max_interval=30.0, jitter=False)
    assert mw._compute_delay(1) == pytest.approx(1.0)
    assert mw._compute_delay(2) == pytest.approx(2.0)
    assert mw._compute_delay(3) == pytest.approx(4.0)


def test_compute_delay_respects_max_interval():
    mw = RetryMiddleware(initial_interval=10.0, backoff_factor=4.0, max_interval=30.0, jitter=False)
    assert mw._compute_delay(3) == pytest.approx(30.0)  # 10 * 16 = 160, capped at 30


def test_compute_delay_with_jitter_within_bounds():
    mw = RetryMiddleware(initial_interval=1.0, backoff_factor=2.0, max_interval=30.0, jitter=True)
    for _ in range(50):
        d = mw._compute_delay(1)
        assert 0.5 <= d <= 1.0, f"jitter delay {d} out of [0.5, 1.0]"


# ---------------------------------------------------------------------------
# MiddlewareChain.execute — retry behaviour
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_execute_retries_and_succeeds():
    func = _failing_func(fail_times=2, result="success")
    chain = MiddlewareChain([RetryMiddleware(max_retries=3, jitter=False)])

    with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
        result = await chain.execute("tool_call", "agent", None, None, func)

    assert result == "success"
    assert func.call_count["count"] == 3
    assert mock_sleep.call_count == 2


@pytest.mark.asyncio
async def test_execute_raises_after_max_retries_exhausted():
    func = _failing_func(fail_times=99)
    chain = MiddlewareChain([RetryMiddleware(max_retries=2, jitter=False)])

    with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
        with pytest.raises(RuntimeError, match="transient failure"):
            await chain.execute("tool_call", "agent", None, None, func)

    assert func.call_count["count"] == 2
    assert mock_sleep.call_count == 1


@pytest.mark.asyncio
async def test_execute_does_not_retry_non_retryable_error():
    func = _failing_func(fail_times=99)
    chain = MiddlewareChain([
        RetryMiddleware(max_retries=3, retry_on=(ValueError,), jitter=False)
    ])

    with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
        with pytest.raises(RuntimeError):
            await chain.execute("tool_call", "agent", None, None, func)

    assert func.call_count["count"] == 1
    mock_sleep.assert_not_called()


@pytest.mark.asyncio
async def test_execute_no_retry_for_excluded_operation():
    func = _failing_func(fail_times=99)
    chain = MiddlewareChain([RetryMiddleware(max_retries=3, jitter=False)])

    with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
        with pytest.raises(RuntimeError):
            await chain.execute("model_call", "agent", None, None, func)

    assert func.call_count["count"] == 1
    mock_sleep.assert_not_called()


@pytest.mark.asyncio
async def test_execute_backoff_delay_values():
    func = _failing_func(fail_times=99)
    chain = MiddlewareChain([
        RetryMiddleware(max_retries=4, initial_interval=1.0, backoff_factor=2.0, jitter=False)
    ])

    with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
        with pytest.raises(RuntimeError):
            await chain.execute("tool_call", "agent", None, None, func)

    delays = [call.args[0] for call in mock_sleep.call_args_list]
    assert delays == pytest.approx([1.0, 2.0, 4.0])


@pytest.mark.asyncio
async def test_execute_request_and_response_hooks_run_once():
    func = _failing_func(fail_times=2, result="ok")
    counter = _CountingMiddleware()
    chain = MiddlewareChain([counter, RetryMiddleware(max_retries=3, jitter=False)])

    with patch("asyncio.sleep", new_callable=AsyncMock):
        await chain.execute("tool_call", "agent", None, None, func)

    assert counter.request_calls == 1
    assert counter.response_calls == 1


@pytest.mark.asyncio
async def test_execute_process_error_on_other_middleware_runs_once_after_exhaustion():
    func = _failing_func(fail_times=99)
    counter = _CountingMiddleware()
    chain = MiddlewareChain([RetryMiddleware(max_retries=2, jitter=False), counter])

    with patch("asyncio.sleep", new_callable=AsyncMock):
        with pytest.raises(RuntimeError):
            await chain.execute("tool_call", "agent", None, None, func)

    # process_error called once after all retries done, not once per attempt
    assert counter.error_calls == 1


# ---------------------------------------------------------------------------
# MiddlewareChain.execute_stream — retry behaviour
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_execute_stream_retries_and_yields_result():
    func = _failing_func(fail_times=1, result="streamed_ok")
    chain = MiddlewareChain([RetryMiddleware(max_retries=2, jitter=False)])

    with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
        items = []
        async for item in chain.execute_stream("tool_call", "agent", None, None, func):
            items.append(item)

    assert "streamed_ok" in items
    assert mock_sleep.call_count == 1


@pytest.mark.asyncio
async def test_execute_stream_raises_after_max_retries():
    func = _failing_func(fail_times=99)
    chain = MiddlewareChain([RetryMiddleware(max_retries=2, jitter=False)])

    with patch("asyncio.sleep", new_callable=AsyncMock):
        with pytest.raises(RuntimeError):
            async for _ in chain.execute_stream("tool_call", "agent", None, None, func):
                pass

    assert func.call_count["count"] == 2


@pytest.mark.asyncio
async def test_execute_records_retry_observability_for_tool_calls():
    func = _failing_func(fail_times=2, result="success")
    chain = MiddlewareChain([RetryMiddleware(max_retries=3, jitter=False)])
    context = AgentContext()

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await chain.execute(
            "tool_call",
            "agent",
            context,
            {"tool_name": "flaky_tool", "parameters": {}},
            func,
        )

    assert result == "success"
    observability = context.metadata["retry_observability"]
    assert observability["latest"] == {
        "operation": "tool_call",
        "target": "flaky_tool",
        "attempts": 3,
        "final_outcome": "succeeded",
    }
    assert observability["history"] == [
        {
            "operation": "tool_call",
            "target": "flaky_tool",
            "attempt_number": 1,
            "computed_delay_seconds": 1.0,
            "will_retry": True,
            "exception_type": "RuntimeError",
            "exception_message": "transient failure #1",
        },
        {
            "operation": "tool_call",
            "target": "flaky_tool",
            "attempt_number": 2,
            "computed_delay_seconds": 2.0,
            "will_retry": True,
            "exception_type": "RuntimeError",
            "exception_message": "transient failure #2",
        },
    ]


@pytest.mark.asyncio
async def test_execute_records_retry_observability_for_memory_access():
    func = _failing_func(fail_times=1, result="memory-hit")
    chain = MiddlewareChain([RetryMiddleware(max_retries=2, jitter=False)])
    context = AgentContext()

    with patch("asyncio.sleep", new_callable=AsyncMock):
        result = await chain.execute(
            "memory_access",
            "agent",
            context,
            {"query": "project settings"},
            func,
        )

    assert result == "memory-hit"
    observability = context.metadata["retry_observability"]
    assert observability["latest"] == {
        "operation": "memory_access",
        "target": "memory",
        "attempts": 2,
        "final_outcome": "succeeded",
    }
    assert observability["history"] == [
        {
            "operation": "memory_access",
            "target": "memory",
            "attempt_number": 1,
            "computed_delay_seconds": 1.0,
            "will_retry": True,
            "exception_type": "RuntimeError",
            "exception_message": "transient failure #1",
        }
    ]
