#!/usr/bin/env python3
"""
LockStock Thin Middleware Tests
--------------------------------
Unit tests for the pure ASGI thin middleware that calls /v1/stamp.

Tests:
1. Path task resolution (longest prefix match)
2. Stamp call and header injection
3. 409 retry with backoff
4. Fire-and-forget response stamping
5. Path exclusion (None value)
6. Block on rejection mode
7. Response task auto-inference (None/explicit/{} semantics)
8. Default fail-closed behavior (block_on_rejection=True)
"""

import asyncio
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from dataclasses import dataclass
from typing import Dict, List, Optional, Any

import httpx

from lockstock_fastapi.thin import LockStockThinMiddleware, StampResult


# Test fixtures

@pytest.fixture
def mock_app():
    """Create a mock ASGI app that returns 200 OK."""
    async def app(scope, receive, send):
        await send({
            "type": "http.response.start",
            "status": 200,
            "headers": [(b"content-type", b"application/json")],
        })
        await send({
            "type": "http.response.body",
            "body": b'{"message": "ok"}',
            "more_body": False,
        })
    return app


@pytest.fixture
def http_scope():
    """Create a basic HTTP request scope."""
    return {
        "type": "http",
        "method": "POST",
        "path": "/v1/chat/completions",
        "headers": [],
    }


@pytest.fixture
def middleware(mock_app):
    """Create middleware instance with test configuration."""
    return LockStockThinMiddleware(
        mock_app,
        agent_id="agent_test123",
        api_key="lsk_test_key",
        server_url="https://test.example.com",
        path_tasks={
            "/v1/chat/completions": "chatcompletion",
            "/v1/models": "network",
            "/health": None,  # Excluded
        },
    )


# Helper to capture sent messages

class MessageCapture:
    """Helper to capture ASGI send messages."""

    def __init__(self):
        self.messages: List[Dict[str, Any]] = []

    async def __call__(self, message: Dict[str, Any]) -> None:
        self.messages.append(message)

    def get_headers(self) -> Dict[bytes, bytes]:
        """Extract headers from response.start message."""
        for msg in self.messages:
            if msg.get("type") == "http.response.start":
                return dict(msg.get("headers", []))
        return {}


# Tests

class TestPathResolution:
    """Test path task resolution logic."""

    def test_exact_match(self, middleware):
        """Exact path match returns correct task."""
        task = middleware._resolve_task("/v1/chat/completions")
        assert task == "chatcompletion"

    def test_prefix_match(self, middleware):
        """Path with trailing segments still matches."""
        task = middleware._resolve_task("/v1/chat/completions/extra")
        assert task == "chatcompletion"

    def test_no_match(self, middleware):
        """Unknown path returns None."""
        task = middleware._resolve_task("/unknown/path")
        assert task is None

    def test_excluded_path(self, middleware):
        """Path with None value is excluded."""
        task = middleware._resolve_task("/health")
        assert task is None

    def test_longest_prefix_wins(self):
        """Longest prefix match takes priority."""
        async def app(scope, receive, send):
            pass

        mw = LockStockThinMiddleware(
            app,
            agent_id="test",
            api_key="test",
            path_tasks={
                "/v1": "base",
                "/v1/chat": "chat",
                "/v1/chat/completions": "chatcompletion",
            },
        )

        assert mw._resolve_task("/v1/other") == "base"
        assert mw._resolve_task("/v1/chat/other") == "chat"
        assert mw._resolve_task("/v1/chat/completions/stream") == "chatcompletion"


class TestStampResult:
    """Test StampResult dataclass."""

    def test_authorized_result(self):
        """Authorized result has all fields."""
        result = StampResult(
            authorized=True,
            sequence=42,
            state_hash="abc123",
            trace_id="trace-1",
            timestamp=1234567890,
            http_status=200,
        )
        assert result.authorized is True
        assert result.sequence == 42
        assert result.state_hash == "abc123"
        assert result.reason is None

    def test_rejected_result(self):
        """Rejected result has reason."""
        result = StampResult(
            authorized=False,
            reason="Task not allowed",
            http_status=403,
        )
        assert result.authorized is False
        assert result.reason == "Task not allowed"
        assert result.sequence is None


