#!/usr/bin/env python3
# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
AI Model Selection Test - Prints prioritized fallback model lists for all interface calls

Tests all main interface methods in interfaceAiObjects.py and shows which models
are selected for each type of AI operation (text generation, image analysis, 
image generation, web research, etc.).
"""

import asyncio
import os
import sys
import base64

# Add the gateway to path (go up 2 levels from tests/functional/)
_gateway_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
if _gateway_path not in sys.path:
    sys.path.insert(0, _gateway_path)

from modules.services import getInterface as getServices
from modules.datamodels.datamodelAi import (
    AiCallOptions,
    AiCallRequest,
    AiModelCall,
    OperationTypeEnum,
    PriorityEnum,
    ProcessingModeEnum,
)
from modules.datamodels.datamodelUam import User
from modules.aicore.aicoreModelRegistry import modelRegistry
from modules.aicore.aicoreModelSelector import modelSelector


class ModelSelectionTester:
    def __init__(self) -> None:
        testUser = User(
            id="test_user_models",
            username="test_models",
            email="test@example.com",
            fullName="Test Models",
            language="en",
            mandateId="test_mandate",
        )
        self.services = getServices(testUser, None)

    async def initialize(self) -> None:
        from modules.services.serviceAi.mainServiceAi import AiService
        from modules.interfaces.interfaceAiObjects import AiObjects

        self.services.ai = await AiService.create(self.services)
        self.aiObjects = await AiObjects.create()

    async def _printFallbackListWithContext(self, title: str, prompt: str, context: str, options: AiCallOptions) -> None:
        print(f"\n{'='*80}")
        print(f"{title}")
        print(f"{'='*80}")
        print(
            f"Operation={options.operationType.name}, Priority={options.priority.name}, ProcessingMode={options.processingMode.name}"
        )

        # Show context and prompt sizes
        promptSize = len(prompt.encode("utf-8"))
        contextSize = len(context.encode("utf-8"))
        totalSize = promptSize + contextSize
        print(f"Prompt size: {promptSize} bytes, Context size: {contextSize} bytes, Total: {totalSize} bytes")

        availableModels = modelRegistry.getAvailableModels()
        failoverModelList = modelSelector.getFailoverModelList(
            prompt=prompt,
            context=context,
            options=options,
            availableModels=availableModels,
        )

        if not failoverModelList:
            print("No suitable models found (capability filter returned empty list).")
            return

        print("Prioritized fallback model sequence (name | quality | speed | $/1k in | ctx | score):")
        for idx, m in enumerate(failoverModelList, 1):
            costIn = getattr(m, "costPer1kTokensInput", 0.0)
            # Calculate detailed score breakdown
            promptSize = len(prompt.encode("utf-8"))
            contextSize = len(context.encode("utf-8"))
            totalSize = promptSize + contextSize
            
            # Get detailed scoring
            sizeRating = modelSelector._getSizeRating(m, totalSize)
            processingModeRating = modelSelector._getProcessingModeRating(m.processingMode, options.processingMode)
            priorityRating = modelSelector._getPriorityRating(m, options.priority)
            totalScore = sizeRating + processingModeRating + priorityRating
            
            print(
                f" {idx:>2}. {m.displayName} | Q={getattr(m, 'qualityRating', 0)} | S={getattr(m, 'speedRating', 0)} | ${costIn:.4f} | ctx={getattr(m, 'contextLength', 0)} | score={totalScore:.3f}"
            )
            print(f"      Size: {sizeRating:.3f}, ProcessingMode: {processingModeRating:.3f}, Priority: {priorityRating:.3f}")

    async def _printFallbackList(self, title: str, prompt: str, options: AiCallOptions) -> None:
        print(f"\n{'='*80}")
        print(f"{title}")
        print(f"{'='*80}")
        print(
            f"Operation={options.operationType.name}, Priority={options.priority.name}, ProcessingMode={options.processingMode.name}"
        )

        # Show context and prompt sizes
        context = ""  # Currently using empty context
        promptSize = len(prompt.encode("utf-8"))
        contextSize = len(context.encode("utf-8"))
        totalSize = promptSize + contextSize
        print(f"Prompt size: {promptSize} bytes, Context size: {contextSize} bytes, Total: {totalSize} bytes")

        availableModels = modelRegistry.getAvailableModels()
        failoverModelList = modelSelector.getFailoverModelList(
            prompt=prompt,
            context=context,
            options=options,
            availableModels=availableModels,
        )

        if not failoverModelList:
            print("No suitable models found (capability filter returned empty list).")
            return

        print("Prioritized fallback model sequence (name | quality | speed | $/1k in | ctx | score):")
        for idx, m in enumerate(failoverModelList, 1):
            costIn = getattr(m, "costPer1kTokensInput", 0.0)
            # Calculate detailed score breakdown
            promptSize = len(prompt.encode("utf-8"))
            contextSize = len(context.encode("utf-8"))
            totalSize = promptSize + contextSize
            
            # Get detailed scoring
            sizeRating = modelSelector._getSizeRating(m, totalSize)
            processingModeRating = modelSelector._getProcessingModeRating(m.processingMode, options.processingMode)
            priorityRating = modelSelector._getPriorityRating(m, options.priority)
            totalScore = sizeRating + processingModeRating + priorityRating
            
            print(
                f" {idx:>2}. {m.displayName} | Q={getattr(m, 'qualityRating', 0)} | S={getattr(m, 'speedRating', 0)} | ${costIn:.4f} | ctx={getattr(m, 'contextLength', 0)} | score={totalScore:.3f}"
            )
            print(f"      Size: {sizeRating:.3f}, ProcessingMode: {processingModeRating:.3f}, Priority: {priorityRating:.3f}")

    async def run(self) -> None:
        """Test model selection for all interface methods."""
        print("=" * 100)
        print("AI INTERFACE MODEL SELECTION TEST")
        print("=" * 100)
        print("Testing model selection for all interface methods in interfaceAiObjects.py")
        print("=" * 100)

        # Test 1: Text Generation (call method)
        await self._testTextGeneration()

        # Test 2: Image Analysis (callImage method)
        await self._testImageAnalysis()

        # Test 3: Image Generation (generateImage method)
        await self._testImageGeneration()

        # Test 4: Web Search (searchWebsites method)
        await self._testWebSearch()

        # Test 5: Web Crawling (crawlWebsites method)
        await self._testWebCrawling()

        # Test 6: Web Research (webQuery method)
        await self._testWebResearch()

        # Test 7: Content Analysis with Chunking
        await self._testContentAnalysis()

        # Test 8: Website Selection
        await self._testWebsiteSelection()

        # Test 9: Actual Interface Calls
        await self._testActualInterfaceCalls()

        # Show model registry summary
        await self._showModelSummary()

        print("\n" + "=" * 100)
        print("ALL INTERFACE TESTS COMPLETED")
        print("=" * 100)

    async def _testTextGeneration(self) -> None:
        """Test model selection for text generation calls."""
        print(f"\n{'='*80}")
        print("1. TEXT GENERATION (call method)")
        print(f"{'='*80}")

        # Test different text generation scenarios
        scenarios = [
            ("Text Analysis", "Write a summary about artificial intelligence trends.", OperationTypeEnum.DATA_ANALYSE),
            ("Planning Task", "Create a project plan for software development.", OperationTypeEnum.PLAN),
            ("Analysis Task", "Analyze the pros and cons of cloud computing.", OperationTypeEnum.DATA_ANALYSE),
        ]

        for title, prompt, operation_type in scenarios:
            options = AiCallOptions(
                operationType=operation_type,
                priority=PriorityEnum.BALANCED,
                processingMode=ProcessingModeEnum.ADVANCED,
                maxCost=0.05,
                maxProcessingTime=30,
            )
            await self._printFallbackList(f"  {title}", prompt, options)

    async def _testImageAnalysis(self) -> None:
        """Test model selection for image analysis calls."""
        print(f"\n{'='*80}")
        print("2. IMAGE ANALYSIS (callImage method)")
        print(f"{'='*80}")

        # Create a small test image (1x1 pixel PNG)
        test_image_data = base64.b64encode(b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\tpHYs\x00\x00\x0b\x13\x00\x00\x0b\x13\x01\x00\x9a\x9c\x18\x00\x00\x00\nIDATx\x9cc```\x00\x00\x00\x04\x00\x01\xdd\x8d\xb4\x1c\x00\x00\x00\x00IEND\xaeB`\x82').decode('utf-8')

        options = AiCallOptions(
            operationType=OperationTypeEnum.IMAGE_ANALYSE,
            priority=PriorityEnum.BALANCED,
            processingMode=ProcessingModeEnum.ADVANCED,
            maxCost=0.02,
            maxProcessingTime=20,
        )

        prompt = "Describe what you see in this image."
        await self._printFallbackList("  Image Analysis", prompt, options)

    async def _testImageGeneration(self) -> None:
        """Test model selection for image generation calls."""
        print(f"\n{'='*80}")
        print("3. IMAGE GENERATION (generateImage method)")
        print(f"{'='*80}")

        options = AiCallOptions(
            operationType=OperationTypeEnum.IMAGE_GENERATE,
            priority=PriorityEnum.QUALITY,
            processingMode=ProcessingModeEnum.DETAILED,
            maxCost=0.10,
            maxProcessingTime=60,
        )

        prompt = "A futuristic cityscape with flying cars and neon lights."
        await self._printFallbackList("  Image Generation", prompt, options)

    async def _testWebResearch(self) -> None:
        """Test model selection for web research calls."""
        print(f"\n{'='*80}")
        print("6. WEB RESEARCH (webQuery method)")
        print(f"{'='*80}")

        options = AiCallOptions(
            operationType=OperationTypeEnum.WEB_SEARCH,
            priority=PriorityEnum.BALANCED,
            processingMode=ProcessingModeEnum.ADVANCED,
            maxCost=0.05,
            maxProcessingTime=30,
        )

        prompt = "What are the latest trends in artificial intelligence?"
        await self._printFallbackList("  Web Research", prompt, options)

    async def _testWebSearch(self) -> None:
        """Test model selection for web search calls."""
        print(f"\n{'='*80}")
        print("4. WEB SEARCH (searchWebsites method)")
        print(f"{'='*80}")

        options = AiCallOptions(
            operationType=OperationTypeEnum.WEB_SEARCH,
            priority=PriorityEnum.BALANCED,
            processingMode=ProcessingModeEnum.BASIC,
            maxCost=0.01,
            maxProcessingTime=30,
        )

        prompt = "Search for artificial intelligence companies"
        await self._printFallbackList("  Web Search", prompt, options)

    async def _testWebCrawling(self) -> None:
        """Test model selection for web crawling calls."""
        print(f"\n{'='*80}")
        print("5. WEB CRAWLING (crawlWebsites method)")
        print(f"{'='*80}")

        options = AiCallOptions(
            operationType=OperationTypeEnum.WEB_CRAWL,
            priority=PriorityEnum.BALANCED,
            processingMode=ProcessingModeEnum.BASIC,
            maxCost=0.02,
            maxProcessingTime=60,
        )

        prompt = "Crawl content from these URLs"
        await self._printFallbackList("  Web Crawling", prompt, options)

    async def _testContentAnalysis(self) -> None:
        """Test model selection for content analysis with chunking."""
        print(f"\n{'='*80}")
        print("7. CONTENT ANALYSIS WITH CHUNKING")
        print(f"{'='*80}")

        # Test with large content to trigger chunking
        large_content = {
            "https://example.com/page1": "This is a large document about artificial intelligence. " * 1000,
            "https://example.com/page2": "This is another large document about machine learning. " * 1000,
        }

        options = AiCallOptions(
            operationType=OperationTypeEnum.DATA_ANALYSE,
            priority=PriorityEnum.BALANCED,
            processingMode=ProcessingModeEnum.ADVANCED,
            maxCost=0.10,
            maxProcessingTime=60,
        )

        prompt = "Analyze this content and provide key insights."
        await self._printFallbackList("  Content Analysis", prompt, options)

    async def _testWebsiteSelection(self) -> None:
        """Test model selection for website selection."""
        print(f"\n{'='*80}")
        print("8. WEBSITE SELECTION (selectRelevantWebsites method)")
        print(f"{'='*80}")

        # This method uses webQuery internally, so it uses the same model selection as web research
        options = AiCallOptions(
            operationType=OperationTypeEnum.WEB_SEARCH,
            priority=PriorityEnum.BALANCED,
            processingMode=ProcessingModeEnum.ADVANCED,
            maxCost=0.03,
            maxProcessingTime=20,
        )

        prompt = "Select the most relevant websites from this list for AI research."
        await self._printFallbackList("  Website Selection", prompt, options)

    async def _testActualInterfaceCalls(self) -> None:
        """Test actual interface calls to show real model selection."""
        print(f"\n{'='*80}")
        print("9. ACTUAL INTERFACE CALLS (Real Model Selection)")
        print(f"{'='*80}")

        # Test 1: Text generation call
        print("\n  Testing: aiObjects.call() - Text Generation")
        try:
            request = AiCallRequest(
                prompt="Write a short summary about machine learning.",
                context="",
                options=AiCallOptions(
                    operationType=OperationTypeEnum.DATA_ANALYSE,
                    priority=PriorityEnum.BALANCED,
                    processingMode=ProcessingModeEnum.ADVANCED,
                    maxCost=0.05,
                    maxProcessingTime=30,
                )
            )
            
            # Get the model selection that would be used
            availableModels = modelRegistry.getAvailableModels()
            failoverModelList = modelSelector.getFailoverModelList(
                prompt=request.prompt,
                context=request.context,
                options=request.options,
                availableModels=availableModels,
            )
            
            if failoverModelList:
                print(f"    Selected model: {failoverModelList[0].displayName}")
                print(f"    Fallback models: {[m.displayName for m in failoverModelList[1:3]]}")
            else:
                print("    No suitable models found")
                
        except Exception as e:
            print(f"    Error: {e}")

        # Test 2: Image analysis call
        print("\n  Testing: aiObjects.callImage() - Image Analysis")
        try:
            options = AiCallOptions(
                operationType=OperationTypeEnum.IMAGE_ANALYSE,
                priority=PriorityEnum.BALANCED,
                processingMode=ProcessingModeEnum.ADVANCED,
                maxCost=0.02,
                maxProcessingTime=20,
            )
            
            availableModels = modelRegistry.getAvailableModels()
            failoverModelList = modelSelector.getFailoverModelList(
                prompt="Describe this image",
                context="",
                options=options,
                availableModels=availableModels,
            )
            
            if failoverModelList:
                print(f"    Selected model: {failoverModelList[0].displayName}")
                print(f"    Fallback models: {[m.displayName for m in failoverModelList[1:3]]}")
            else:
                print("    No suitable models found")
                
        except Exception as e:
            print(f"    Error: {e}")

        # Test 3: Image generation call
        print("\n  Testing: aiObjects.generateImage() - Image Generation")
        try:
            options = AiCallOptions(
                operationType=OperationTypeEnum.IMAGE_GENERATE,
                priority=PriorityEnum.QUALITY,
                processingMode=ProcessingModeEnum.DETAILED,
                maxCost=0.10,
                maxProcessingTime=60,
            )
            
            availableModels = modelRegistry.getAvailableModels()
            failoverModelList = modelSelector.getFailoverModelList(
                prompt="A futuristic cityscape",
                context="",
                options=options,
                availableModels=availableModels,
            )
            
            if failoverModelList:
                print(f"    Selected model: {failoverModelList[0].displayName}")
                print(f"    Fallback models: {[m.displayName for m in failoverModelList[1:3]]}")
            else:
                print("    No suitable models found")
                
        except Exception as e:
            print(f"    Error: {e}")

        # Test 4: Web research call
        print("\n  Testing: aiObjects.webQuery() - Web Research")
        try:
            options = AiCallOptions(
                operationType=OperationTypeEnum.WEB_SEARCH,
                priority=PriorityEnum.BALANCED,
                processingMode=ProcessingModeEnum.ADVANCED,
                maxCost=0.05,
                maxProcessingTime=30,
            )
            
            availableModels = modelRegistry.getAvailableModels()
            failoverModelList = modelSelector.getFailoverModelList(
                prompt="What are AI trends?",
                context="",
                options=options,
                availableModels=availableModels,
            )
            
            if failoverModelList:
                print(f"    Selected model: {failoverModelList[0].displayName}")
                print(f"    Fallback models: {[m.displayName for m in failoverModelList[1:3]]}")
            else:
                print("    No suitable models found")
                
        except Exception as e:
            print(f"    Error: {e}")

    async def _showModelSummary(self) -> None:
        """Show summary of all available models and their capabilities."""
        print(f"\n{'='*80}")
        print("MODEL REGISTRY SUMMARY")
        print(f"{'='*80}")

        availableModels = modelRegistry.getAvailableModels()
        print(f"Total models available: {len(availableModels)}")
        
        # Group by connector type
        by_connector = {}
        for model in availableModels:
            connector_type = getattr(model, 'connectorType', 'unknown')
            if connector_type not in by_connector:
                by_connector[connector_type] = []
            by_connector[connector_type].append(model)
        
        print(f"\nModels by connector type:")
        for connector_type, models in by_connector.items():
            print(f"  {connector_type}: {len(models)} models")
            for model in models:
                capabilities = getattr(model, 'capabilities', [])
                print(f"    - {model.displayName}: {capabilities}")
        
        # Show operation type support
        print(f"\nOperation type support:")
        for op_type in OperationTypeEnum:
            supported_models = [m for m in availableModels if hasattr(m, 'operationTypes') and any(ot.operationType == op_type for ot in m.operationTypes)]
            print(f"  {op_type.name}: {len(supported_models)} models")
            if supported_models:
                model_names = [m.displayName for m in supported_models[:3]]  # Show first 3 models
                print(f"    Models: {', '.join(model_names)}")


async def main() -> None:
    tester = ModelSelectionTester()
    await tester.initialize()
    await tester.run()


if __name__ == "__main__":
    asyncio.run(main())