"""Tests for middleware chain and built-in middleware implementations."""

from __future__ import annotations

from collections.abc import AsyncGenerator
from typing import Any, Optional

import pytest

from agentbyte.context import AgentContext
from agentbyte.agents import ToolApprovalEvent
from agentbyte.middleware import (
    ApprovalMiddleware,
    BaseMiddleware,
    ContextCompactionMiddleware,
    GuardrailMiddleware,
    LoggingMiddleware,
    MetricsMiddleware,
    MiddlewareChain,
    MiddlewareContext,
    OperationType,
    PIIRedactionMiddleware,
    RateLimitMiddleware,
)
from agentbyte.messages import AssistantMessage, Usage, UserMessage
from agentbyte.llm.types import ChatCompletionResult


class TraceMiddleware(BaseMiddleware):
    """Test middleware that traces request/response order."""

    def __init__(self, name: str):
        self.name = name

    async def process_request(self, context: MiddlewareContext) -> MiddlewareContext:
        context.metadata.setdefault("trace", []).append(f"request:{self.name}")
        return context

    async def process_response(self, context: MiddlewareContext, result: Any) -> Any:
        context.metadata.setdefault("trace", []).append(f"response:{self.name}")
        return {
            "value": result,
            "trace": context.metadata["trace"],
        }

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


class RecoveryMiddleware(BaseMiddleware):
    """Test middleware that recovers from errors."""

    async def process_request(self, context: MiddlewareContext) -> MiddlewareContext:
        return context

    async def process_response(self, context: MiddlewareContext, result: Any) -> Any:
        return result

    async def process_error(
        self,
        context: MiddlewareContext,
        error: Exception,
    ) -> Optional[Any]:
        return {"recovered": True, "error": str(error)}


class MutateDataMiddleware(BaseMiddleware):
    """Test middleware that mutates request payload data."""

    async def process_request(self, context: MiddlewareContext) -> MiddlewareContext:
        updated = context.data.copy()
        updated["value"] = updated["value"] + 1
        context.data = updated
        return context

    async def process_response(self, context: MiddlewareContext, result: Any) -> Any:
        return result

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


class FailingRequestMiddleware(BaseMiddleware):
    """Test middleware that fails during request phase."""

    async def process_request(self, context: MiddlewareContext) -> MiddlewareContext:
        raise RuntimeError("request middleware failed")

    async def process_response(self, context: MiddlewareContext, result: Any) -> Any:
        return result

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


class FailingErrorMiddleware(BaseMiddleware):
    """Test middleware that fails while handling errors."""

    async def process_request(self, context: MiddlewareContext) -> MiddlewareContext:
        return context

    async def process_response(self, context: MiddlewareContext, result: Any) -> Any:
        return result

    async def process_error(
        self,
        context: MiddlewareContext,
        error: Exception,
    ) -> Optional[Any]:
        raise RuntimeError("error middleware failed")


class DummyEvent:
    """Simple event-like object for execute_stream event assertions."""

    def __init__(self, event_type: str, source: str):
        self.event_type = event_type
        self.source = source


class StreamMiddleware(BaseMiddleware):
    """Generator-style middleware for execute_stream tests."""

    async def process_request(
        self,
        context: MiddlewareContext,
    ) -> AsyncGenerator[Any, None]:
        yield DummyEvent("request_event", "stream_middleware")
        yield context

    async def process_response(
        self,
        context: MiddlewareContext,
        result: Any,
    ) -> AsyncGenerator[Any, None]:
        yield DummyEvent("response_event", "stream_middleware")
        yield {"value": result, "stream": True}

    async def process_error(
        self,
        context: MiddlewareContext,
        error: Exception,
    ) -> AsyncGenerator[Any, None]:
        yield DummyEvent("error_event", "stream_middleware")
        yield {"recovered": True, "error": str(error)}


@pytest.mark.asyncio
async def test_middleware_chain_request_response_order() -> None:
    chain = MiddlewareChain([TraceMiddleware("mw1"), TraceMiddleware("mw2")])

    async def operation() -> str:
        return "ok"

    result = await chain.execute(
        operation="model_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"payload": "x"},
        func=operation,
    )

    assert result["trace"] == [
        "request:mw1",
        "request:mw2",
        "response:mw2",
        "response:mw1",
    ]


