"""Integration tests for the AttestMCP proxy.

These tests verify the core product guarantees:
1. Zero behavioral change — upstream responses pass through unchanged
2. Audit completeness — every proxied request produces an audit record
3. Chain integrity — the hash chain is valid after all operations
4. Identity extraction — agent identity appears on every audit record
"""

from __future__ import annotations

import base64
import json
from typing import TYPE_CHECKING

import httpx
import pytest
from fastapi import FastAPI, Request, Response
from fastapi.testclient import TestClient

from attestmcp.config import Settings
from attestmcp.identity import IdentityExtractor
from attestmcp.proxy import create_app
from attestmcp.store import SQLiteAuditStore

if TYPE_CHECKING:
    from pathlib import Path


# ──────────────────────────────────────────────
# Mock MCP Memory Server
# ──────────────────────────────────────────────

def create_mock_memory_server() -> FastAPI:
    """A minimal MCP memory server that responds to standard methods."""
    mock = FastAPI()

    @mock.post("/mcp")
    async def handle_mcp(request: Request) -> Response:
        body = await request.body()
        parsed = json.loads(body)
        method = parsed.get("method", "")
        request_id = parsed.get("id")

        if method == "initialize":
            return Response(
                content=json.dumps({
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": {
                        "protocolVersion": "2025-11-25",
                        "capabilities": {"tools": {}},
                        "serverInfo": {
                            "name": "mock-memory",
                            "version": "1.0.0",
                        },
                    },
                }),
                media_type="application/json",
            )

        if method == "tools/list":
            return Response(
                content=json.dumps({
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": {
                        "tools": [
                            {
                                "name": "store_memory",
                                "description": "Store a memory",
                                "inputSchema": {
                                    "type": "object",
                                    "properties": {
                                        "content": {"type": "string"},
                                    },
                                },
                            },
                            {
                                "name": "search_memories",
                                "description": "Search memories",
                                "inputSchema": {
                                    "type": "object",
                                    "properties": {
                                        "query": {"type": "string"},
                                    },
                                },
                            },
                        ],
                    },
                }),
                media_type="application/json",
            )

        if method == "tools/call":
            tool_name = parsed.get("params", {}).get("name", "unknown")
            arguments = parsed.get("params", {}).get("arguments", {})

            if tool_name == "store_memory":
                return Response(
                    content=json.dumps({
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "result": {
                            "content": [
                                {
                                    "type": "text",
                                    "text": f"Stored: {arguments.get('content', '')}",
                                }
                            ],
                        },
                    }),
                    media_type="application/json",
                )

            if tool_name == "search_memories":
                return Response(
                    content=json.dumps({
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "result": {
                            "content": [
                                {
                                    "type": "text",
                                    "text": "Found 2 relevant memories",
                                }
                            ],
                        },
                    }),
                    media_type="application/json",
                )

        # Unknown method
        return Response(
            content=json.dumps({
                "jsonrpc": "2.0",
                "id": request_id,
                "error": {
                    "code": -32601,
                    "message": f"Method not found: {method}",
                },
            }),
            media_type="application/json",
            status_code=200,
        )

    return mock


# ──────────────────────────────────────────────
# Test Fixtures
# ──────────────────────────────────────────────

@pytest.fixture
def mock_transport() -> httpx.ASGITransport:
    """Transport that routes requests to the mock memory server."""
    mock_app = create_mock_memory_server()
    return httpx.ASGITransport(app=mock_app)  # type: ignore[arg-type]


@pytest.fixture
def store(tmp_path: Path) -> SQLiteAuditStore:
    """Audit store backed by a temp database."""
    return SQLiteAuditStore(db_path=str(tmp_path / "test_audit.db"))


@pytest.fixture
def proxy_app(
    tmp_path: Path, mock_transport: httpx.ASGITransport, store: SQLiteAuditStore
) -> FastAPI:
    """Proxy app with injected mock transport and temp store."""
    settings = Settings(
        upstream_url="http://mock-memory/mcp",
        db_path=str(tmp_path / "test_audit.db"),
        jwks_url=None,
    )
    mock_http_client = httpx.AsyncClient(transport=mock_transport)
    return create_app(
        settings,
        store=store,
        identity_extractor=IdentityExtractor(),
        http_client=mock_http_client,
    )


