#!/usr/bin/env python3
# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Test script for methodAi operations.
Tests all OperationType's with various prompts through the workflow action interface.
"""

import asyncio
import sys
import os
from datetime import datetime
from typing import Dict, Any, List

# 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.datamodels.datamodelAi import OperationTypeEnum
from modules.datamodels.datamodelChatbot import ChatWorkflow, ChatDocument, WorkflowModeEnum
from modules.datamodels.datamodelUam import User


class MethodAiOperationsTester:
    """Test all operation types through methodAi.process() action."""
    
    def __init__(self):
        # Use root user for testing (has full access to everything)
        from modules.interfaces.interfaceDbAppObjects import getRootInterface
        rootInterface = getRootInterface()
        self.testUser = rootInterface.currentUser
        
        self.services = None
        self.methodAi = None
        self.testResults = []
        
        # Create logs directory if it doesn't exist (go up 1 level from gateway/)
        _gateway_dir = os.path.dirname(_gateway_path)
        self.logsDir = os.path.join(_gateway_dir, "local", "logs")
        os.makedirs(self.logsDir, exist_ok=True)
        
        # Create modeltest subdirectory
        self.modelTestDir = os.path.join(self.logsDir, "modeltest")
        os.makedirs(self.modelTestDir, exist_ok=True)
        
        # Test prompts for each operation type
        self.testPrompts = {
            OperationTypeEnum.PLAN: {
                "aiPrompt": "Create a 5-step plan to organize a project meeting and include the manual for the project management office.",
                "resultType": "json"
            },
            OperationTypeEnum.DATA_ANALYSE: {
                "aiPrompt": "Analyze the following text and extract the main topics and key points: 'Machine learning is transforming healthcare by enabling early disease detection through pattern recognition in medical images.'",
                "resultType": "json"
            },
            OperationTypeEnum.DATA_GENERATE: {
                # "aiPrompt": "Generate the first 1000 prime numbers.",
                "aiPrompt": "Schreibe einen Bericht über die Verhaltensweisen in Finnalnd's Saunas in 50 Kapiteln",
                "resultType": "docx"
            },
            OperationTypeEnum.DATA_EXTRACT: {
                "aiPrompt": "Extract all email addresses and phone numbers from the following text: 'Contact us at support@example.com or call 123-456-7890. For sales, email sales@example.com or call 987-654-3210.'",
                "resultType": "json"
            },
            OperationTypeEnum.IMAGE_ANALYSE: {
                "aiPrompt": "Analyze this image and describe what you see, including any text or numbers visible.",
                "resultType": "json",
                # documentList should contain document references resolvable by workflow service
                # The test image will be uploaded and referenced during initialization
                "documentList": []  # Will be populated in initialize() if test image is available
            },
            OperationTypeEnum.IMAGE_GENERATE: {
                "aiPrompt": "A beautiful sunset over the ocean with purple and orange hues",
                "resultType": "png"
            },
            OperationTypeEnum.WEB_SEARCH_DATA: {
                "aiPrompt": "Who works in valueon ag in switzerland?",
                "resultType": "json"
            },
            OperationTypeEnum.WEB_CRAWL: {
                "aiPrompt": "Extract who works in this company",
                "resultType": "json",
                "documentList": ["https://www.valueon.ch"]
            }
        }
    
    async def initialize(self):
        """Initialize services and methodAi."""
        print("🔧 Initializing services...")
        
        # Set logging level to DEBUG to see debug messages
        import logging
        logging.getLogger().setLevel(logging.DEBUG)
        
        # Import and initialize services - use the same approach as routeChatPlayground
        import modules.interfaces.interfaceDbChatbot as interfaceDbChatbot
        interfaceDbChat = interfaceDbChatbot.getInterface(self.testUser)
        
        # Import and initialize services
        from modules.services import getInterface as getServices
        
        # Get services first
        self.services = getServices(self.testUser, None)
        
        # Now create AND SAVE workflow in database using the interface
        import uuid
        import time
        currentTimestamp = time.time()
        
        testWorkflow = ChatWorkflow(
            id=str(uuid.uuid4()),
            name="Test Workflow",
            status="running",
            startedAt=currentTimestamp,
            lastActivity=currentTimestamp,
            currentRound=1,
            currentTask=0,
            currentAction=0,
            totalTasks=0,
            totalActions=0,
            mandateId=self.testUser.mandateId,
            messageIds=[],
            workflowMode=WorkflowModeEnum.WORKFLOW_DYNAMIC,
            maxSteps=5
        )
        
        # SAVE workflow to database so it exists for access control
        # Convert ChatWorkflow to dict for createWorkflow
        workflowDict = testWorkflow.model_dump()
        interfaceDbChat.createWorkflow(workflowDict)
        
        # Set the workflow in services (Services class uses .workflow, not .currentWorkflow)
        self.services.workflow = testWorkflow
        
        # Debug: Print workflow status
        print(f"Debug: services.workflow is set: {hasattr(self.services, 'workflow') and self.services.workflow is not None}")
        if self.services.workflow:
            print(f"Debug: Workflow ID: {self.services.workflow.id}")
        
        # Import and initialize methodAi AFTER setting workflow
        from modules.workflows.methods.methodAi import MethodAi
        self.methodAi = MethodAi(self.services)
        
        # Verify methodAi has access to the workflow
        if hasattr(self.methodAi, 'services'):
            print(f"Debug: methodAi.services.workflow is set: {hasattr(self.methodAi.services, 'workflow') and self.methodAi.services.workflow is not None}")
        
        # Prepare test image document for IMAGE_ANALYSE if available
        await self._prepareTestImageDocument()
        
        print("✅ Services initialized")
        print(f"📁 Results will be saved to: {self.modelTestDir}")
    
    async def _prepareTestImageDocument(self):
        """Upload test image as a document for IMAGE_ANALYSE testing."""
        try:
            # Path to test image (relative to gateway directory)
            testImagePath = os.path.join(
                os.path.dirname(__file__),  # tests/functional/
                "..",  # tests/
                "testdata",  # tests/testdata/
                "Foto20250906_125903.jpg"
            )
            testImagePath = os.path.abspath(testImagePath)
            
            if not os.path.exists(testImagePath):
                print(f"⚠️  Test image not found at: {testImagePath}")
                print("   IMAGE_ANALYSE tests will be skipped or will fail")
                return
            
            # Read image file
            with open(testImagePath, 'rb') as f:
                imageData = f.read()
            
            # Create a ChatDocument
            from modules.datamodels.datamodelChatbot import ChatDocument
            import uuid
            
            testImageDoc = ChatDocument(
                id=str(uuid.uuid4()),
                documentName="Foto20250906_125903.jpg",
                mimeType="image/jpeg",
                documentData=imageData,
                workflowId=self.services.workflow.id if self.services.workflow else None
            )
            
            # Create a message with this document
            from modules.datamodels.datamodelChatbot import ChatMessage
            import time
            
            testMessage = ChatMessage(
                id=str(uuid.uuid4()),
                workflowId=self.services.workflow.id if self.services.workflow else None,
                role="user",
                content="Test image for IMAGE_ANALYSE",
                language="en",
                timestamp=time.time(),
                documents=[testImageDoc]
            )
            
            # Save message to database
            if self.services.workflow:
                import modules.interfaces.interfaceDbChatbot as interfaceDbChatbot
                interfaceDbChat = interfaceDbChatbot.getInterface(self.testUser)
                messageDict = testMessage.model_dump()
                interfaceDbChat.createMessage(messageDict)
                
                # Update workflow messageIds
                if self.services.workflow.messageIds is None:
                    self.services.workflow.messageIds = []
                self.services.workflow.messageIds.append(testMessage.id)
                
                # Update documentList for IMAGE_ANALYSE test
                # Format: messageId:label (using documentName as label)
                docRef = f"{testMessage.id}:{testImageDoc.documentName}"
                self.testPrompts[OperationTypeEnum.IMAGE_ANALYSE]["documentList"] = [docRef]
                
                print(f"✅ Test image uploaded: {testImageDoc.documentName}")
                print(f"   Document reference: {docRef}")
            else:
                print("⚠️  No workflow available, cannot upload test image")
                
        except Exception as e:
            print(f"⚠️  Failed to prepare test image document: {str(e)}")
            print("   IMAGE_ANALYSE tests may fail")
    
    async def testOperation(self, operationType: OperationTypeEnum) -> Dict[str, Any]:
        """Test a specific operation type."""
        print(f"\n{'='*80}")
        print(f"TESTING OPERATION: {operationType.value}")
        print(f"{'='*80}")
        
        startTime = asyncio.get_event_loop().time()
        
        # Get test prompt for this operation
        testConfig = self.testPrompts.get(operationType, {})
        
        if not testConfig:
            result = {
                "operationType": operationType.value,
                "status": "ERROR",
                "error": "No test configuration found for this operation type",
                "processingTime": 0.0
            }
            self.testResults.append(result)
            return result
        
        print(f"Prompt: {testConfig.get('aiPrompt', 'N/A')}")
        print(f"Result Type: {testConfig.get('resultType', 'txt')}")
        
        try:
            # Prepare parameters
            parameters = {
                "aiPrompt": testConfig.get("aiPrompt"),
                "resultType": testConfig.get("resultType", "txt")
            }
            
            # Add document list if provided
            if "documentList" in testConfig and testConfig["documentList"]:
                parameters["documentList"] = testConfig["documentList"]
            
            # Ensure workflow is still set in both self.services AND methodAi.services
            if not self.services.workflow or (hasattr(self, 'methodAi') and hasattr(self.methodAi, 'services') and not self.methodAi.services.workflow):
                print(f"⚠️  Warning: Workflow is None, trying to re-set it...")
                import time
                import uuid
                currentTimestamp = time.time()
                testWorkflow = ChatWorkflow(
                    id=str(uuid.uuid4()),
                    name="Test Workflow",
                    status="running",
                    startedAt=currentTimestamp,
                    lastActivity=currentTimestamp,
                    currentRound=1,
                    currentTask=0,
                    currentAction=0,
                    totalTasks=0,
                    totalActions=0,
                    mandateId=self.testUser.mandateId,
                    messageIds=[],
                    workflowMode=WorkflowModeEnum.WORKFLOW_DYNAMIC,
                    maxSteps=5
                )
                # Save workflow to database
                import modules.interfaces.interfaceDbChatbot as interfaceDbChatbot
                interfaceDbChat = interfaceDbChatbot.getInterface(self.testUser)
                workflowDict = testWorkflow.model_dump()
                interfaceDbChat.createWorkflow(workflowDict)
                
                self.services.workflow = testWorkflow
                # Also set in methodAi.services if it exists
                if hasattr(self, 'methodAi') and hasattr(self.methodAi, 'services'):
                    self.methodAi.services.workflow = testWorkflow
            
            # Call methodAi.process()
            print(f"Calling methodAi.process()...")
            print(f"Debug: Current workflow ID before call: {self.services.workflow.id if self.services.workflow else 'None'}")
            print(f"Debug: methodAi.services.workflow: {self.methodAi.services.workflow.id if hasattr(self.methodAi, 'services') and self.methodAi.services.workflow else 'None/NotSet'}")
            print(f"Debug: Is same services object? {self.services is self.methodAi.services}")
            print(f"Debug: services id: {id(self.services)}")
            print(f"Debug: methodAi.services id: {id(self.methodAi.services)}")
            
            actionResult = await self.methodAi.process(parameters)
            
            endTime = asyncio.get_event_loop().time()
            processingTime = endTime - startTime
            
            # Analyze result
            result = {
                "operationType": operationType.value,
                "status": "SUCCESS" if actionResult.success else "ERROR",
                "processingTime": round(processingTime, 2),
                "hasDocuments": len(actionResult.documents) > 0 if actionResult.documents else False,
                "documentCount": len(actionResult.documents) if actionResult.documents else 0,
                "error": actionResult.error if not actionResult.success else None
            }
            
            # Extract document information
            if actionResult.documents:
                doc = actionResult.documents[0]
                result["documentName"] = doc.documentName
                result["mimeType"] = doc.mimeType
                result["dataSize"] = len(doc.documentData) if doc.documentData else 0
                result["dataPreview"] = str(doc.documentData)[:200] + "..." if len(str(doc.documentData)) > 200 else str(doc.documentData)
            
            print(f"✅ Status: {result['status']}")
            print(f"⏱️  Processing time: {result['processingTime']}s")
            print(f"📄 Documents: {result.get('documentCount', 0)}")
            
            if actionResult.success:
                if result.get('documentName'):
                    print(f"📄 Saved: {result['documentName']}")
                    print(f"📄 MIME type: {result.get('mimeType')}")
                    print(f"📄 Size: {result.get('dataSize')} bytes")
                
                # Try to decode if it's JSON
                if result.get('mimeType') == 'application/json':
                    try:
                        import json
                        jsonData = json.loads(actionResult.documents[0].documentData)
                        result["isValidJson"] = True
                        result["jsonKeys"] = list(jsonData.keys()) if isinstance(jsonData, dict) else "Not a dict"
                        print(f"✅ Valid JSON with keys: {result['jsonKeys']}")
                    except:
                        result["isValidJson"] = False
                        print(f"⚠️  Not valid JSON")
            else:
                print(f"❌ Error: {result.get('error')}")
            
            self.testResults.append(result)
            return result
            
        except Exception as e:
            endTime = asyncio.get_event_loop().time()
            processingTime = endTime - startTime
            
            result = {
                "operationType": operationType.value,
                "status": "EXCEPTION",
                "processingTime": round(processingTime, 2),
                "error": str(e),
                "hasDocuments": False
            }
            
            print(f"💥 EXCEPTION: {str(e)}")
            self.testResults.append(result)
            return result
    
    async def testAllOperations(self):
        """Test all operation types."""
        print(f"\n{'='*80}")
        print("STARTING METHODAI OPERATIONS TESTS - ALL OPERATION TYPES")
        print(f"{'='*80}")
        
        # Get all operation types
        allOperationTypes = list(OperationTypeEnum)
        
        # Filter to only operation types that have test configurations
        operationTypesToTest = [
            opType for opType in allOperationTypes 
            if opType in self.testPrompts
        ]
        
        print(f"Testing {len(operationTypesToTest)} operation type(s):")
        for i, opType in enumerate(operationTypesToTest, 1):
            print(f"  {i}. {opType.name}")
        
        print(f"\n{'='*80}")
        print("STARTING TESTS")
        print(f"{'='*80}\n")
        
        # Test each operation type
        for i, operationType in enumerate(operationTypesToTest, 1):
            print(f"\n{'─'*80}")
            print(f"[{i}/{len(operationTypesToTest)}] Testing: {operationType.name}")
            print(f"{'─'*80}")
            
            await self.testOperation(operationType)
            
            if i < len(operationTypesToTest):
                print(f"\n{'─'*80}")
        
        # Print summary
        self.printSummary()
    
    def printSummary(self):
        """Print test summary."""
        print(f"\n{'='*80}")
        print("TEST SUMMARY")
        print(f"{'='*80}")
        
        successfulTests = [r for r in self.testResults if r["status"] == "SUCCESS"]
        failedTests = [r for r in self.testResults if r["status"] == "ERROR"]
        exceptionTests = [r for r in self.testResults if r["status"] == "EXCEPTION"]
        
        print(f"\nTotal tests: {len(self.testResults)}")
        print(f"✅ Successful: {len(successfulTests)}")
        print(f"❌ Failed: {len(failedTests)}")
        print(f"💥 Exceptions: {len(exceptionTests)}")
        
        if successfulTests:
            print(f"\n{'─'*80}")
            print("SUCCESSFUL TESTS")
            print(f"{'─'*80}")
            for result in successfulTests:
                print(f"✅ {result['operationType']}: {result['processingTime']}s")
        
        if failedTests:
            print(f"\n{'─'*80}")
            print("FAILED TESTS")
            print(f"{'─'*80}")
            for result in failedTests:
                print(f"❌ {result['operationType']}: {result.get('error', 'Unknown error')}")
        
        if exceptionTests:
            print(f"\n{'─'*80}")
            print("EXCEPTIONS")
            print(f"{'─'*80}")
            for result in exceptionTests:
                print(f"💥 {result['operationType']}: {result.get('error', 'Unknown error')}")
        
        # Save results
        import json
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        resultsFile = os.path.join(self.modelTestDir, f"method_ai_operations_test_{timestamp}.json")
        
        with open(resultsFile, 'w', encoding='utf-8') as f:
            json.dump({
                "timestamp": timestamp,
                "summary": {
                    "total": len(self.testResults),
                    "successful": len(successfulTests),
                    "failed": len(failedTests),
                    "exceptions": len(exceptionTests)
                },
                "results": self.testResults
            }, f, indent=2, ensure_ascii=False)
        
        print(f"\n📄 Results saved to: {resultsFile}")


async def main():
    """Run methodAI operations tests."""
    tester = MethodAiOperationsTester()
    
    await tester.initialize()
    await tester.testAllOperations()
    
    print(f"\n{'='*80}")
    print("TESTING COMPLETED")
    print(f"{'='*80}")


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