#!/usr/bin/env python3
"""
Test script to verify the pipeline fixes work correctly
"""

import sys
import os
import traceback
from typing import Dict, Any
from PIL import Image
import numpy as np

# Add the project root to the path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# Import the pipeline at module level
try:
    from core.ai_pipeline import MonsterGenerationPipeline
    PIPELINE_AVAILABLE = True
except ImportError as e:
    print(f"⚠️ Warning: Could not import MonsterGenerationPipeline: {e}")
    PIPELINE_AVAILABLE = False

def test_pipeline_fixes():
    """Test the pipeline with improved error handling"""
    
    print("🧪 Testing Monster Generation Pipeline Fixes")
    print("=" * 50)
    
    try:
        # Initialize pipeline
        print("🔧 Initializing pipeline...")
        pipeline = MonsterGenerationPipeline(device="cpu")  # Use CPU for testing
        print("✅ Pipeline initialized successfully")
        
        # Test with a simple text input
        print("\n🚀 Testing monster generation...")
        test_input = "Create a friendly fire monster with wings"
        
        result = pipeline.generate_monster(
            text_input=test_input,
            user_id="test_user"
        )
        
        print(f"\n📊 Generation Result:")
        print(f"Status: {result.get('status', 'unknown')}")
        print(f"Success: {result.get('generation_log', {}).get('success', False)}")
        print(f"Stages completed: {result.get('generation_log', {}).get('stages_completed', [])}")
        print(f"Fallbacks used: {result.get('generation_log', {}).get('fallbacks_used', [])}")
        print(f"Errors: {result.get('generation_log', {}).get('errors', [])}")
        
        if result.get('traits'):
            print(f"Monster name: {result.get('traits', {}).get('name', 'Unknown')}")
            print(f"Monster element: {result.get('traits', {}).get('element', 'Unknown')}")
        
        if result.get('dialogue'):
            print(f"Monster dialogue: {result.get('dialogue', '')}")
        
        print(f"Download files: {result.get('download_files', [])}")
        
        # Clean up
        pipeline.cleanup()
        print("\n🧹 Pipeline cleaned up successfully")
        
        return True
        
    except Exception as e:
        print(f"❌ Test failed with error: {e}")
        print(f"Error type: {type(e).__name__}")
        print("Full traceback:")
        traceback.print_exc()
        return False

def test_fallback_manager():
    """Test the fallback manager"""
    
    print("\n🧪 Testing Fallback Manager")
    print("=" * 30)
    
    try:
        from utils.fallbacks import FallbackManager
        
        fallback = FallbackManager()
        
        # Test text generation fallback
        print("📝 Testing text generation fallback...")
        traits, dialogue = fallback.handle_text_gen_failure("Create a water monster")
        print(f"✅ Generated traits: {traits.get('name', 'Unknown')}")
        print(f"✅ Generated dialogue: {dialogue}")
        
        # Test image generation fallback
        print("🎨 Testing image generation fallback...")
        image = fallback.handle_image_gen_failure("Create a fire monster")
        print(f"✅ Generated image: {type(image)}")
        
        # Test 3D generation fallback
        print("🔲 Testing 3D generation fallback...")
        model_3d = fallback.handle_3d_gen_failure(image)
        print(f"✅ Generated 3D model: {type(model_3d)}")
        
        print("✅ All fallback tests passed!")
        return True
        
    except Exception as e:
        print(f"❌ Fallback test failed: {e}")
        traceback.print_exc()
        return False

def test_pipeline_timeout():
    """Test that the pipeline handles 3D generation timeout gracefully"""
    
    if not PIPELINE_AVAILABLE:
        print("⚠️ Skipping pipeline timeout test - pipeline not available")
        return False
    
    print("🧪 Testing pipeline timeout handling...")
    
    # Create a simple test image
    test_image = Image.new('RGB', (512, 512), color='red')
    
    # Initialize pipeline
    pipeline = MonsterGenerationPipeline(device="cpu")  # Use CPU for testing
    
    # Test with a simple text input
    result = pipeline.generate_monster(
        text_input="Create a simple red monster",
        user_id="test_user"
    )
    
    print(f"📊 Pipeline result status: {result.get('status', 'unknown')}")
    print(f"📊 Stages completed: {result.get('generation_log', {}).get('stages_completed', [])}")
    print(f"📊 Fallbacks used: {result.get('generation_log', {}).get('fallbacks_used', [])}")
    print(f"📊 Errors: {result.get('generation_log', {}).get('errors', [])}")
    
    # Check if we got a result
    if result.get('status') in ['success', 'fallback']:
        print("✅ Pipeline completed successfully!")
        if result.get('model_3d'):
            print(f"✅ 3D model generated: {result['model_3d']}")
        if result.get('image'):
            print(f"✅ Image generated: {type(result['image'])}")
        if result.get('traits'):
            print(f"✅ Monster traits: {result['traits'].get('name', 'Unknown')}")
    else:
        print("❌ Pipeline failed")
        return False
    
    return True

def test_3d_generator_timeout():
    """Test the 3D generator timeout mechanism directly"""
    
    print("\n🧪 Testing 3D generator timeout mechanism...")
    
    try:
        from models.model_3d_generator import Hunyuan3DGenerator
        
        # Create a test image
        test_image = Image.new('RGB', (512, 512), color='blue')
        
        # Initialize 3D generator with short timeout for testing
        generator = Hunyuan3DGenerator(device="cpu")
        generator.api_timeout = 10  # 10 seconds timeout for testing
        
        print("⏱️ Testing with 10-second timeout...")
        
        # This should either complete quickly or timeout
        result = generator.image_to_3d(test_image)
        
        print(f"✅ 3D generation completed: {type(result)}")
        return True
        
    except Exception as e:
        print(f"❌ 3D generation failed: {e}")
        if "timeout" in str(e).lower():
            print("✅ Timeout mechanism working correctly")
            return True
        else:
            print("❌ Unexpected error")
            return False

def main():
    """Main test function"""
    
    print("🔍 Starting Pipeline Fix Verification")
    print("=" * 50)
    
    # Test fallback manager first (doesn't require heavy models)
    fallback_success = test_fallback_manager()
    
    # Test 3D generator timeout mechanism
    timeout_success = test_3d_generator_timeout()
    
    # Test pipeline timeout handling
    pipeline_timeout_success = test_pipeline_timeout()
    
    # Test full pipeline (may fail due to missing models, but should show better error handling)
    pipeline_success = test_pipeline_fixes()
    
    print("\n" + "=" * 50)
    print("📋 Test Results Summary:")
    print(f"Fallback Manager: {'✅ PASSED' if fallback_success else '❌ FAILED'}")
    print(f"3D Generator Timeout: {'✅ PASSED' if timeout_success else '❌ FAILED'}")
    print(f"Pipeline Timeout: {'✅ PASSED' if pipeline_timeout_success else '❌ FAILED'}")
    print(f"Full Pipeline: {'✅ PASSED' if pipeline_success else '❌ FAILED'}")
    
    if fallback_success and timeout_success:
        print("\n🎉 Core timeout and fallback mechanisms are working!")
        if pipeline_success:
            print("🎉 Full pipeline is working correctly!")
        else:
            print("⚠️ Pipeline may need model dependencies, but timeout handling is functional.")
    else:
        print("\n❌ Some core tests failed. Check the error messages above.")
    
    return fallback_success and timeout_success

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)