@pytest.fixture
def client(proxy_app: FastAPI) -> TestClient:
    """Test client for the proxy."""
    with TestClient(proxy_app) as c:
        yield c


# ──────────────────────────────────────────────
# Tests: Transparency
# ──────────────────────────────────────────────

class TestTransparency:
    """The proxy must not modify request or response payloads."""

    def test_initialize_passes_through(self, client: TestClient) -> None:
        response = client.post(
            "/mcp",
            json={
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-11-25",
                    "capabilities": {},
                    "clientInfo": {"name": "test", "version": "1.0"},
                },
            },
        )
        assert response.status_code == 200
        data = response.json()
        assert data["result"]["serverInfo"]["name"] == "mock-memory"

    def test_tools_list_passes_through(self, client: TestClient) -> None:
        response = client.post(
            "/mcp",
            json={"jsonrpc": "2.0", "id": 2, "method": "tools/list"},
        )
        assert response.status_code == 200
        data = response.json()
        tools = data["result"]["tools"]
        assert len(tools) == 2
        assert tools[0]["name"] == "store_memory"
        assert tools[1]["name"] == "search_memories"

    def test_tools_call_store_memory(self, client: TestClient) -> None:
        response = client.post(
            "/mcp",
            json={
                "jsonrpc": "2.0",
                "id": 3,
                "method": "tools/call",
                "params": {
                    "name": "store_memory",
                    "arguments": {"content": "The auth uses JWT with refresh tokens"},
                },
            },
        )
        assert response.status_code == 200
        data = response.json()
        assert "Stored:" in data["result"]["content"][0]["text"]

    def test_tools_call_search_memories(self, client: TestClient) -> None:
        response = client.post(
            "/mcp",
            json={
                "jsonrpc": "2.0",
                "id": 4,
                "method": "tools/call",
                "params": {
                    "name": "search_memories",
                    "arguments": {"query": "authentication"},
                },
            },
        )
        assert response.status_code == 200
        data = response.json()
        assert "Found 2" in data["result"]["content"][0]["text"]

    def test_error_response_passes_through(self, client: TestClient) -> None:
        response = client.post(
            "/mcp",
            json={
                "jsonrpc": "2.0",
                "id": 5,
                "method": "nonexistent/method",
            },
        )
        assert response.status_code == 200
        data = response.json()
        assert "error" in data
        assert data["error"]["code"] == -32601


# ──────────────────────────────────────────────
# Tests: Audit Completeness
# ──────────────────────────────────────────────

class TestAuditCompleteness:
    """Every proxied request must produce exactly one audit record."""

    def test_single_request_creates_one_record(
        self, client: TestClient, store: SQLiteAuditStore
    ) -> None:
        initial_count = store.count()
        client.post(
            "/mcp",
            json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
        )
        assert store.count() == initial_count + 1

    def test_multiple_requests_create_multiple_records(
        self, client: TestClient, store: SQLiteAuditStore
    ) -> None:
        initial_count = store.count()

        for i in range(5):
            client.post(
                "/mcp",
                json={
                    "jsonrpc": "2.0",
                    "id": i + 10,
                    "method": "tools/call",
                    "params": {
                        "name": "store_memory",
                        "arguments": {"content": f"memory {i}"},
                    },
                },
            )

        assert store.count() == initial_count + 5

    def test_audit_record_captures_method(
        self, client: TestClient, store: SQLiteAuditStore
    ) -> None:
        client.post(
            "/mcp",
            json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
        )
        records = store.query()
        tools_list_records = [r for r in records if r.method == "tools/list"]
        assert len(tools_list_records) >= 1

    def test_audit_record_captures_tool_name(
        self, client: TestClient, store: SQLiteAuditStore
    ) -> None:
        client.post(
            "/mcp",
            json={
                "jsonrpc": "2.0",
                "id": 1,
                "method": "tools/call",
                "params": {
                    "name": "search_memories",
                    "arguments": {"query": "test"},
                },
            },
        )
        records = store.query()
        search_records = [r for r in records if r.tool_name == "search_memories"]
        assert len(search_records) >= 1

    def test_audit_record_has_result_hash(
        self, client: TestClient, store: SQLiteAuditStore
    ) -> None:
        client.post(
            "/mcp",
            json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
        )
        records = store.query()
        latest = records[-1]
        assert latest.result_hash is not None
        assert len(latest.result_hash) == 64

    def test_audit_record_has_duration(
        self, client: TestClient, store: SQLiteAuditStore
    ) -> None:
        client.post(
            "/mcp",
            json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
        )
        records = store.query()
        latest = records[-1]
        assert latest.duration_ms >= 0


