#!/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
"""

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):
            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):
            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):
            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, middleware, http_scope):
        """Skipped stamp (not blocking) injects skip 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=False,
                reason="Connection error",
                http_status=None,
            )

            await middleware(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 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"])
