#!/usr/bin/env python3
# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Code Generation Formats Test 11 - Tests code generation in JSON, CSV, and XML formats
Tests code generation with structured data formats including validation and formatting.
"""

import asyncio
import json
import sys
import os
import time
import csv
import io
import xml.etree.ElementTree as ET
from typing import Dict, Any, List, Optional

# 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)

# Import the service initialization
from modules.services import getInterface as getServices
from modules.aichat.datamodelFeatureAiChat import UserInputRequest, WorkflowModeEnum
from modules.datamodels.datamodelUam import User
from modules.workflows.automation import chatStart
import modules.aichat.interfaceFeatureAiChat as interfaceFeatureAiChat


class CodeGenerationFormatsTester11:
    def __init__(self):
        # Use root user for testing (has full access to everything)
        from modules.interfaces.interfaceDbApp import getRootInterface
        rootInterface = getRootInterface()
        self.testUser = rootInterface.currentUser
        
        # Initialize services using the existing system
        self.services = getServices(self.testUser, None)  # Test user, no workflow
        self.workflow = None
        self.testResults = {}
        self.generatedDocuments = {}
    
    async def initialize(self):
        """Initialize the test environment."""
        # Enable debug file logging for tests
        from modules.shared.configuration import APP_CONFIG
        APP_CONFIG.set("APP_DEBUG_CHAT_WORKFLOW_ENABLED", True)
        
        # Set logging level to INFO to see workflow progress
        import logging
        logging.getLogger().setLevel(logging.INFO)
        
        print(f"Initialized test with user: {self.testUser.id}")
        print(f"Mandate ID: {self.testUser.mandateId}")
        print(f"Debug logging enabled: {APP_CONFIG.get('APP_DEBUG_CHAT_WORKFLOW_ENABLED', False)}")
    
    def createTestPrompt(self, format: str) -> str:
        """Create a test prompt for code generation in the specified format.
        
        The prompt requests 3 files for each format:
        - Structured data generation appropriate for the format
        - Proper formatting and validation
        """
        formatPrompts = {
            "json": (
                "Generate 3 JSON code files for a customer management system:\n"
                "1) Create a config.json file with:\n"
                "   - Application name: 'Customer Manager'\n"
                "   - Version: '1.0.0'\n"
                "   - Database settings: host, port, name\n"
                "   - API settings: baseUrl, timeout\n"
                "2) Create a customers.json file with an array of customer objects:\n"
                "   - Each customer should have: id, name, email, phone, address\n"
                "   - Include at least 3 sample customers\n"
                "3) Create a settings.json file with:\n"
                "   - Theme settings: darkMode, fontSize, language\n"
                "   - Notification settings: email, sms, push\n"
                "   - Feature flags: enableAnalytics, enableReports\n\n"
                "Format all files as valid JSON with proper indentation."
            ),
            "csv": (
                "Generate 3 CSV code files for expense tracking:\n"
                "1) Create an expenses.csv file with:\n"
                "   - Header row: Documentname, Datum, Händler, Kreditkartennummer, Gesamtbetrag, Währung, MWST-Satz\n"
                "   - Data rows with at least 5 expense entries\n"
                "   - Use consistent date format (DD.MM.YYYY)\n"
                "   - Use CHF as currency\n"
                "   - Use 7.7% as VAT rate\n"
                "2) Create a categories.csv file with:\n"
                "   - Header row: CategoryID, CategoryName, Description, ParentCategory\n"
                "   - Data rows with at least 8 categories\n"
                "3) Create a vendors.csv file with:\n"
                "   - Header row: VendorID, VendorName, ContactPerson, Email, Phone, Address\n"
                "   - Data rows with at least 6 vendors\n\n"
                "Format all files as valid CSV with proper header row and consistent column count."
            ),
            "xml": (
                "Generate 3 XML code files for a product catalog:\n"
                "1) Create a products.xml file with:\n"
                "   - Root element: <catalog>\n"
                "   - Each product as <product> element with:\n"
                "     - <id>, <name>, <description>, <price>, <category>\n"
                "   - Include at least 4 products\n"
                "2) Create a categories.xml file with:\n"
                "   - Root element: <categories>\n"
                "   - Each category as <category> element with:\n"
                "     - <id>, <name>, <description>, <parentId>\n"
                "   - Include at least 5 categories\n"
                "3) Create a suppliers.xml file with:\n"
                "   - Root element: <suppliers>\n"
                "   - Each supplier as <supplier> element with:\n"
                "     - <id>, <name>, <contact>, <address>\n"
                "   - Include at least 3 suppliers\n\n"
                "Format all files as valid XML with proper indentation and structure."
            )
        }
        
        return formatPrompts.get(format.lower(), formatPrompts["json"])
    
    async def generateCodeInFormat(self, format: str) -> Dict[str, Any]:
        """Generate code in the specified format using workflow."""
        print("\n" + "="*80)
        print(f"GENERATING CODE IN {format.upper()} FORMAT")
        print("="*80)
        
        prompt = self.createTestPrompt(format)
        print(f"Prompt: {prompt[:200]}...")
        
        # Create user input request
        userInput = UserInputRequest(
            prompt=prompt,
            listFileId=[],
            userLanguage="en"
        )
        
        # Start workflow
        print(f"\nStarting workflow for {format.upper()} code generation...")
        workflow = await chatStart(
            currentUser=self.testUser,
            userInput=userInput,
            workflowMode=WorkflowModeEnum.WORKFLOW_DYNAMIC,
            workflowId=None
        )
        
        if not workflow:
            return {
                "success": False,
                "error": "Failed to start workflow"
            }
        
        self.workflow = workflow
        print(f"Workflow started: {workflow.id}")
        
        # Wait for workflow completion (no timeout - wait indefinitely)
        print(f"Waiting for workflow completion...")
        completed = await self.waitForWorkflowCompletion(timeout=None)
        
        if not completed:
            return {
                "success": False,
                "error": "Workflow did not complete",
                "workflowId": workflow.id,
                "status": workflow.status if workflow else "unknown"
            }
        
        # Analyze results
        results = self.analyzeWorkflowResults()
        
        # Extract documents for this format
        documents = results.get("documents", [])
        formatDocuments = [d for d in documents if d.get("fileName", "").endswith(f".{format.lower()}")]
        
        return {
            "success": True,
            "format": format,
            "workflowId": workflow.id,
            "status": results.get("status"),
            "documentCount": len(formatDocuments),
            "documents": formatDocuments,
            "results": results
        }
    
    async def waitForWorkflowCompletion(self, timeout: Optional[int] = None, checkInterval: int = 2) -> bool:
        """Wait for workflow to complete."""
        if not self.workflow:
            return False
        
        startTime = time.time()
        lastStatus = None
        
        interfaceDbChat = interfaceDbChat.getInterface(self.testUser)
        
        if timeout is None:
            print("Waiting indefinitely (no timeout)")
        
        while True:
            # Check timeout only if specified
            if timeout is not None and time.time() - startTime > timeout:
                print(f"\n⏱️ Timeout after {timeout} seconds")
                return False
            
            # Get current workflow status
            try:
                currentWorkflow = interfaceDbChat.getWorkflow(self.workflow.id)
                if not currentWorkflow:
                    print("\n❌ Workflow not found")
                    return False
                
                currentStatus = currentWorkflow.status
                elapsed = int(time.time() - startTime)
                
                # Print status if it changed
                if currentStatus != lastStatus:
                    print(f"Workflow status: {currentStatus} (elapsed: {elapsed}s)")
                    lastStatus = currentStatus
                
                # Check if workflow is complete
                if currentStatus in ["completed", "stopped", "failed"]:
                    self.workflow = currentWorkflow
                    statusIcon = "✅" if currentStatus == "completed" else "❌"
                    print(f"\n{statusIcon} Workflow finished with status: {currentStatus} (elapsed: {elapsed}s)")
                    return currentStatus == "completed"
                
                # Wait before next check
                await asyncio.sleep(checkInterval)
                
            except Exception as e:
                print(f"\n⚠️ Error checking workflow status: {str(e)}")
                await asyncio.sleep(checkInterval)
    
    def analyzeWorkflowResults(self) -> Dict[str, Any]:
        """Analyze workflow results and extract information."""
        if not self.workflow:
            return {"error": "No workflow to analyze"}
        
        interfaceDbChat = interfaceDbChat.getInterface(self.testUser)
        workflow = interfaceDbChat.getWorkflow(self.workflow.id)
        
        if not workflow:
            return {"error": "Workflow not found"}
        
        # Get unified chat data
        chatData = interfaceDbChat.getUnifiedChatData(workflow.id, None)
        
        # Count messages
        messages = chatData.get("messages", [])
        userMessages = [m for m in messages if m.get("role") == "user"]
        assistantMessages = [m for m in messages if m.get("role") == "assistant"]
        
        # Count documents
        documents = chatData.get("documents", [])
        
        # Get logs
        logs = chatData.get("logs", [])
        
        results = {
            "workflowId": workflow.id,
            "status": workflow.status,
            "workflowMode": str(workflow.workflowMode) if hasattr(workflow, 'workflowMode') else None,
            "currentRound": workflow.currentRound,
            "totalTasks": workflow.totalTasks,
            "totalActions": workflow.totalActions,
            "messageCount": len(messages),
            "userMessageCount": len(userMessages),
            "assistantMessageCount": len(assistantMessages),
            "documentCount": len(documents),
            "logCount": len(logs),
            "documents": documents,
            "logs": logs
        }
        
        print(f"\nWorkflow Results:")
        print(f"  Status: {results['status']}")
        print(f"  Tasks: {results['totalTasks']}")
        print(f"  Actions: {results['totalActions']}")
        print(f"  Messages: {results['messageCount']}")
        print(f"  Documents: {results['documentCount']}")
        
        # Print document details
        if documents:
            print(f"\nGenerated Documents:")
            for doc in documents:
                fileName = doc.get("fileName", "unknown")
                fileSize = doc.get("fileSize", 0)
                mimeType = doc.get("mimeType", "unknown")
                print(f"  - {fileName} ({fileSize} bytes, {mimeType})")
        
        return results
    
    
    def verifyCodeFormat(self, document: Dict[str, Any], expectedFormat: str) -> Dict[str, Any]:
        """Verify that a code file matches the expected format and is valid."""
        fileName = document.get("fileName", "")
        mimeType = document.get("mimeType", "")
        fileSize = document.get("fileSize", 0)
        
        # Expected MIME types
        expectedMimeTypes = {
            "json": ["application/json"],
            "csv": ["text/csv"],
            "xml": ["application/xml", "text/xml"]
        }
        
        # Expected file extensions
        expectedExtensions = {
            "json": [".json"],
            "csv": [".csv"],
            "xml": [".xml"]
        }
        
        formatLower = expectedFormat.lower()
        expectedMimes = expectedMimeTypes.get(formatLower, [])
        expectedExts = expectedExtensions.get(formatLower, [])
        
        # Check file extension
        hasCorrectExtension = any(fileName.lower().endswith(ext) for ext in expectedExts)
        
        # Check MIME type
        hasCorrectMimeType = any(mimeType.lower() == mime.lower() for mime in expectedMimes)
        
        # Check file size (should be > 0)
        hasValidSize = fileSize > 0
        
        # Try to read and validate content
        isValidContent = False
        validationError = None
        
        try:
            # Get file content from fileId
            fileId = document.get("fileId")
            if fileId and hasattr(self.services, 'interfaceDbComponent'):
                fileData = self.services.interfaceDbComponent.getFileData(fileId)
                if fileData:
                    content = fileData.decode('utf-8') if isinstance(fileData, bytes) else fileData
                    
                    # Validate format-specific syntax
                    if formatLower == "json":
                        try:
                            json.loads(content)
                            isValidContent = True
                        except json.JSONDecodeError as e:
                            validationError = f"Invalid JSON: {str(e)}"
                    
                    elif formatLower == "csv":
                        try:
                            reader = csv.reader(io.StringIO(content))
                            rows = list(reader)
                            if len(rows) > 0:
                                # Check header row exists
                                headerCount = len(rows[0])
                                # Check all rows have same column count
                                allRowsValid = all(len(row) == headerCount for row in rows)
                                isValidContent = allRowsValid
                                if not allRowsValid:
                                    validationError = "CSV rows have inconsistent column counts"
                            else:
                                validationError = "CSV file is empty"
                        except Exception as e:
                            validationError = f"CSV parsing error: {str(e)}"
                    
                    elif formatLower == "xml":
                        try:
                            ET.fromstring(content)
                            isValidContent = True
                        except ET.ParseError as e:
                            validationError = f"Invalid XML: {str(e)}"
                else:
                    validationError = "Could not read file data"
            else:
                validationError = "No fileId available"
                
        except Exception as e:
            validationError = f"Error reading/validating file: {str(e)}"
        
        verification = {
            "format": expectedFormat,
            "fileName": fileName,
            "mimeType": mimeType,
            "fileSize": fileSize,
            "hasCorrectExtension": hasCorrectExtension,
            "hasCorrectMimeType": hasCorrectMimeType,
            "hasValidSize": hasValidSize,
            "isValidContent": isValidContent,
            "validationError": validationError,
            "isValid": hasCorrectExtension and hasValidSize and hasCorrectMimeType,
            "isComplete": hasCorrectExtension and hasValidSize and hasCorrectMimeType and isValidContent
        }
        
        return verification
    
    async def testAllFormats(self) -> Dict[str, Any]:
        """Test code generation in JSON, CSV, and XML formats."""
        print("\n" + "="*80)
        print("TESTING CODE GENERATION IN ALL FORMATS")
        print("="*80)
        
        # Test all code formats
        formats = ["json", "csv", "xml"]
        results = {}
        
        for format in formats:
            try:
                print(f"\n{'='*80}")
                print(f"Testing {format.upper()} format...")
                print(f"{'='*80}")
                
                result = await self.generateCodeInFormat(format)
                results[format] = result
                
                if result.get("success"):
                    documents = result.get("documents", [])
                    if documents:
                        # Verify all documents (expecting 3 files per format)
                        verifications = []
                        for doc in documents:
                            verification = self.verifyCodeFormat(doc, format)
                            verifications.append(verification)
                        
                        result["verifications"] = verifications
                        
                        # Count valid documents
                        validCount = sum(1 for v in verifications if v.get("isValid"))
                        contentValidCount = sum(1 for v in verifications if v.get("isValidContent"))
                        
                        print(f"\n✅ {format.upper()} generation successful!")
                        print(f"   Documents: {len(documents)} (expected: 3)")
                        print(f"   Valid Format: {validCount}/{len(documents)}")
                        print(f"   Valid Content: {contentValidCount}/{len(documents)}")
                        
                        # Print details for each file
                        for i, verification in enumerate(verifications, 1):
                            statusIcon = "✅" if verification.get("isValid") else "❌"
                            contentIcon = "✅" if verification.get("isValidContent") else "❌"
                            print(f"   File {i}: {statusIcon} Format, {contentIcon} Content - {verification.get('fileName', 'unknown')}")
                            if verification.get("validationError"):
                                print(f"      Error: {verification['validationError']}")
                    else:
                        print(f"\n⚠️ {format.upper()} generation completed but no documents found")
                else:
                    error = result.get("error", "Unknown error")
                    print(f"\n❌ {format.upper()} generation failed: {error}")
                
                # Small delay between tests
                await asyncio.sleep(2)
                
            except Exception as e:
                import traceback
                print(f"\n❌ Error testing {format.upper()}: {str(e)}")
                print(traceback.format_exc())
                results[format] = {
                    "success": False,
                    "error": str(e),
                    "traceback": traceback.format_exc()
                }
        
        return results
    
    async def runTest(self):
        """Run the complete test."""
        print("\n" + "="*80)
        print("CODE GENERATION FORMATS TEST 11 - JSON, CSV, XML")
        print("="*80)
        
        try:
            # Initialize
            await self.initialize()
            
            # Test all formats
            formatResults = await self.testAllFormats()
            
            # Summary
            print("\n" + "="*80)
            print("TEST SUMMARY")
            print("="*80)
            
            # Format tests summary
            print("\nFormat Tests:")
            successCount = 0
            failCount = 0
            completeCount = 0  # Files with valid content
            
            for format, result in formatResults.items():
                if result.get("success"):
                    successCount += 1
                    verifications = result.get("verifications", [])
                    docCount = result.get("documentCount", 0)
                    
                    # Count valid files
                    validCount = sum(1 for v in verifications if v.get("isValid"))
                    contentValidCount = sum(1 for v in verifications if v.get("isValidContent"))
                    completeCount += contentValidCount
                    
                    # Overall status (all files valid)
                    allValid = len(verifications) > 0 and all(v.get("isValid") for v in verifications)
                    allContentValid = len(verifications) > 0 and all(v.get("isValidContent") for v in verifications)
                    
                    statusIcon = "✅" if allValid else "⚠️"
                    contentIcon = "✅" if allContentValid else "❌"
                    
                    print(f"{statusIcon} {format.upper():6s}: {'PASS' if allValid else 'PARTIAL'} - {docCount} file(s) ({validCount} valid format, {contentValidCount} valid content)")
                    
                    # Print errors if any
                    for v in verifications:
                        if v.get("validationError"):
                            print(f"      {v.get('fileName', 'unknown')}: {v['validationError']}")
                else:
                    failCount += 1
                    error = result.get("error", "Unknown error")
                    print(f"❌ {format.upper():6s}: FAIL - {error}")
            
            print(f"\nFormat Tests: {successCount} passed, {failCount} failed out of {len(formatResults)} formats")
            print(f"Valid Content Files: {completeCount} total files with valid content")
            
            self.testResults = {
                "success": failCount == 0,
                "formatTests": {
                    "successCount": successCount,
                    "failCount": failCount,
                    "completeCount": completeCount,
                    "totalFormats": len(formatResults),
                    "results": formatResults
                },
                "totalSuccess": successCount,
                "totalFail": failCount
            }
            
            return self.testResults
            
        except Exception as e:
            import traceback
            print(f"\n❌ Test failed with error: {type(e).__name__}: {str(e)}")
            print(f"Traceback:\n{traceback.format_exc()}")
            self.testResults = {
                "success": False,
                "error": str(e),
                "traceback": traceback.format_exc()
            }
            return self.testResults


async def main():
    """Run code generation formats test 11."""
    tester = CodeGenerationFormatsTester11()
    results = await tester.runTest()
    
    # Print final results as JSON for easy parsing
    print("\n" + "="*80)
    print("FINAL RESULTS (JSON)")
    print("="*80)
    print(json.dumps(results, indent=2, default=str))


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