"""Integration tests for POST /query endpoint."""

import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import uuid4

import pytest
from fastapi.testclient import TestClient
from openai.types.chat import ChatCompletion, ChatCompletionMessage
from openai.types.chat.chat_completion import Choice
from sqlalchemy import select

from src.api.main import app
from src.db.models import ApiKey, Document
from src.db.session import get_sync_session
from src.services.auth_service import hash_api_key


@pytest.fixture
def test_api_key():
    """Create a test API key in the database."""
    test_key = "test_api_key_for_query_12345"
    key_hash = hash_api_key(test_key)

    with get_sync_session() as session:
        # Clean up any existing test API key
        existing = session.execute(
            select(ApiKey).where(ApiKey.key_hash == key_hash)
        ).scalar_one_or_none()
        if existing:
            session.delete(existing)
            session.commit()

        # Create new test API key
        api_key = ApiKey(
            key_hash=key_hash,
            name="Test Query API Key",
            rate_limit_per_minute=60,
            is_active=True,
        )
        session.add(api_key)
        session.commit()
        session.refresh(api_key)

        yield test_key, api_key.id

        # Cleanup: delete test API key
        if api_key_obj := session.execute(
            select(ApiKey).where(ApiKey.key_hash == key_hash)
        ).scalar_one_or_none():
            session.delete(api_key_obj)
            session.commit()


@pytest.fixture
def test_document():
    """Create a test completed document in the database."""
    doc_id = uuid4()

    with get_sync_session() as session:
        # Create test document
        document = Document(
            id=doc_id,
            filename="test_document.xlsx",
            status="completed",
            file_size=1024,
            mime_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        )
        session.add(document)
        session.commit()

        yield doc_id

        # Cleanup: delete test document
        if doc_obj := session.execute(
            select(Document).where(Document.id == doc_id)
        ).scalar_one_or_none():
            session.delete(doc_obj)
            session.commit()


@pytest.fixture
def test_client():
    """Create a FastAPI test client."""
    return TestClient(app)


@pytest.fixture
def mock_milvus_search():
    """Mock Milvus search results."""
    with patch("src.retrieval.semantic_retriever.connections") as mock_conn:
        with patch("src.retrieval.semantic_retriever.Collection") as mock_coll_class:
            # Create mock collection
            mock_collection = MagicMock()
            mock_collection.load = MagicMock()

            # Create mock search result
            mock_hit = MagicMock()
            mock_hit.distance = 0.85
            mock_hit.entity.get = lambda key, default="": {
                "filename": "test_document.xlsx",
                "sheet_name": "Sheet1",
                "col_range": "A1:D1",
                "text_content": "Item Code: 2024-4_0019, Status: Active",
            }.get(key, default)

            mock_collection.search.return_value = [[mock_hit]]
            mock_coll_class.return_value = mock_collection

            yield mock_collection


@pytest.fixture
def mock_azure_openai():
    """Mock Azure OpenAI response."""
    with patch("src.retrieval.answer_generator.AsyncAzureOpenAI") as mock_class:
        mock_client = MagicMock()

        # Create mock response
        mock_message = ChatCompletionMessage(
            role="assistant",
            content="Item code 2024-4_0019 has status Active as shown in Sheet1.",
        )
        mock_choice = Choice(index=0, message=mock_message, finish_reason="stop")
        mock_completion = ChatCompletion(
            id="test-completion",
            model="gpt-4",
            object="chat.completion",
            created=1234567890,
            choices=[mock_choice],
        )

        mock_client.chat.completions.create = AsyncMock(return_value=mock_completion)
        mock_class.return_value = mock_client

        yield mock_client


@pytest.fixture
def mock_embeddings():
    """Mock embedding service."""
    with patch("src.retrieval.semantic_retriever.EmbeddingService") as mock_class:
        mock_service = MagicMock()
        mock_service.embed_batch = AsyncMock(return_value=[[0.1] * 3072])
        mock_class.return_value = mock_service
        yield mock_service


