"""Integration tests for audit trail -- proxy dispatch logging and export API.

Verifies that:
1. Proxy requests generate audit entries in the database.
2. The audit export endpoint returns correct data in JSONL and CSV formats.
3. User tracking fields (api_key_id, user_id) are populated.
4. Enterprise plan gating on the export endpoint works correctly.
"""

from __future__ import annotations

import json
from datetime import datetime, timezone
from typing import TYPE_CHECKING
from unittest.mock import AsyncMock, patch

import pytest
from httpx import ASGITransport, AsyncClient

from llmhosts.audit.logger import AuditLogger
from llmhosts.audit.models import AuditEntry, AuditQuery, AuditSummary
from llmhosts.proxy.app import create_app
from llmhosts.proxy.dispatcher import BackendDispatcher, BackendEntry
from llmhosts.proxy.models import UnifiedMessage, UnifiedRequest, UnifiedResponse

if TYPE_CHECKING:
    from pathlib import Path

    from llmhosts.proxy.auth import ProxyKeyStore

# Rebuild Pydantic models with forward refs (same pattern as conftest.py)
AuditEntry.model_rebuild()
AuditQuery.model_rebuild()
AuditSummary.model_rebuild()


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture()
async def audit_logger(tmp_path: Path) -> AuditLogger:
    """Create and initialise a fresh AuditLogger for testing."""
    db_path = tmp_path / "test_audit.db"
    logger = AuditLogger(db_path)
    await logger.initialize()
    yield logger  # type: ignore[misc]
    await logger.close()


@pytest.fixture()
def dispatcher_with_audit(audit_logger: AuditLogger) -> BackendDispatcher:
    """Create a BackendDispatcher with a mock Core backend and audit logger attached."""
    dispatcher = BackendDispatcher()
    # Add a fake local backend so dispatch can find it
    fake_backend = BackendEntry(
        name="llmhosts-core",
        base_url="local://",
        backend_type="llmhosts-core",
        models=["llama-3-8b"],
        priority=-1,
        healthy=True,
    )
    dispatcher._backends.append(fake_backend)
    # Set up a mock Core engine so the dispatch path enters Core routing
    mock_engine = type("MockEngine", (), {"is_loaded": True})()
    dispatcher._core_engine = mock_engine
    dispatcher.attach_audit_logger(audit_logger)
    return dispatcher


@pytest.fixture()
def test_app(audit_logger: AuditLogger, dispatcher_with_audit: BackendDispatcher, tmp_path: Path):
    """Create a FastAPI app with audit logger and dispatcher pre-wired.

    Keys are created in the app's own key_store (which the AuthMiddleware uses)
    so that Bearer tokens are accepted by the middleware.
    """
    app = create_app()
    app.state.dispatcher = dispatcher_with_audit
    app.state.audit_logger = audit_logger

    # Use the app's own key_store (the one AuthMiddleware references)
    key_store: ProxyKeyStore = app.state.key_store
    _enterprise_key_id, enterprise_raw_key = key_store.create_key("enterprise-test", plan="enterprise")
    _free_key_id, free_raw_key = key_store.create_key("free-test", plan="free")

    # Store raw keys for use in tests
    app.state._test_enterprise_key = enterprise_raw_key
    app.state._test_free_key = free_raw_key
    app.state._test_enterprise_key_id = _enterprise_key_id
    app.state._test_free_key_id = _free_key_id

    return app


@pytest.fixture()
async def client(test_app) -> AsyncClient:
    """Async HTTP client wired to the test app."""
    transport = ASGITransport(app=test_app, raise_app_exceptions=False)
    async with AsyncClient(transport=transport, base_url="http://test") as ac:
        yield ac


# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------


def _make_unified_request(model: str = "llama-3-8b") -> UnifiedRequest:
    """Create a minimal UnifiedRequest for testing."""
    return UnifiedRequest(
        model=model,
        messages=[UnifiedMessage(role="user", content="Hello")],
        stream=False,
    )


def _make_unified_response() -> UnifiedResponse:
    """Create a mock UnifiedResponse."""
    return UnifiedResponse(
        content="Hello! How can I help you?",
        finish_reason="stop",
        model="llama-3-8b",
        prompt_tokens=10,
        completion_tokens=8,
        total_tokens=18,
    )


