"""
Test script for context-aware image generation (Task 7)
Tests the generate_with_context method
"""
import os
import sys
from unittest.mock import Mock, patch
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))

from agent import VisualizationAgent
from google import genai
from google.genai import types


def test_generate_with_context_basic():
    """Test basic context-aware generation"""
    print("\n=== Testing generate_with_context: Basic ===")
    try:
        # Create mock PIL Image
        mock_pil_image = Mock()
        mock_pil_image.save = Mock()
        
        # Create mock response part with inline_data
        mock_part = Mock()
        mock_part.inline_data = Mock()
        mock_part.as_image = Mock(return_value=mock_pil_image)
        
        # Create mock response
        mock_response = Mock()
        mock_response.parts = [mock_part]
        
        # Mock the genai.Client
        with patch('agent.genai.Client') as MockClient:
            mock_client_instance = Mock()
            mock_client_instance.models.generate_content.return_value = mock_response
            MockClient.return_value = mock_client_instance
            
            # Create agent
            agent = VisualizationAgent(api_key="test_key")
            
            # Test data
            construction_data = {
                "building_type": "residential_single_family",
                "location": {"name": "Manila"},
                "risk_data": {
                    "hazards": {
                        "seismic": {
                            "active_fault": {"status": "PRESENT"},
                            "ground_shaking": {"status": "HIGH"}
                        },
                        "volcanic": {},
                        "hydrometeorological": {}
                    }
                }
            }
            
            # Call generate_with_context
            result = agent.generate_with_context(
                prompt="A disaster-resistant building",
                construction_data=construction_data
            )
            
            # Verify success
            assert result["success"] is True, "Generation should succeed"
            assert "image_path" in result, "Missing image_path"
            
            # Verify that generate_content was called
            assert mock_client_instance.models.generate_content.called, "API should be called"
            
            print("✓ Basic context-aware generation successful")
            return True
            
    except Exception as e:
        print(f"✗ Basic context-aware test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_generate_with_context_building_type():
    """Test that building type is included in enhanced prompt"""
    print("\n=== Testing generate_with_context: Building Type ===")
    try:
        # Create mock PIL Image
        mock_pil_image = Mock()
        mock_pil_image.save = Mock()
        
        # Create mock response part
        mock_part = Mock()
        mock_part.inline_data = Mock()
        mock_part.as_image = Mock(return_value=mock_pil_image)
        
        mock_response = Mock()
        mock_response.parts = [mock_part]
        
        with patch('agent.genai.Client') as MockClient:
            mock_client_instance = Mock()
            mock_client_instance.models.generate_content.return_value = mock_response
            MockClient.return_value = mock_client_instance
            
            agent = VisualizationAgent(api_key="test_key")
            
            construction_data = {
                "building_type": "commercial_office"
            }
            
            result = agent.generate_with_context(
                prompt="Modern building",
                construction_data=construction_data
            )
            
            assert result["success"] is True, "Generation should succeed"
            
            # Check that the API was called with enhanced prompt
            call_args = mock_client_instance.models.generate_content.call_args
            prompt_used = call_args[1]["contents"]
            
            assert "Modern office building" in prompt_used, "Building type description should be in prompt"
            
            print("✓ Building type included in enhanced prompt")
            return True
            
    except Exception as e:
        print(f"✗ Building type test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_generate_with_context_location():
    """Test that location context is included"""
    print("\n=== Testing generate_with_context: Location ===")
    try:
        mock_pil_image = Mock()
        mock_pil_image.save = Mock()
        
        mock_part = Mock()
        mock_part.inline_data = Mock()
        mock_part.as_image = Mock(return_value=mock_pil_image)
        
        mock_response = Mock()
        mock_response.parts = [mock_part]
        
        with patch('agent.genai.Client') as MockClient:
            mock_client_instance = Mock()
            mock_client_instance.models.generate_content.return_value = mock_response
            MockClient.return_value = mock_client_instance
            
            agent = VisualizationAgent(api_key="test_key")
            
            construction_data = {
                "location": {"name": "Quezon City"}
            }
            
            result = agent.generate_with_context(
                prompt="Building design",
                construction_data=construction_data
            )
            
            assert result["success"] is True, "Generation should succeed"
            
            # Check prompt includes location
            call_args = mock_client_instance.models.generate_content.call_args
            prompt_used = call_args[1]["contents"]
            
            assert "Quezon City" in prompt_used, "Location should be in prompt"
            assert "Philippines" in prompt_used, "Philippines context should be in prompt"
            
            print("✓ Location context included in enhanced prompt")
            return True
            
    except Exception as e:
        print(f"✗ Location test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_generate_with_context_risk_data():
    """Test that disaster risk information is included"""
    print("\n=== Testing generate_with_context: Risk Data ===")
    try:
        mock_pil_image = Mock()
        mock_pil_image.save = Mock()
        
        mock_part = Mock()
        mock_part.inline_data = Mock()
        mock_part.as_image = Mock(return_value=mock_pil_image)
        
        mock_response = Mock()
        mock_response.parts = [mock_part]
        
        with patch('agent.genai.Client') as MockClient:
            mock_client_instance = Mock()
            mock_client_instance.models.generate_content.return_value = mock_response
            MockClient.return_value = mock_client_instance
            
            agent = VisualizationAgent(api_key="test_key")
            
            construction_data = {
                "risk_data": {
                    "hazards": {
                        "seismic": {
                            "active_fault": {"status": "PRESENT"},
                            "liquefaction": {"status": "HIGH"}
                        },
                        "volcanic": {},
                        "hydrometeorological": {
                            "flood": {"status": "HIGH"}
                        }
                    }
                }
            }
            
            result = agent.generate_with_context(
                prompt="Disaster-resistant structure",
                construction_data=construction_data
            )
            
            assert result["success"] is True, "Generation should succeed"
            
            # Check prompt includes disaster-resistant features
            call_args = mock_client_instance.models.generate_content.call_args
            prompt_used = call_args[1]["contents"]
            
            assert "Disaster-Resistant Features:" in prompt_used, "Features section should be in prompt"
            # Should have features related to seismic and flood risks
            assert any(keyword in prompt_used.lower() for keyword in ["seismic", "earthquake", "foundation", "flood", "elevated"]), \
                "Risk-specific features should be in prompt"
            
            print("✓ Risk data included in enhanced prompt")
            return True
            
    except Exception as e:
        print(f"✗ Risk data test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_generate_with_context_recommendations():
    """Test that recommendations are included"""
    print("\n=== Testing generate_with_context: Recommendations ===")
    try:
        mock_pil_image = Mock()
        mock_pil_image.save = Mock()
        
        mock_part = Mock()
        mock_part.inline_data = Mock()
        mock_part.as_image = Mock(return_value=mock_pil_image)
        
        mock_response = Mock()
        mock_response.parts = [mock_part]
        
        with patch('agent.genai.Client') as MockClient:
            mock_client_instance = Mock()
            mock_client_instance.models.generate_content.return_value = mock_response
            MockClient.return_value = mock_client_instance
            
            agent = VisualizationAgent(api_key="test_key")
            
            construction_data = {
                "recommendations": {
                    "priority_actions": [
                        "Use reinforced concrete",
                        "Install seismic bracing",
                        "Elevate foundation"
                    ]
                }
            }
            
            result = agent.generate_with_context(
                prompt="Safe building",
                construction_data=construction_data
            )
            
            assert result["success"] is True, "Generation should succeed"
            
            # Check prompt includes recommendations
            call_args = mock_client_instance.models.generate_content.call_args
            prompt_used = call_args[1]["contents"]
            
            assert "Construction Recommendations:" in prompt_used, "Recommendations section should be in prompt"
            
            print("✓ Recommendations included in enhanced prompt")
            return True
            
    except Exception as e:
        print(f"✗ Recommendations test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_generate_with_context_complete():
    """Test complete context with all data"""
    print("\n=== Testing generate_with_context: Complete Context ===")
    try:
        mock_pil_image = Mock()
        mock_pil_image.save = Mock()
        
        mock_part = Mock()
        mock_part.inline_data = Mock()
        mock_part.as_image = Mock(return_value=mock_pil_image)
        
        mock_response = Mock()
        mock_response.parts = [mock_part]
        
        with patch('agent.genai.Client') as MockClient:
            mock_client_instance = Mock()
            mock_client_instance.models.generate_content.return_value = mock_response
            MockClient.return_value = mock_client_instance
            
            agent = VisualizationAgent(api_key="test_key")
            
            # Complete construction data
            construction_data = {
                "building_type": "institutional_school",
                "location": {"name": "Cebu City"},
                "risk_data": {
                    "hazards": {
                        "seismic": {
                            "ground_shaking": {"status": "HIGH"}
                        },
                        "volcanic": {
                            "ashfall": {"status": "MODERATE"}
                        },
                        "hydrometeorological": {
                            "severe_winds": {"status": "HIGH"}
                        }
                    }
                },
                "recommendations": {
                    "priority_actions": [
                        "Implement seismic design",
                        "Use wind-resistant materials"
                    ]
                }
            }
            
            result = agent.generate_with_context(
                prompt="Educational facility",
                construction_data=construction_data
            )
            
            assert result["success"] is True, "Generation should succeed"
            
            # Check all context elements are in prompt
            call_args = mock_client_instance.models.generate_content.call_args
            prompt_used = call_args[1]["contents"]
            
            assert "School building" in prompt_used, "Building type should be in prompt"
            assert "Cebu City" in prompt_used, "Location should be in prompt"
            assert "Disaster-Resistant Features:" in prompt_used, "Risk features should be in prompt"
            assert "Construction Recommendations:" in prompt_used, "Recommendations should be in prompt"
            
            print("✓ Complete context included in enhanced prompt")
            print(f"  Prompt length: {len(prompt_used)} characters")
            return True
            
    except Exception as e:
        print(f"✗ Complete context test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_generate_with_context_config():
    """Test that config is passed through correctly"""
    print("\n=== Testing generate_with_context: Config Passing ===")
    try:
        mock_pil_image = Mock()
        mock_pil_image.save = Mock()
        
        mock_part = Mock()
        mock_part.inline_data = Mock()
        mock_part.as_image = Mock(return_value=mock_pil_image)
        
        mock_response = Mock()
        mock_response.parts = [mock_part]
        
        with patch('agent.genai.Client') as MockClient:
            mock_client_instance = Mock()
            mock_client_instance.models.generate_content.return_value = mock_response
            MockClient.return_value = mock_client_instance
            
            agent = VisualizationAgent(api_key="test_key")
            
            construction_data = {
                "building_type": "mixed_use"
            }
            
            config = {
                "aspect_ratio": "16:9",
                "image_size": "2K"
            }
            
            result = agent.generate_with_context(
                prompt="Mixed use building",
                construction_data=construction_data,
                config=config
            )
            
            assert result["success"] is True, "Generation should succeed"
            
            # Verify config was passed to API
            call_args = mock_client_instance.models.generate_content.call_args
            api_config = call_args[1]["config"]
            
            assert api_config.image_config.aspect_ratio == "16:9", "Aspect ratio should be passed"
            assert api_config.image_config.image_size == "2K", "Image size should be passed"
            
            print("✓ Config passed through correctly")
            return True
            
    except Exception as e:
        print(f"✗ Config passing test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_generate_with_context_error_handling():
    """Test error handling in generate_with_context"""
    print("\n=== Testing generate_with_context: Error Handling ===")
    try:
        with patch('agent.genai.Client') as MockClient:
            mock_client_instance = Mock()
            mock_client_instance.models.generate_content.side_effect = Exception("API Error")
            MockClient.return_value = mock_client_instance
            
            agent = VisualizationAgent(api_key="test_key")
            
            construction_data = {
                "building_type": "residential_high_rise"
            }
            
            result = agent.generate_with_context(
                prompt="High rise building",
                construction_data=construction_data
            )
            
            assert result["success"] is False, "Should fail on API error"
            assert "error" in result, "Should have error field"
            assert "code" in result["error"], "Error should have code"
            assert "message" in result["error"], "Error should have message"
            
            print("✓ Error handling works correctly")
            return True
            
    except Exception as e:
        print(f"✗ Error handling test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    print("=" * 60)
    print("Context-Aware Generation Test Suite (Task 7)")
    print("=" * 60)
    
    # Set test API key if not present
    if not os.getenv("GEMINI_API_KEY"):
        print("\n⚠ Warning: GEMINI_API_KEY not found in environment")
        print("Using mock API key for unit tests")
        os.environ["GEMINI_API_KEY"] = "test_api_key_for_unit_tests"
    
    # Run tests
    test_results = []
    test_results.append(("Basic Context-Aware Generation", test_generate_with_context_basic()))
    test_results.append(("Building Type Inclusion", test_generate_with_context_building_type()))
    test_results.append(("Location Context", test_generate_with_context_location()))
    test_results.append(("Risk Data Inclusion", test_generate_with_context_risk_data()))
    test_results.append(("Recommendations Inclusion", test_generate_with_context_recommendations()))
    test_results.append(("Complete Context", test_generate_with_context_complete()))
    test_results.append(("Config Passing", test_generate_with_context_config()))
    test_results.append(("Error Handling", test_generate_with_context_error_handling()))
    
    # Summary
    print("\n" + "=" * 60)
    print("TEST SUMMARY")
    print("=" * 60)
    
    passed = sum(1 for _, result in test_results if result)
    total = len(test_results)
    
    for test_name, result in test_results:
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"  {status}: {test_name}")
    
    print(f"\nTotal: {passed}/{total} tests passed")
    print("=" * 60)
    
    if passed == total:
        print("\n✅ All tests passed!")
        sys.exit(0)
    else:
        print(f"\n❌ {total - passed} test(s) failed")
        sys.exit(1)