# ──────────────────────────────────────────────
# Tests: Chain Integrity After Proxy Operations
# ──────────────────────────────────────────────

class TestChainIntegrity:
    """The hash chain must be valid after all proxy operations."""

    def test_chain_valid_after_single_request(
        self, client: TestClient, store: SQLiteAuditStore
    ) -> None:
        client.post(
            "/mcp",
            json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
        )
        result = store.verify_chain()
        assert result.valid is True

    def test_chain_valid_after_mixed_operations(
        self, client: TestClient, store: SQLiteAuditStore
    ) -> None:
        """Run a realistic sequence of operations and verify the chain."""
        # Initialize
        client.post(
            "/mcp",
            json={
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-11-25",
                    "capabilities": {},
                    "clientInfo": {"name": "test-agent", "version": "1.0"},
                },
            },
        )

        # List tools
        client.post(
            "/mcp",
            json={"jsonrpc": "2.0", "id": 2, "method": "tools/list"},
        )

        # Store some memories
        for i in range(3):
            client.post(
                "/mcp",
                json={
                    "jsonrpc": "2.0",
                    "id": 10 + i,
                    "method": "tools/call",
                    "params": {
                        "name": "store_memory",
                        "arguments": {"content": f"Decision #{i}: use PostgreSQL"},
                    },
                },
            )

        # Search memories
        client.post(
            "/mcp",
            json={
                "jsonrpc": "2.0",
                "id": 20,
                "method": "tools/call",
                "params": {
                    "name": "search_memories",
                    "arguments": {"query": "database decision"},
                },
            },
        )

        # Verify the entire chain
        result = store.verify_chain()
        assert result.valid is True
        assert result.records_checked >= 6


# ──────────────────────────────────────────────
# Tests: Identity on Audit Records
# ──────────────────────────────────────────────

class TestIdentityAudit:
    """Agent identity must appear on every audit record."""

    def test_anonymous_identity_when_no_auth(
        self, client: TestClient, store: SQLiteAuditStore
    ) -> None:
        client.post(
            "/mcp",
            json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
        )
        records = store.query()
        latest = records[-1]
        assert latest.agent_identity.source.value == "anonymous"
        assert latest.agent_identity.trust_level.value == "anonymous"

    def test_claimed_identity_with_jwt(
        self, client: TestClient, store: SQLiteAuditStore
    ) -> None:
        # Create a minimal unsigned JWT
        header = base64.urlsafe_b64encode(
            b'{"alg":"none","typ":"JWT"}'
        ).rstrip(b"=")
        payload = base64.urlsafe_b64encode(
            b'{"email":"agent@acme.com","azp":"support-bot","org":"acme"}'
        ).rstrip(b"=")
        token = f"{header.decode()}.{payload.decode()}."

        client.post(
            "/mcp",
            json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
            headers={"Authorization": f"Bearer {token}"},
        )

        records = store.query()
        latest = records[-1]
        assert latest.agent_identity.source.value == "oauth_jwt"
        assert latest.agent_identity.user == "agent@acme.com"
        assert latest.agent_identity.agent_name == "support-bot"
        assert latest.agent_identity.org == "acme"
        assert latest.agent_identity.trust_level.value == "claimed"

    def test_session_id_captured(
        self, client: TestClient, store: SQLiteAuditStore
    ) -> None:
        client.post(
            "/mcp",
            json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
            headers={"Mcp-Session-Id": "session-abc-123"},
        )
        records = store.query()
        latest = records[-1]
        assert latest.agent_identity.session_id == "session-abc-123"
        assert latest.session_id == "session-abc-123"


# ──────────────────────────────────────────────
# Tests: Health Endpoint
# ──────────────────────────────────────────────

class TestHealth:
    """Health check returns proxy status."""

    def test_health_check(self, client: TestClient) -> None:
        response = client.get("/health")
        assert response.status_code == 200
        data = response.json()
        assert data["status"] == "ok"
        assert "upstream" in data
        assert "audit_records" in data
        assert "chain_head" in data