def _make_core_response() -> tuple[dict, dict[str, str]]:
    """Create a mock _dispatch_core return value (body, headers)."""
    body = {
        "id": "chatcmpl-test",
        "object": "chat.completion",
        "created": 1000000,
        "model": "llama-3-8b",
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": "Hello! How can I help you?"},
                "finish_reason": "stop",
            }
        ],
        "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18},
    }
    headers = {
        "X-LLMHosts-Backend": "llmhosts-core",
        "X-LLMHosts-Model": "llama-3-8b",
        "X-LLMHosts-Cost": "0.0",
    }
    return body, headers


# ===================================================================
# Test: Proxy requests generate audit entries
# ===================================================================


class TestAuditDispatchIntegration:
    """Test that dispatcher.dispatch() writes audit entries."""

    async def test_dispatch_creates_audit_entry(
        self,
        dispatcher_with_audit: BackendDispatcher,
        audit_logger: AuditLogger,
    ) -> None:
        """A successful dispatch should write one audit entry."""
        request = _make_unified_request()

        with patch.object(
            dispatcher_with_audit, "_dispatch_core", new_callable=AsyncMock, return_value=_make_core_response()
        ):
            await dispatcher_with_audit.dispatch(
                request,
                source_ip="192.168.1.100",
                request_id="test-req-001",
                api_key_id="abc12345...",
                user_id="user-42",
                dialect="openai",
                path="/v1/chat/completions",
            )

        entries = await audit_logger.query(AuditQuery(limit=10))
        assert len(entries) == 1

        entry = entries[0]
        assert entry.request_id == "test-req-001"
        assert entry.source_ip == "192.168.1.100"
        assert entry.model_requested == "llama-3-8b"
        assert entry.model_used == "llama-3-8b"
        assert entry.backend_type == "llmhosts-core"
        assert entry.status_code == 200
        assert entry.input_tokens == 10
        assert entry.output_tokens == 8
        assert entry.latency_ms > 0

    async def test_dispatch_records_user_tracking_fields(
        self,
        dispatcher_with_audit: BackendDispatcher,
        audit_logger: AuditLogger,
    ) -> None:
        """api_key_id and user_id should be recorded in audit entries."""
        request = _make_unified_request()

        with patch.object(
            dispatcher_with_audit, "_dispatch_core", new_callable=AsyncMock, return_value=_make_core_response()
        ):
            await dispatcher_with_audit.dispatch(
                request,
                api_key_id="lmh_abcd...",
                user_id="user-enterprise-99",
            )

        entries = await audit_logger.query(AuditQuery(limit=10))
        assert len(entries) == 1
        assert entries[0].api_key_id == "lmh_abcd..."
        assert entries[0].user_id == "user-enterprise-99"

    async def test_dispatch_records_error_on_failure(
        self,
        dispatcher_with_audit: BackendDispatcher,
        audit_logger: AuditLogger,
    ) -> None:
        """Failed dispatch should still record an audit entry with error info."""
        request = _make_unified_request()

        with (
            patch.object(
                dispatcher_with_audit,
                "_dispatch_core",
                new_callable=AsyncMock,
                side_effect=Exception("Backend exploded"),
            ),
            pytest.raises(Exception, match="Backend exploded"),
        ):
            await dispatcher_with_audit.dispatch(
                request,
                request_id="test-err-001",
            )

        entries = await audit_logger.query(AuditQuery(limit=10))
        assert len(entries) == 1
        assert entries[0].request_id == "test-err-001"
        assert entries[0].status_code == 502
        assert "Backend exploded" in (entries[0].error or "")

    async def test_multiple_dispatches_create_multiple_entries(
        self,
        dispatcher_with_audit: BackendDispatcher,
        audit_logger: AuditLogger,
    ) -> None:
        """Multiple dispatches should create multiple audit entries."""
        with patch.object(
            dispatcher_with_audit, "_dispatch_core", new_callable=AsyncMock, return_value=_make_core_response()
        ):
            for i in range(5):
                request = _make_unified_request()
                await dispatcher_with_audit.dispatch(request, request_id=f"test-multi-{i:03d}")

        entries = await audit_logger.query(AuditQuery(limit=100))
        assert len(entries) == 5
        request_ids = {e.request_id for e in entries}
        assert request_ids == {f"test-multi-{i:03d}" for i in range(5)}

    async def test_dispatch_without_audit_logger_no_error(self) -> None:
        """Dispatcher without audit logger should work without errors."""
        dispatcher = BackendDispatcher()
        fake_backend = BackendEntry(
            name="llmhosts-core",
            base_url="local://",
            backend_type="llmhosts-core",
            models=["llama-3-8b"],
            priority=-1,
            healthy=True,
        )
        dispatcher._backends.append(fake_backend)
        # Set up mock Core engine
        mock_engine = type("MockEngine", (), {"is_loaded": True})()
        dispatcher._core_engine = mock_engine
        # No audit logger attached

        request = _make_unified_request()

        with patch.object(dispatcher, "_dispatch_core", new_callable=AsyncMock, return_value=_make_core_response()):
            result = await dispatcher.dispatch(request)

        assert result.content == "Hello! How can I help you?"

    async def test_dispatch_records_dialect_and_path(
        self,
        dispatcher_with_audit: BackendDispatcher,
        audit_logger: AuditLogger,
    ) -> None:
        """dialect and path fields should be recorded correctly."""
        request = _make_unified_request()

        with patch.object(
            dispatcher_with_audit, "_dispatch_core", new_callable=AsyncMock, return_value=_make_core_response()
        ):
            await dispatcher_with_audit.dispatch(
                request,
                request_id="test-dialect-001",
                dialect="anthropic",
                path="/v1/messages",
            )

        entries = await audit_logger.query(AuditQuery(limit=10))
        assert len(entries) == 1
        assert entries[0].dialect == "anthropic"
        assert entries[0].path == "/v1/messages"