class TestStampCall:
    """Test /v1/stamp API calls."""

    @pytest.mark.asyncio
    async def test_successful_stamp(self, middleware):
        """Successful stamp returns authorized result."""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "authorized": True,
            "sequence": 5,
            "state_hash": "deadbeef",
            "trace_id": "trace-123",
            "timestamp": 1234567890,
        }

        with patch.object(middleware, "_get_client") as mock_get_client:
            mock_client = AsyncMock()
            mock_client.post.return_value = mock_response
            mock_get_client.return_value = mock_client

            result = await middleware._stamp("chatcompletion")

            assert result.authorized is True
            assert result.sequence == 5
            assert result.state_hash == "deadbeef"

            # Verify API call
            mock_client.post.assert_called_once()
            call_args = mock_client.post.call_args
            assert "v1/stamp" in call_args[0][0]
            assert call_args[1]["json"]["agent_id"] == "agent_test123"
            assert call_args[1]["json"]["task"] == "chatcompletion"

    @pytest.mark.asyncio
    async def test_rejected_stamp(self, middleware):
        """Rejected stamp returns unauthorized result."""
        mock_response = MagicMock()
        mock_response.status_code = 403
        mock_response.json.return_value = {
            "authorized": False,
            "reason": "Task not allowed",
        }

        with patch.object(middleware, "_get_client") as mock_get_client:
            mock_client = AsyncMock()
            mock_client.post.return_value = mock_response
            mock_get_client.return_value = mock_client

            result = await middleware._stamp("unauthorized_task")

            assert result.authorized is False
            assert "not allowed" in result.reason

    @pytest.mark.asyncio
    async def test_connection_error(self, middleware):
        """Connection error returns unauthorized with reason."""
        with patch.object(middleware, "_get_client") as mock_get_client:
            mock_client = AsyncMock()
            mock_client.post.side_effect = httpx.ConnectError("Connection refused")
            mock_get_client.return_value = mock_client

            result = await middleware._stamp("chatcompletion")

            assert result.authorized is False
            assert "Connection error" in result.reason


class TestRetryLogic:
    """Test 409 retry with exponential backoff."""

    @pytest.mark.asyncio
    async def test_retry_on_409(self, middleware):
        """409 Conflict triggers retry."""
        call_count = 0

        async def mock_stamp(task, **kwargs):
            nonlocal call_count
            call_count += 1
            if call_count < 3:
                return StampResult(
                    authorized=False,
                    reason="Conflict",
                    http_status=409,
                )
            return StampResult(
                authorized=True,
                sequence=5,
                state_hash="success",
                http_status=200,
            )

        with patch.object(middleware, "_stamp", side_effect=mock_stamp):
            with patch("asyncio.sleep", new_callable=AsyncMock):
                result = await middleware._stamp_with_retry("chatcompletion")

        assert result.authorized is True
        assert call_count == 3

    @pytest.mark.asyncio
    async def test_max_retries_exceeded(self, middleware):
        """Max retries returns last failure."""
        async def always_conflict(task, **kwargs):
            return StampResult(
                authorized=False,
                reason="Conflict",
                http_status=409,
            )

        with patch.object(middleware, "_stamp", side_effect=always_conflict):
            with patch("asyncio.sleep", new_callable=AsyncMock):
                result = await middleware._stamp_with_retry("chatcompletion")

        assert result.authorized is False
        assert result.http_status == 409

    @pytest.mark.asyncio
    async def test_no_retry_on_other_errors(self, middleware):
        """Non-409 errors don't trigger retry."""
        call_count = 0

        async def mock_stamp(task, **kwargs):
            nonlocal call_count
            call_count += 1
            return StampResult(
                authorized=False,
                reason="Forbidden",
                http_status=403,
            )

        with patch.object(middleware, "_stamp", side_effect=mock_stamp):
            result = await middleware._stamp_with_retry("chatcompletion")

        assert call_count == 1  # No retry
        assert result.http_status == 403