class TestQueryEndpoint:
    """Tests for POST /query endpoint."""

    def test_query_valid_request_returns_200(
        self,
        test_client,
        test_api_key,
        mock_milvus_search,
        mock_azure_openai,
        mock_embeddings,
    ):
        """Test successful query returns 200 with answer and sources."""
        test_key, _ = test_api_key

        # Send query request
        response = test_client.post(
            "/query",
            json={
                "query": "What is item code 2024-4_0019?",
                "include_context": True,
            },
            headers={"X-API-Key": test_key},
        )

        # Verify response
        assert response.status_code == 200
        data = response.json()

        assert "query" in data
        assert data["query"] == "What is item code 2024-4_0019?"
        assert "answer" in data
        assert "2024-4_0019" in data["answer"]
        assert "confidence" in data
        assert 0.0 <= data["confidence"] <= 1.0
        assert "sources" in data
        assert len(data["sources"]) == 1
        assert data["sources"][0]["file"] == "test_document.xlsx"
        assert "processing_time_ms" in data
        assert data["processing_time_ms"] > 0

    def test_query_with_document_ids_filter(
        self,
        test_client,
        test_api_key,
        test_document,
        mock_milvus_search,
        mock_azure_openai,
        mock_embeddings,
    ):
        """Test query with document_ids filter."""
        test_key, _ = test_api_key

        # Send query with document_ids filter
        response = test_client.post(
            "/query",
            json={
                "query": "test query",
                "document_ids": [str(test_document)],
                "include_context": True,
            },
            headers={"X-API-Key": test_key},
        )

        # Verify response
        assert response.status_code == 200
        data = response.json()
        assert "answer" in data
        assert "sources" in data

    def test_query_without_context_returns_no_sources(
        self,
        test_client,
        test_api_key,
        mock_milvus_search,
        mock_azure_openai,
        mock_embeddings,
    ):
        """Test query with include_context=False returns no sources."""
        test_key, _ = test_api_key

        # Send query without context
        response = test_client.post(
            "/query",
            json={
                "query": "test query",
                "include_context": False,
            },
            headers={"X-API-Key": test_key},
        )

        # Verify response
        assert response.status_code == 200
        data = response.json()
        assert "answer" in data
        assert "sources" in data
        assert len(data["sources"]) == 0

    def test_query_empty_string_returns_400(self, test_client, test_api_key):
        """Test empty query returns 400 validation error."""
        test_key, _ = test_api_key

        # Send empty query
        response = test_client.post(
            "/query",
            json={"query": ""},
            headers={"X-API-Key": test_key},
        )

        # Verify error response
        assert response.status_code == 422  # Pydantic validation error

    def test_query_whitespace_only_returns_400(self, test_client, test_api_key):
        """Test whitespace-only query returns 400 validation error."""
        test_key, _ = test_api_key

        # Send whitespace query
        response = test_client.post(
            "/query",
            json={"query": "   "},
            headers={"X-API-Key": test_key},
        )

        # Verify error response
        assert response.status_code == 422

    def test_query_too_long_returns_400(self, test_client, test_api_key):
        """Test query exceeding 1000 chars returns 400 validation error."""
        test_key, _ = test_api_key

        # Send query exceeding max length
        long_query = "a" * 1001
        response = test_client.post(
            "/query",
            json={"query": long_query},
            headers={"X-API-Key": test_key},
        )

        # Verify error response
        assert response.status_code == 422

    def test_query_missing_api_key_returns_401(self, test_client):
        """Test missing API key returns 401 unauthorized."""
        # Send query without API key
        response = test_client.post(
            "/query",
            json={"query": "test query"},
        )

        # Verify error response
        assert response.status_code == 403  # API key dependency raises 403

    def test_query_invalid_api_key_returns_401(self, test_client):
        """Test invalid API key returns 401 unauthorized."""
        # Send query with invalid API key
        response = test_client.post(
            "/query",
            json={"query": "test query"},
            headers={"X-API-Key": "invalid_key_12345"},
        )

        # Verify error response
        assert response.status_code == 403

    def test_query_invalid_document_id_returns_404(
        self,
        test_client,
        test_api_key,
        mock_embeddings,
    ):
        """Test query with nonexistent document_id returns 404."""
        test_key, _ = test_api_key

        # Send query with invalid document ID
        fake_doc_id = str(uuid4())
        response = test_client.post(
            "/query",
            json={
                "query": "test query",
                "document_ids": [fake_doc_id],
            },
            headers={"X-API-Key": test_key},
        )

        # Verify error response
        assert response.status_code == 404
        data = response.json()
        assert "detail" in data

    def test_query_processing_time_tracked(
        self,
        test_client,
        test_api_key,
        mock_milvus_search,
        mock_azure_openai,
        mock_embeddings,
    ):
        """Test processing time is tracked and included in response."""
        test_key, _ = test_api_key

        # Send query
        response = test_client.post(
            "/query",
            json={"query": "test query"},
            headers={"X-API-Key": test_key},
        )

        # Verify processing time
        assert response.status_code == 200
        data = response.json()
        assert "processing_time_ms" in data
        assert isinstance(data["processing_time_ms"], int)
        assert data["processing_time_ms"] > 0
        assert data["processing_time_ms"] < 10000  # Should be < 10 seconds

    def test_query_confidence_score_included(
        self,
        test_client,
        test_api_key,
        mock_milvus_search,
        mock_azure_openai,
        mock_embeddings,
    ):
        """Test confidence score is calculated and included."""
        test_key, _ = test_api_key

        # Send query
        response = test_client.post(
            "/query",
            json={"query": "test query"},
            headers={"X-API-Key": test_key},
        )

        # Verify confidence score
        assert response.status_code == 200
        data = response.json()
        assert "confidence" in data
        assert isinstance(data["confidence"], float)
        assert 0.0 <= data["confidence"] <= 1.0

    def test_query_no_results_returns_low_confidence(
        self,
        test_client,
        test_api_key,
        mock_embeddings,
    ):
        """Test query with no results returns low confidence answer."""
        test_key, _ = test_api_key

        with patch("src.retrieval.semantic_retriever.connections"):
            with patch("src.retrieval.semantic_retriever.Collection") as mock_coll:
                # Mock empty search results
                mock_collection = MagicMock()
                mock_collection.load = MagicMock()
                mock_collection.search.return_value = [[]]
                mock_coll.return_value = mock_collection

                # Send query
                response = test_client.post(
                    "/query",
                    json={"query": "nonexistent query"},
                    headers={"X-API-Key": test_key},
                )

                # Verify low confidence response
                assert response.status_code == 200
                data = response.json()
                assert "couldn't find any information" in data["answer"].lower()
                assert data["confidence"] < 0.2

    def test_query_retrieval_error_returns_500(
        self,
        test_client,
        test_api_key,
        mock_embeddings,
    ):
        """Test retrieval error returns 500 internal error."""
        test_key, _ = test_api_key

        with patch("src.retrieval.semantic_retriever.connections") as mock_conn:
            # Mock connection error
            mock_conn.connect.side_effect = Exception("Milvus connection failed")

            # Send query
            response = test_client.post(
                "/query",
                json={"query": "test query"},
                headers={"X-API-Key": test_key},
            )

            # Verify error response
            assert response.status_code == 500
            data = response.json()
            assert "detail" in data

    def test_query_generation_error_returns_500(
        self,
        test_client,
        test_api_key,
        mock_milvus_search,
        mock_embeddings,
    ):
        """Test answer generation error returns 500 internal error."""
        test_key, _ = test_api_key

        with patch("src.retrieval.answer_generator.AsyncAzureOpenAI") as mock_class:
            mock_client = MagicMock()
            mock_client.chat.completions.create = AsyncMock(
                side_effect=Exception("OpenAI API error")
            )
            mock_class.return_value = mock_client

            # Send query
            response = test_client.post(
                "/query",
                json={"query": "test query"},
                headers={"X-API-Key": test_key},
            )

            # Verify error response
            assert response.status_code == 500
            data = response.json()
            assert "detail" in data

    def test_query_strips_whitespace(
        self,
        test_client,
        test_api_key,
        mock_milvus_search,
        mock_azure_openai,
        mock_embeddings,
    ):
        """Test query whitespace is stripped."""
        test_key, _ = test_api_key

        # Send query with leading/trailing whitespace
        response = test_client.post(
            "/query",
            json={"query": "  test query  "},
            headers={"X-API-Key": test_key},
        )

        # Verify whitespace was stripped
        assert response.status_code == 200
        data = response.json()
        assert data["query"] == "test query"

    def test_query_multiple_document_ids(
        self,
        test_client,
        test_api_key,
        test_document,
        mock_milvus_search,
        mock_azure_openai,
        mock_embeddings,
    ):
        """Test query with multiple document_ids."""
        test_key, _ = test_api_key

        # Create second test document
        doc_id_2 = uuid4()
        with get_sync_session() as session:
            document_2 = Document(
                id=doc_id_2,
                filename="test_document_2.xlsx",
                status="completed",
                file_size=2048,
                mime_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            )
            session.add(document_2)
            session.commit()

        try:
            # Send query with multiple document IDs
            response = test_client.post(
                "/query",
                json={
                    "query": "test query",
                    "document_ids": [str(test_document), str(doc_id_2)],
                },
                headers={"X-API-Key": test_key},
            )

            # Verify response
            assert response.status_code == 200
            data = response.json()
            assert "answer" in data

        finally:
            # Cleanup second document
            with get_sync_session() as session:
                if doc_obj := session.execute(
                    select(Document).where(Document.id == doc_id_2)
                ).scalar_one_or_none():
                    session.delete(doc_obj)
                    session.commit()