# ===================================================================
# Test: Audit export endpoint
# ===================================================================


class TestAuditExportEndpoint:
    """Test the GET /v1/audit/export endpoint."""

    async def _seed_audit_entries(self, audit_logger: AuditLogger, count: int = 5) -> list[str]:
        """Insert test audit entries and return their request IDs."""
        request_ids: list[str] = []
        for i in range(count):
            entry = AuditEntry(
                request_id=f"export-req-{i:03d}",
                timestamp=datetime(2026, 2, 20, 10, i, tzinfo=timezone.utc),
                source_ip="10.0.0.1",
                method="POST",
                path="/v1/chat/completions",
                dialect="openai",
                model_requested="llama-3-8b",
                model_used="llama-3-8b",
                backend_type="ollama",
                routing_tier="rule",
                routing_reasoning="",
                routing_confidence=0.9,
                input_tokens=100 + i,
                output_tokens=50 + i,
                actual_cost=0.001 * i,
                cloud_equivalent_cost=0.01 * i,
                cached=False,
                latency_ms=50.0 + i,
                status_code=200,
                api_key_id="lmh_test...",
                user_id="user-test",
            )
            await audit_logger.log(entry)
            request_ids.append(entry.request_id)
        return request_ids

    async def test_export_requires_enterprise_plan(self, client: AsyncClient, test_app) -> None:
        """Free plan users should get 403 on audit export."""
        free_key = test_app.state._test_free_key
        resp = await client.get(
            "/v1/audit/export",
            headers={"Authorization": f"Bearer {free_key}"},
        )
        assert resp.status_code == 403
        body = resp.json()
        assert body["error"]["code"] == "enterprise_required"

    async def test_export_anonymous_gets_403(self, client: AsyncClient) -> None:
        """Anonymous (no auth) users should get 403."""
        resp = await client.get("/v1/audit/export")
        assert resp.status_code == 403
        body = resp.json()
        assert body["error"]["code"] == "enterprise_required"

    async def test_export_jsonl_enterprise(self, client: AsyncClient, test_app, audit_logger: AuditLogger) -> None:
        """Enterprise plan users should get JSONL audit data."""
        await self._seed_audit_entries(audit_logger, count=3)

        enterprise_key = test_app.state._test_enterprise_key
        resp = await client.get(
            "/v1/audit/export?format=jsonl",
            headers={"Authorization": f"Bearer {enterprise_key}"},
        )
        assert resp.status_code == 200
        assert "application/x-ndjson" in resp.headers.get("content-type", "")

        lines = resp.text.strip().split("\n")
        assert len(lines) == 3

        # Each line is valid JSON
        for line in lines:
            parsed = json.loads(line)
            assert "request_id" in parsed
            assert "timestamp" in parsed
            assert "api_key_id" in parsed
            assert "user_id" in parsed

    async def test_export_csv_enterprise(self, client: AsyncClient, test_app, audit_logger: AuditLogger) -> None:
        """Enterprise plan users should get CSV audit data."""
        await self._seed_audit_entries(audit_logger, count=3)

        enterprise_key = test_app.state._test_enterprise_key
        resp = await client.get(
            "/v1/audit/export?format=csv",
            headers={"Authorization": f"Bearer {enterprise_key}"},
        )
        assert resp.status_code == 200
        assert "text/csv" in resp.headers.get("content-type", "")

        lines = resp.text.strip().split("\n")
        # Header + 3 data rows
        assert len(lines) == 4
        header = lines[0]
        assert "request_id" in header
        assert "api_key_id" in header

    async def test_export_with_model_filter(self, client: AsyncClient, test_app, audit_logger: AuditLogger) -> None:
        """Model filter should limit results."""
        # Seed entries with different models
        for i, model in enumerate(["gpt-4", "llama-3-8b", "gpt-4", "llama-3-8b", "gpt-4"]):
            entry = AuditEntry(
                request_id=f"filter-req-{i:03d}",
                timestamp=datetime(2026, 2, 20, 10, i, tzinfo=timezone.utc),
                source_ip="10.0.0.1",
                method="POST",
                path="/v1/chat/completions",
                dialect="openai",
                model_requested=model,
                model_used=model,
                backend_type="ollama",
                routing_tier="rule",
                routing_reasoning="",
                routing_confidence=0.9,
                input_tokens=100,
                output_tokens=50,
                actual_cost=0.001,
                cloud_equivalent_cost=0.01,
                cached=False,
                latency_ms=50.0,
                status_code=200,
                api_key_id="lmh_test...",
                user_id="user-test",
            )
            await audit_logger.log(entry)

        enterprise_key = test_app.state._test_enterprise_key
        resp = await client.get(
            "/v1/audit/export?format=jsonl&model=gpt-4",
            headers={"Authorization": f"Bearer {enterprise_key}"},
        )
        assert resp.status_code == 200
        lines = resp.text.strip().split("\n")
        assert len(lines) == 3  # 3 gpt-4 entries

    async def test_export_with_since_filter(self, client: AsyncClient, test_app, audit_logger: AuditLogger) -> None:
        """Since filter should only return entries after the given time."""
        await self._seed_audit_entries(audit_logger, count=5)
        # Entries are at 10:00, 10:01, 10:02, 10:03, 10:04
        # since=10:02 should return 10:02, 10:03, 10:04 = 3 entries

        enterprise_key = test_app.state._test_enterprise_key
        resp = await client.get(
            "/v1/audit/export?format=jsonl&since=2026-02-20T10:02:00%2B00:00",
            headers={"Authorization": f"Bearer {enterprise_key}"},
        )
        assert resp.status_code == 200
        lines = resp.text.strip().split("\n")
        assert len(lines) == 3

    async def test_export_with_limit(self, client: AsyncClient, test_app, audit_logger: AuditLogger) -> None:
        """Limit parameter should cap the number of results."""
        await self._seed_audit_entries(audit_logger, count=10)

        enterprise_key = test_app.state._test_enterprise_key
        resp = await client.get(
            "/v1/audit/export?format=jsonl&limit=3",
            headers={"Authorization": f"Bearer {enterprise_key}"},
        )
        assert resp.status_code == 200
        lines = resp.text.strip().split("\n")
        assert len(lines) == 3

    async def test_export_empty_result(self, client: AsyncClient, test_app) -> None:
        """Export with no entries returns empty body."""
        enterprise_key = test_app.state._test_enterprise_key
        resp = await client.get(
            "/v1/audit/export?format=jsonl",
            headers={"Authorization": f"Bearer {enterprise_key}"},
        )
        assert resp.status_code == 200
        # Empty result -- no data lines
        assert resp.text.strip() == ""