class TestHeaderInjection:
    """Test X-LockStock-* header injection."""

    @pytest.mark.asyncio
    async def test_headers_on_success(self, middleware, http_scope):
        """Successful stamp injects headers."""
        capture = MessageCapture()

        async def mock_receive():
            return {"type": "http.request", "body": b""}

        with patch.object(middleware, "_stamp_with_retry") as mock_stamp:
            mock_stamp.return_value = StampResult(
                authorized=True,
                sequence=42,
                state_hash="abc123def456",
                trace_id="trace-789",
                http_status=200,
            )

            await middleware(http_scope, mock_receive, capture)

        headers = capture.get_headers()
        assert headers.get(b"x-lockstock-status") == b"accepted"
        assert headers.get(b"x-lockstock-seq") == b"42"
        assert headers.get(b"x-lockstock-hash") == b"abc123def456"
        assert headers.get(b"x-lockstock-trace") == b"trace-789"

    @pytest.mark.asyncio
    async def test_headers_on_skip(self, mock_app, http_scope):
        """Skipped stamp (not blocking) injects skip headers."""
        # Explicit block_on_rejection=False to test non-blocking skip behavior
        mw = LockStockThinMiddleware(
            mock_app,
            agent_id="agent_test123",
            api_key="lsk_test_key",
            server_url="https://test.example.com",
            path_tasks={
                "/v1/chat/completions": "chatcompletion",
                "/v1/models": "network",
                "/health": None,
            },
            block_on_rejection=False,
        )

        capture = MessageCapture()

        async def mock_receive():
            return {"type": "http.request", "body": b""}

        with patch.object(mw, "_stamp_with_retry") as mock_stamp:
            mock_stamp.return_value = StampResult(
                authorized=False,
                reason="Connection error",
                http_status=None,
            )

            await mw(http_scope, mock_receive, capture)

        headers = capture.get_headers()
        assert headers.get(b"x-lockstock-status") == b"skipped"
        assert b"x-lockstock-reason" in headers


class TestBlockOnRejection:
    """Test block_on_rejection mode."""

    @pytest.mark.asyncio
    async def test_blocking_returns_503(self, mock_app):
        """Block mode returns 503 on rejection."""
        mw = LockStockThinMiddleware(
            mock_app,
            agent_id="test",
            api_key="test",
            path_tasks={"/v1/chat/completions": "chatcompletion"},
            block_on_rejection=True,
        )

        capture = MessageCapture()
        scope = {"type": "http", "method": "POST", "path": "/v1/chat/completions"}

        async def mock_receive():
            return {"type": "http.request", "body": b""}

        with patch.object(mw, "_stamp_with_retry") as mock_stamp:
            mock_stamp.return_value = StampResult(
                authorized=False,
                reason="Forbidden",
                http_status=403,
            )

            await mw(scope, mock_receive, capture)

        # Check 503 response
        start_msg = capture.messages[0]
        assert start_msg["type"] == "http.response.start"
        assert start_msg["status"] == 503

        headers = dict(start_msg["headers"])
        assert headers.get(b"x-lockstock-status") == b"rejected"


class TestResponseTaskAutoInference:
    """Test response_tasks auto-inference from path_tasks."""

    def test_auto_infer_when_none(self):
        """response_tasks=None auto-infers TASK+RESPONSE convention."""
        async def app(scope, receive, send):
            pass

        mw = LockStockThinMiddleware(
            app,
            agent_id="test",
            api_key="test",
            path_tasks={
                "/v1/chat/completions": "chatcompletion",
                "/v1/models": "network",
                "/health": None,  # Excluded
            },
            # response_tasks not passed — defaults to None (auto-infer)
        )

        assert mw.response_tasks == {
            "chatcompletion": "chatcompletionRESPONSE",
            "network": "networkRESPONSE",
        }

    def test_auto_infer_skips_none_tasks(self):
        """Auto-inference skips paths with None (excluded) tasks."""
        async def app(scope, receive, send):
            pass

        mw = LockStockThinMiddleware(
            app,
            agent_id="test",
            api_key="test",
            path_tasks={
                "/health": None,
                "/metrics": None,
            },
        )

        assert mw.response_tasks == {}

    def test_explicit_dict_used_as_is(self):
        """Explicit response_tasks dict is used without modification."""
        async def app(scope, receive, send):
            pass

        explicit = {"chatcompletion": "CUSTOM_RESPONSE"}
        mw = LockStockThinMiddleware(
            app,
            agent_id="test",
            api_key="test",
            path_tasks={"/v1/chat/completions": "chatcompletion"},
            response_tasks=explicit,
        )

        assert mw.response_tasks == {"chatcompletion": "CUSTOM_RESPONSE"}

    def test_explicit_empty_dict_is_opt_out(self):
        """Explicit empty dict {} disables response stamping."""
        async def app(scope, receive, send):
            pass

        mw = LockStockThinMiddleware(
            app,
            agent_id="test",
            api_key="test",
            path_tasks={"/v1/chat/completions": "chatcompletion"},
            response_tasks={},  # Conscious opt-out
        )

        assert mw.response_tasks == {}

    def test_empty_path_tasks_no_inference(self):
        """No path_tasks means no response_tasks to infer."""
        async def app(scope, receive, send):
            pass

        mw = LockStockThinMiddleware(
            app,
            agent_id="test",
            api_key="test",
        )

        assert mw.response_tasks == {}