@pytest.mark.asyncio
async def test_middleware_chain_error_recovery() -> None:
    chain = MiddlewareChain([RecoveryMiddleware()])

    async def operation() -> str:
        raise RuntimeError("boom")

    result = await chain.execute(
        operation="tool_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"tool": "calc"},
        func=operation,
    )

    assert result["recovered"] is True
    assert "boom" in result["error"]


@pytest.mark.asyncio
async def test_guardrail_middleware_blocks_forbidden_input() -> None:
    chain = MiddlewareChain([GuardrailMiddleware(blocked_patterns=[r"password", r"api[_-]?key", r"secret"])])

    async def operation() -> str:
        return "should-not-run"

    with pytest.raises(ValueError, match="blocked pattern"):
        await chain.execute(
            operation="model_call",
            agent_name="assistant",
            agent_context=AgentContext(),
            data={"messages": [UserMessage(content="my password is 123", source="user")]},
            func=operation,
        )


@pytest.mark.asyncio
async def test_pii_redaction_middleware_redacts_embedding_input_payload() -> None:
    chain = MiddlewareChain([PIIRedactionMiddleware()])

    async def operation(data: dict[str, Any]) -> dict[str, Any]:
        return data

    result = await chain.execute(
        operation=OperationType.EMBEDDING_CALL.value,
        agent_name="embedding_client",
        agent_context=None,
        data={"input": ["contact me at foo@example.com"], "kwargs": {}},
        func=operation,
    )

    assert result["input"] == ["contact me at [EMAIL-REDACTED]"]


@pytest.mark.asyncio
async def test_rate_limit_middleware_enforces_limit() -> None:
    middleware = RateLimitMiddleware(max_requests=1, window_seconds=0.01)
    chain = MiddlewareChain([middleware])

    async def operation() -> str:
        return "ok"

    await chain.execute(
        operation="model_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"request": 1},
        func=operation,
    )

    second_result = await chain.execute(
        operation="model_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"request": 2},
        func=operation,
    )

    assert second_result == "ok"


@pytest.mark.asyncio
async def test_logging_middleware_passes_through() -> None:
    chain = MiddlewareChain([LoggingMiddleware()])

    async def operation() -> dict[str, bool]:
        return {"ok": True}

    result = await chain.execute(
        operation="model_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"test": True},
        func=operation,
    )

    assert result == {"ok": True}


@pytest.mark.asyncio
async def test_middleware_chain_request_data_mutation_affects_operation_input() -> None:
    chain = MiddlewareChain([MutateDataMiddleware()])

    async def operation(payload: dict[str, int]) -> int:
        return payload["value"]

    result = await chain.execute(
        operation="model_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"value": 1},
        func=operation,
    )

    assert result == 2


@pytest.mark.asyncio
async def test_middleware_chain_request_phase_error_can_recover() -> None:
    chain = MiddlewareChain([FailingRequestMiddleware(), RecoveryMiddleware()])

    async def operation() -> str:
        return "should-not-run"

    result = await chain.execute(
        operation="model_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"payload": "x"},
        func=operation,
    )

    assert result["recovered"] is True
    assert "request middleware failed" in result["error"]


@pytest.mark.asyncio
async def test_middleware_chain_continues_when_error_middleware_fails() -> None:
    chain = MiddlewareChain([RecoveryMiddleware(), FailingErrorMiddleware()])

    async def operation() -> str:
        raise RuntimeError("boom")

    result = await chain.execute(
        operation="tool_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"tool": "calc"},
        func=operation,
    )

    assert result["recovered"] is True
    assert "boom" in result["error"]


@pytest.mark.asyncio
async def test_middleware_chain_execute_stream_yields_events_and_result() -> None:
    chain = MiddlewareChain([StreamMiddleware()])

    async def operation() -> str:
        return "ok"

    items: list[Any] = []
    async for item in chain.execute_stream(
        operation="model_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"payload": "x"},
        func=operation,
    ):
        items.append(item)

    assert len(items) == 3
    assert items[0].event_type == "request_event"
    assert items[1].event_type == "response_event"
    assert items[2] == {"value": "ok", "stream": True}


@pytest.mark.asyncio
async def test_approval_middleware_emits_event_for_matching_tool() -> None:
    chain = MiddlewareChain([ApprovalMiddleware(tool_names=["sensitive_tool"])])

    async def operation(payload: dict[str, Any]) -> dict[str, Any]:
        return payload

    items: list[Any] = []
    async for item in chain.execute_stream(
        operation="tool_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"tool_name": "sensitive_tool", "parameters": {}, "call_id": "abc"},
        func=operation,
    ):
        items.append(item)

    assert len(items) == 1
    assert isinstance(items[0], ToolApprovalEvent)
    assert items[0].approval_request.tool_name == "sensitive_tool"


