#!/usr/bin/env python3
"""
Unit tests for AI service (mainServiceAi.py)
Tests callAiContent, callAiPlanning, and related functionality.
"""

import pytest
from unittest.mock import Mock, AsyncMock, patch

from modules.datamodels.datamodelAi import AiCallOptions, OperationTypeEnum, PriorityEnum, ProcessingModeEnum
from modules.datamodels.datamodelExtraction import ContentPart
from modules.datamodels.datamodelWorkflow import AiResponse


class TestAiServiceCallAiContent:
    """Test callAiContent method (mocked)"""
    
    @pytest.mark.asyncio
    async def test_callAiContent_requires_operationType(self):
        """Test that callAiContent requires operationType to be set"""
        from modules.services.serviceAi.mainServiceAi import AiService
        
        # Create mock services
        mockServices = Mock()
        mockServices.workflow = None
        mockServices.chat = Mock()
        mockServices.chat.progressLogStart = Mock()
        mockServices.chat.progressLogUpdate = Mock()
        mockServices.chat.progressLogFinish = Mock()
        mockServices.chat.storeWorkflowStat = Mock()
        
        aiService = AiService(mockServices)
        
        # Mock aiObjects initialization
        aiService.aiObjects = Mock()
        aiService._ensureAiObjectsInitialized = AsyncMock()
        
        # Test with missing operationType - should analyze prompt
        options = AiCallOptions()  # operationType not set
        options.operationType = None
        
        # Mock _analyzePromptAndCreateOptions
        analyzedOptions = AiCallOptions()
        analyzedOptions.operationType = OperationTypeEnum.DATA_ANALYSE
        aiService._analyzePromptAndCreateOptions = AsyncMock(return_value=analyzedOptions)
        
        # Mock _callAiWithLooping
        aiService._callAiWithLooping = AsyncMock(return_value="Test response")
        
        # Mock aiObjects.call
        mockResponse = Mock()
        mockResponse.content = "Test response"
        aiService.aiObjects.call = AsyncMock(return_value=mockResponse)
        
        # Call should work (will analyze prompt if operationType not set)
        result = await aiService.callAiContent(
            prompt="Test prompt",
            options=options
        )
        
        # Should have analyzed prompt and set operationType
        assert result is not None
        assert isinstance(result, AiResponse)


class TestAiServiceCallAiPlanning:
    """Test callAiPlanning method (mocked)"""
    
    @pytest.mark.asyncio
    async def test_callAiPlanning_basic(self):
        """Test basic callAiPlanning call"""
        from modules.services.serviceAi.mainServiceAi import AiService
        
        # Create mock services
        mockServices = Mock()
        mockServices.workflow = None
        mockServices.utils = Mock()
        mockServices.utils.writeDebugFile = Mock()
        
        aiService = AiService(mockServices)
        
        # Mock aiObjects
        aiService.aiObjects = Mock()
        mockResponse = Mock()
        mockResponse.content = '{"result": "plan"}'
        aiService.aiObjects.call = AsyncMock(return_value=mockResponse)
        aiService._ensureAiObjectsInitialized = AsyncMock()
        
        # Call planning
        result = await aiService.callAiPlanning(
            prompt="Test planning prompt"
        )
        
        assert result == '{"result": "plan"}'


class TestAiServiceOperationTypeHandling:
    """Test operationType handling in callAiContent"""
    
    @pytest.mark.asyncio
    async def test_callAiContent_with_outputFormat_sets_documentGenerate(self):
        """Test that outputFormat sets operationType to DOCUMENT_GENERATE"""
        from modules.services.serviceAi.mainServiceAi import AiService
        
        mockServices = Mock()
        mockServices.workflow = None
        mockServices.chat = Mock()
        mockServices.chat.progressLogStart = Mock()
        mockServices.chat.progressLogUpdate = Mock()
        mockServices.chat.progressLogFinish = Mock()
        mockServices.utils = Mock()
        mockServices.utils.jsonExtractString = Mock(return_value='{"documents": []}')
        
        aiService = AiService(mockServices)
        aiService.aiObjects = Mock()
        aiService._ensureAiObjectsInitialized = AsyncMock()
        
        # Mock _callAiWithLooping
        aiService._callAiWithLooping = AsyncMock(return_value='{"documents": []}')
        
        # Mock generation service
        with patch('modules.services.serviceGeneration.mainServiceGeneration.GenerationService') as mockGenService:
            mockGenInstance = Mock()
            mockGenInstance.renderReport = AsyncMock(return_value=(b"content", "application/pdf"))
            mockGenService.return_value = mockGenInstance
            
            options = AiCallOptions()  # operationType not set
            options.operationType = None
            
            # Should set operationType to DOCUMENT_GENERATE when outputFormat is provided
            try:
                result = await aiService.callAiContent(
                    prompt="Generate document",
                    options=options,
                    outputFormat="pdf"
                )
                # If it gets here, operationType was set correctly
                assert options.operationType == OperationTypeEnum.DOCUMENT_GENERATE
            except Exception:
                # If it fails, that's okay for unit test - we're testing the logic
                pass


if __name__ == "__main__":
    pytest.main([__file__, "-v"])