#!/usr/bin/env python3
"""
Simple test script to verify Vertex AI multimodal embeddings work correctly.
Run this script to test the embedding generation before deploying.
"""

import os
import sys
import logging
from PIL import Image
import io
import numpy as np

# Add the current directory to the path so we can import main
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def create_test_image():
    """Create a simple test image"""
    # Create a simple 100x100 RGB image with a red square
    img = Image.new('RGB', (100, 100), color='white')
    pixels = img.load()
    
    # Draw a red square in the center
    for i in range(25, 75):
        for j in range(25, 75):
            pixels[i, j] = (255, 0, 0)  # Red
    
    # Convert to bytes
    img_bytes = io.BytesIO()
    img.save(img_bytes, format='PNG')
    return img_bytes.getvalue()

def test_vertex_ai_embeddings():
    """Test Vertex AI embedding generation"""
    try:
        # Set required environment variables for testing
        os.environ['GOOGLE_CLOUD_PROJECT'] = os.environ.get('GOOGLE_CLOUD_PROJECT', 'your-project-id')
        os.environ['VERTEX_AI_LOCATION'] = os.environ.get('VERTEX_AI_LOCATION', 'us-central1')
        
        # Import the function we want to test
        from main import generate_image_embeddings
        
        # Create test image
        logger.info("Creating test image...")
        test_image_data = create_test_image()
        logger.info(f"Created test image with {len(test_image_data)} bytes")
        
        # Generate embeddings
        logger.info("Generating embeddings using Vertex AI...")
        embeddings = generate_image_embeddings(test_image_data)
        
        if embeddings is None:
            logger.error("Failed to generate embeddings!")
            return False
        
        # Verify embeddings
        logger.info(f"Generated embeddings with shape: {embeddings.shape}")
        logger.info(f"Embeddings dtype: {embeddings.dtype}")
        logger.info(f"Embeddings range: [{embeddings.min():.4f}, {embeddings.max():.4f}]")
        logger.info(f"Embeddings norm: {np.linalg.norm(embeddings):.4f}")
        
        # Basic validation
        assert isinstance(embeddings, np.ndarray), "Embeddings should be numpy array"
        assert embeddings.dtype == np.float32, "Embeddings should be float32"
        assert len(embeddings.shape) == 1, "Embeddings should be 1D array"
        assert embeddings.shape[0] == 1408, f"Expected 1408 dimensions, got {embeddings.shape[0]}"
        
        # Check if normalized (should be close to 1.0)
        norm = np.linalg.norm(embeddings)
        assert 0.9 <= norm <= 1.1, f"Embeddings should be normalized, norm is {norm}"
        
        logger.info("✅ All tests passed! Vertex AI embeddings are working correctly.")
        return True
        
    except ImportError as e:
        logger.error(f"Import error: {e}")
        logger.error("Make sure you have installed the required dependencies:")
        logger.error("pip install google-cloud-aiplatform")
        return False
    except Exception as e:
        logger.error(f"Test failed with error: {e}")
        return False

if __name__ == "__main__":
    logger.info("Testing Vertex AI multimodal embeddings...")
    
    # Check if required environment variables are set
    project_id = os.environ.get('GOOGLE_CLOUD_PROJECT')
    if not project_id:
        logger.error("Please set GOOGLE_CLOUD_PROJECT environment variable")
        logger.error("Example: export GOOGLE_CLOUD_PROJECT=your-project-id")
        sys.exit(1)
    
    logger.info(f"Using project: {project_id}")
    
    success = test_vertex_ai_embeddings()
    
    if success:
        logger.info("🎉 Test completed successfully!")
        sys.exit(0)
    else:
        logger.error("❌ Test failed!")
        sys.exit(1)