@pytest.mark.asyncio
async def test_approval_middleware_passthrough_for_non_matching_tool() -> None:
    chain = MiddlewareChain([ApprovalMiddleware(tool_names=["sensitive_tool"])])

    async def operation(payload: dict[str, Any]) -> str:
        return payload["tool_name"]

    items: list[Any] = []
    async for item in chain.execute_stream(
        operation="tool_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"tool_name": "safe_tool", "parameters": {}, "call_id": "x1"},
        func=operation,
    ):
        items.append(item)

    assert items == ["safe_tool"]


@pytest.mark.asyncio
async def test_approval_middleware_composes_with_logging_middleware() -> None:
    chain = MiddlewareChain(
        [
            LoggingMiddleware(),
            ApprovalMiddleware(tool_names=["sensitive_tool"]),
        ]
    )

    async def operation(payload: dict[str, Any]) -> dict[str, Any]:
        return payload

    items: list[Any] = []
    async for item in chain.execute_stream(
        operation="tool_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"tool_name": "sensitive_tool", "parameters": {}, "call_id": "abc"},
        func=operation,
    ):
        items.append(item)

    assert len(items) == 1
    assert isinstance(items[0], ToolApprovalEvent)


@pytest.mark.asyncio
async def test_guardrail_middleware_blocks_tool_and_pattern() -> None:
    chain = MiddlewareChain(
        [
            GuardrailMiddleware(
                blocked_tools=["danger_tool"],
                blocked_patterns=[r"secret"],
            )
        ]
    )

    async def operation(payload: dict[str, Any]) -> str:
        return payload["tool_name"]

    with pytest.raises(ValueError, match="blocked by guardrails"):
        await chain.execute(
            operation="tool_call",
            agent_name="assistant",
            agent_context=AgentContext(),
            data={"tool_name": "danger_tool", "parameters": {}, "call_id": "x"},
            func=operation,
        )

    with pytest.raises(ValueError, match="blocked pattern"):
        await chain.execute(
            operation="tool_call",
            agent_name="assistant",
            agent_context=AgentContext(),
            data={"tool_name": "safe_tool", "parameters": {"text": "my secret"}, "call_id": "x"},
            func=operation,
        )


@pytest.mark.asyncio
async def test_pii_redaction_middleware_redacts_model_request_and_response() -> None:
    chain = MiddlewareChain([PIIRedactionMiddleware()])

    async def operation(payload: dict[str, Any]) -> ChatCompletionResult:
        messages = payload["messages"]
        assert "[EMAIL-REDACTED]" in messages[0].content
        assert "john@example.com" not in messages[0].content
        return ChatCompletionResult(
            message=AssistantMessage(content="Email me at jane@example.com", source="assistant"),
            usage=Usage(tokens_input=1, tokens_output=1),
            model="fake",
        )

    result = await chain.execute(
        operation="model_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={
            "messages": [
                AssistantMessage(content="Contact john@example.com", source="user"),
            ]
        },
        func=operation,
    )

    assert "[EMAIL-REDACTED]" in result.message.content
    assert "jane@example.com" not in result.message.content


@pytest.mark.asyncio
async def test_metrics_middleware_tracks_operation_counts_and_errors() -> None:
    middleware = MetricsMiddleware()
    chain = MiddlewareChain([middleware])

    async def ok_operation() -> str:
        return "ok"

    async def failing_operation() -> str:
        raise RuntimeError("boom")

    result = await chain.execute(
        operation="model_call",
        agent_name="assistant",
        agent_context=AgentContext(),
        data={"request": 1},
        func=ok_operation,
    )
    assert result == "ok"

    with pytest.raises(RuntimeError, match="boom"):
        await chain.execute(
            operation="tool_call",
            agent_name="assistant",
            agent_context=AgentContext(),
            data={"tool": "x"},
            func=failing_operation,
        )

    metrics = middleware.get_metrics()
    assert metrics["total_operations"] == 2
    assert metrics["operations_by_type"]["model_call"] == 1
    assert metrics["operations_by_type"]["tool_call"] == 1
    assert metrics["errors_by_type"]["RuntimeError"] == 1
    assert metrics["average_duration"] >= 0


# ===========================================================================
# OperationType enum
# ===========================================================================