class TestDefaultFailClosed:
    """Test that block_on_rejection defaults to True (fail-closed)."""

    @pytest.mark.asyncio
    async def test_default_blocks_on_rejection(self, mock_app):
        """Default middleware blocks rejected requests with 503."""
        mw = LockStockThinMiddleware(
            mock_app,
            agent_id="test",
            api_key="test",
            path_tasks={"/v1/chat/completions": "chatcompletion"},
            # block_on_rejection not passed — defaults to True
        )

        capture = MessageCapture()
        scope = {"type": "http", "method": "POST", "path": "/v1/chat/completions"}

        async def mock_receive():
            return {"type": "http.request", "body": b""}

        with patch.object(mw, "_stamp_with_retry") as mock_stamp:
            mock_stamp.return_value = StampResult(
                authorized=False,
                reason="Server unreachable",
                http_status=None,
            )

            await mw(scope, mock_receive, capture)

        start_msg = capture.messages[0]
        assert start_msg["type"] == "http.response.start"
        assert start_msg["status"] == 503

    @pytest.mark.asyncio
    async def test_explicit_false_allows_passthrough(self, mock_app):
        """Explicit block_on_rejection=False allows unstamped requests."""
        mw = LockStockThinMiddleware(
            mock_app,
            agent_id="test",
            api_key="test",
            path_tasks={"/v1/chat/completions": "chatcompletion"},
            block_on_rejection=False,
        )

        capture = MessageCapture()
        scope = {"type": "http", "method": "POST", "path": "/v1/chat/completions"}

        async def mock_receive():
            return {"type": "http.request", "body": b""}

        with patch.object(mw, "_stamp_with_retry") as mock_stamp:
            mock_stamp.return_value = StampResult(
                authorized=False,
                reason="Server unreachable",
                http_status=None,
            )

            await mw(scope, mock_receive, capture)

        # Should pass through to app, not 503
        start_msg = capture.messages[0]
        assert start_msg["type"] == "http.response.start"
        assert start_msg["status"] == 200

        headers = capture.get_headers()
        assert headers.get(b"x-lockstock-status") == b"skipped"


class TestExcludedPaths:
    """Test path exclusion."""

    @pytest.mark.asyncio
    async def test_excluded_path_no_stamp(self, middleware):
        """Excluded paths bypass stamping entirely."""
        scope = {"type": "http", "method": "GET", "path": "/health"}
        capture = MessageCapture()

        async def mock_receive():
            return {"type": "http.request", "body": b""}

        with patch.object(middleware, "_stamp_with_retry") as mock_stamp:
            await middleware(scope, mock_receive, capture)

        # _stamp_with_retry should not be called
        mock_stamp.assert_not_called()

        # Response should still be sent (pass-through)
        assert any(m["type"] == "http.response.start" for m in capture.messages)


class TestNonHttpRequests:
    """Test non-HTTP request handling."""

    @pytest.mark.asyncio
    async def test_websocket_passthrough(self, middleware):
        """WebSocket requests pass through unchanged."""
        scope = {"type": "websocket", "path": "/ws"}
        capture = MessageCapture()

        async def mock_receive():
            return {"type": "websocket.connect"}

        with patch.object(middleware, "_stamp_with_retry") as mock_stamp:
            # Mock app that handles websocket
            middleware.app = AsyncMock()
            await middleware(scope, mock_receive, capture)

        mock_stamp.assert_not_called()
        middleware.app.assert_called_once()


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