# ===================================================================
# Test: User tracking fields populated
# ===================================================================


class TestUserTrackingFields:
    """Test that api_key_id and user_id are stored and retrieved correctly."""

    async def test_api_key_id_stored_and_queryable(self, audit_logger: AuditLogger) -> None:
        """api_key_id field persists through write and read cycle."""
        entry = AuditEntry(
            request_id="key-test-001",
            timestamp=datetime(2026, 2, 20, tzinfo=timezone.utc),
            source_ip="10.0.0.1",
            method="POST",
            path="/v1/chat/completions",
            dialect="openai",
            model_requested="llama-3-8b",
            model_used="llama-3-8b",
            backend_type="ollama",
            routing_tier="rule",
            routing_reasoning="",
            routing_confidence=0.9,
            input_tokens=100,
            output_tokens=50,
            actual_cost=0.001,
            cloud_equivalent_cost=0.01,
            cached=False,
            latency_ms=50.0,
            status_code=200,
            api_key_id="lmh_abcd...",
            user_id="",
        )
        await audit_logger.log(entry)

        entries = await audit_logger.query(AuditQuery(limit=10))
        assert len(entries) == 1
        assert entries[0].api_key_id == "lmh_abcd..."

    async def test_user_id_stored_and_queryable(self, audit_logger: AuditLogger) -> None:
        """user_id field persists through write and read cycle."""
        entry = AuditEntry(
            request_id="user-test-001",
            timestamp=datetime(2026, 2, 20, tzinfo=timezone.utc),
            source_ip="10.0.0.1",
            method="POST",
            path="/v1/chat/completions",
            dialect="openai",
            model_requested="llama-3-8b",
            model_used="llama-3-8b",
            backend_type="ollama",
            routing_tier="rule",
            routing_reasoning="",
            routing_confidence=0.9,
            input_tokens=100,
            output_tokens=50,
            actual_cost=0.001,
            cloud_equivalent_cost=0.01,
            cached=False,
            latency_ms=50.0,
            status_code=200,
            api_key_id="",
            user_id="user-42",
        )
        await audit_logger.log(entry)

        entries = await audit_logger.query(AuditQuery(limit=10))
        assert len(entries) == 1
        assert entries[0].user_id == "user-42"

    async def test_default_empty_tracking_fields(self, audit_logger: AuditLogger) -> None:
        """When not provided, api_key_id and user_id default to empty string."""
        entry = AuditEntry(
            request_id="default-test-001",
            timestamp=datetime(2026, 2, 20, tzinfo=timezone.utc),
            source_ip="10.0.0.1",
            method="POST",
            path="/v1/chat/completions",
            dialect="openai",
            model_requested="llama-3-8b",
            model_used="llama-3-8b",
            backend_type="ollama",
            routing_tier="rule",
            routing_reasoning="",
            routing_confidence=0.9,
            input_tokens=100,
            output_tokens=50,
            actual_cost=0.001,
            cloud_equivalent_cost=0.01,
            cached=False,
            latency_ms=50.0,
            status_code=200,
        )
        await audit_logger.log(entry)

        entries = await audit_logger.query(AuditQuery(limit=10))
        assert len(entries) == 1
        assert entries[0].api_key_id == ""
        assert entries[0].user_id == ""

    async def test_tracking_fields_in_model_dump(self) -> None:
        """api_key_id and user_id appear in the serialized model."""
        entry = AuditEntry(
            request_id="dump-test",
            timestamp=datetime(2026, 2, 20, tzinfo=timezone.utc),
            source_ip="10.0.0.1",
            method="POST",
            path="/v1/chat/completions",
            dialect="openai",
            model_requested="test",
            model_used="test",
            backend_type="ollama",
            routing_tier="rule",
            routing_reasoning="",
            routing_confidence=0.9,
            input_tokens=0,
            output_tokens=0,
            actual_cost=0.0,
            cloud_equivalent_cost=0.0,
            cached=False,
            latency_ms=0.0,
            status_code=200,
            api_key_id="lmh_1234...",
            user_id="user-test",
        )
        data = entry.model_dump(mode="json")
        assert data["api_key_id"] == "lmh_1234..."
        assert data["user_id"] == "user-test"


# ===================================================================
# Test: Enterprise plan gating
# ===================================================================


class TestEnterprisePlanGating:
    """Dedicated tests for the enterprise plan check on audit export."""

    async def test_free_plan_blocked(self, client: AsyncClient, test_app) -> None:
        """Free plan gets 403."""
        free_key = test_app.state._test_free_key
        resp = await client.get(
            "/v1/audit/export",
            headers={"Authorization": f"Bearer {free_key}"},
        )
        assert resp.status_code == 403
        assert resp.json()["error"]["code"] == "enterprise_required"

    async def test_enterprise_plan_allowed(self, client: AsyncClient, test_app) -> None:
        """Enterprise plan gets 200 (even with empty data)."""
        enterprise_key = test_app.state._test_enterprise_key
        resp = await client.get(
            "/v1/audit/export",
            headers={"Authorization": f"Bearer {enterprise_key}"},
        )
        assert resp.status_code == 200

    async def test_no_auth_treated_as_free(self, client: AsyncClient) -> None:
        """Unauthenticated requests are treated as free plan and blocked."""
        resp = await client.get("/v1/audit/export")
        assert resp.status_code == 403