def test_operation_type_values_are_strings():
    """OperationType values must equal plain strings for backward compat."""
    assert OperationType.MODEL_CALL == "model_call"
    assert OperationType.TOOL_CALL == "tool_call"
    assert OperationType.MEMORY_ACCESS == "memory_access"
    assert OperationType.APPROVAL == "approval"


def test_operation_type_interchangeable_with_str():
    """MiddlewareContext should accept either the enum or a bare string."""
    ctx_enum = MiddlewareContext(
        operation=OperationType.MODEL_CALL,
        agent_name="a",
        agent_context=AgentContext(),
        data={},
    )
    ctx_str = MiddlewareContext(
        operation="model_call",
        agent_name="a",
        agent_context=AgentContext(),
        data={},
    )
    assert ctx_enum.operation == ctx_str.operation


# ===========================================================================
# ContextCompactionMiddleware
# ===========================================================================


def _make_messages(n: int) -> list:
    """Create *n* UserMessages with ~100 chars each for token estimation."""
    return [UserMessage(content="x" * 100, source="user") for _ in range(n)]


@pytest.mark.asyncio
async def test_compaction_not_applied_below_threshold():
    """No compaction when token estimate is under max_tokens."""
    mw = ContextCompactionMiddleware(max_tokens=100_000, keep_last=5)
    messages = _make_messages(3)
    ctx = MiddlewareContext(
        operation=OperationType.MODEL_CALL,
        agent_name="a",
        agent_context=AgentContext(),
        data={"messages": messages},
    )
    result = await mw.process_request(ctx)
    assert result.data["messages"] is messages
    assert not result.metadata.get("compaction_applied")


@pytest.mark.asyncio
async def test_compaction_fires_when_over_threshold():
    """Compaction is applied when estimated tokens exceed max_tokens."""
    # 50 messages × 100 chars = 5000 chars ÷ 4 = 1250 estimated tokens
    mw = ContextCompactionMiddleware(max_tokens=500, keep_last=3)
    messages = _make_messages(50)
    ctx = MiddlewareContext(
        operation=OperationType.MODEL_CALL,
        agent_name="a",
        agent_context=AgentContext(),
        data={"messages": messages},
    )
    result = await mw.process_request(ctx)
    assert result.metadata.get("compaction_applied") is True
    assert result.metadata["original_message_count"] == 50
    compacted = result.data["messages"]
    # Structure: [original_first, summary, *tail(keep_last=3)]
    assert len(compacted) == 2 + 3  # system + placeholder + 3 tail messages


@pytest.mark.asyncio
async def test_compaction_preserves_first_message():
    """After compaction the first element is the original first message."""
    mw = ContextCompactionMiddleware(max_tokens=200, keep_last=2)
    first_msg = UserMessage(content="system instructions " * 5, source="user")
    messages = [first_msg] + _make_messages(20)
    ctx = MiddlewareContext(
        operation=OperationType.MODEL_CALL,
        agent_name="a",
        agent_context=AgentContext(),
        data={"messages": messages},
    )
    result = await mw.process_request(ctx)
    assert result.data["messages"][0] is first_msg


@pytest.mark.asyncio
async def test_compaction_not_applied_for_tool_call():
    """Compaction must NOT fire for tool_call operations."""
    mw = ContextCompactionMiddleware(max_tokens=1, keep_last=2)
    messages = _make_messages(50)
    ctx = MiddlewareContext(
        operation=OperationType.TOOL_CALL,
        agent_name="a",
        agent_context=AgentContext(),
        data={"messages": messages},
    )
    result = await mw.process_request(ctx)
    assert result.data["messages"] is messages
    assert not result.metadata.get("compaction_applied")


@pytest.mark.asyncio
async def test_compaction_init_validation():
    """ContextCompactionMiddleware rejects non-positive constructor args."""
    with pytest.raises(ValueError):
        ContextCompactionMiddleware(max_tokens=0)
    with pytest.raises(ValueError):
        ContextCompactionMiddleware(keep_last=0)


@pytest.mark.asyncio
async def test_compaction_process_response_passthrough():
    """process_response and process_error are no-op pass-throughs."""
    mw = ContextCompactionMiddleware()
    ctx = MiddlewareContext(
        operation=OperationType.MODEL_CALL,
        agent_name="a",
        agent_context=AgentContext(),
        data={},
    )
    sentinel = object()
    assert await mw.process_response(ctx, sentinel) is sentinel
    assert await mw.process_error(ctx, ValueError("x")) is None
