import pytest
from unittest.mock import AsyncMock, patch, MagicMock
from uuid import uuid4
import aiohttp
from tg_bot.application.services.rag_service import RAGService
from tests.metrics.test_hit_at_5 import calculate_hit_at_5, calculate_average_hit_at_5


class TestRAGIntegration:

    @pytest.fixture
    def rag_service(self):
        service = RAGService()
        from tg_bot.infrastructure.external.deepseek_client import DeepSeekClient
        service.deepseek_client = DeepSeekClient()
        return service

    @pytest.mark.asyncio
    async def test_rag_service_search_documents_real_implementation(self, rag_service):
        user_telegram_id = "123456789"
        query = "трудовой договор"
        user_uuid = str(uuid4())
        collection_id = str(uuid4())
        
        mock_user_response = AsyncMock()
        mock_user_response.status = 200
        mock_user_response.json = AsyncMock(return_value={"user_id": user_uuid})
        
        mock_collections_response = AsyncMock()
        mock_collections_response.status = 200
        mock_collections_response.json = AsyncMock(return_value=[
            {"collection_id": collection_id, "name": "Законы"}
        ])
        
        mock_documents_response = AsyncMock()
        mock_documents_response.status = 200
        mock_documents_response.json = AsyncMock(return_value=[
            {
                "document_id": str(uuid4()),
                "title": "Трудовой кодекс",
                "content": "Содержание о трудовых договорах"
            }
        ])
        
        with patch('aiohttp.ClientSession') as mock_session_class:
            mock_session = AsyncMock()
            mock_session.__aenter__ = AsyncMock(return_value=mock_session)
            mock_session.__aexit__ = AsyncMock(return_value=None)
            mock_session.get = AsyncMock(side_effect=[
                mock_user_response,
                mock_collections_response,
                mock_documents_response
            ])
            mock_session_class.return_value = mock_session
            
            result = await rag_service.search_documents_in_collections(user_telegram_id, query)
            
            assert isinstance(result, list)
            if result:
                assert "document_id" in result[0]
                assert "title" in result[0]

    @pytest.mark.asyncio
    async def test_rag_service_generate_answer_real_implementation(self, rag_service):
        question = "Какие права имеет работник?"
        user_telegram_id = "123456789"
        
        mock_documents = [
            {
                "document_id": str(uuid4()),
                "title": "Трудовой кодекс",
                "content": "Работник имеет право на...",
                "collection_name": "Законы"
            }
            for _ in range(3)
        ]
        
        with patch.object(rag_service, 'search_documents_in_collections') as mock_search, \
             patch.object(rag_service.deepseek_client, 'chat_completion') as mock_deepseek:
            
            mock_search.return_value = mock_documents
            mock_deepseek.return_value = {
                "content": "Работник имеет следующие права...",
                "usage": {"prompt_tokens": 100, "completion_tokens": 200, "total_tokens": 300}
            }
            
            result = await rag_service.generate_answer_with_rag(question, user_telegram_id)
            
            assert "answer" in result
            assert "sources" in result
            assert "usage" in result
            assert len(result["sources"]) <= 5
            assert result["answer"] != ""
            
            for source in result["sources"]:
                assert "title" in source
                assert "collection" in source
                assert "document_id" in source

    @pytest.mark.asyncio
    async def test_rag_service_limits_to_top5_documents(self, rag_service):
        question = "Тестовый вопрос"
        user_telegram_id = "123456789"
        
        many_documents = [
            {
                "document_id": str(uuid4()),
                "title": f"Документ {i}",
                "content": f"Содержание {i}",
                "collection_name": "Коллекция"
            }
            for i in range(10)
        ]
        
        with patch.object(rag_service, 'search_documents_in_collections') as mock_search, \
             patch.object(rag_service.deepseek_client, 'chat_completion') as mock_deepseek:
            
            mock_search.return_value = many_documents
            mock_deepseek.return_value = {
                "content": "Ответ",
                "usage": {}
            }
            
            result = await rag_service.generate_answer_with_rag(question, user_telegram_id)
            
            assert len(result["sources"]) == 5

    @pytest.mark.asyncio
    async def test_rag_service_handles_empty_search_results(self, rag_service):
        question = "Вопрос без документов"
        user_telegram_id = "123456789"
        
        with patch.object(rag_service, 'search_documents_in_collections') as mock_search, \
             patch.object(rag_service.deepseek_client, 'chat_completion') as mock_deepseek:
            
            mock_search.return_value = []
            mock_deepseek.return_value = {
                "content": "Ответ",
                "usage": {}
            }
            
            result = await rag_service.generate_answer_with_rag(question, user_telegram_id)
            
            assert result["sources"] == []
            assert "Релевантные документы не найдены" in result.get("answer", "") or \
                   result["answer"] == "No relevant documents found"