#!/usr/bin/env python3
"""
DeepFace Integration Test Suite for PunimTag
Tests the complete integration of DeepFace into the application
"""

import os
import sys
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

# Suppress TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import warnings
warnings.filterwarnings('ignore')

from src.core.database import DatabaseManager
from src.core.face_processing import FaceProcessor
from src.core.config import DEEPFACE_DETECTOR_BACKEND, DEEPFACE_MODEL_NAME


def test_face_detection():
    """Test 1: Face detection with DeepFace"""
    print("\n" + "="*60)
    print("Test 1: DeepFace Face Detection")
    print("="*60)
    
    try:
        db = DatabaseManager(":memory:", verbose=0)  # In-memory database for testing
        processor = FaceProcessor(db, verbose=1)
        
        # Test with a sample image
        test_image = "demo_photos/2019-11-22_0011.jpg"
        if not os.path.exists(test_image):
            print(f"❌ Test image not found: {test_image}")
            print("   Please ensure demo photos are available")
            return False
        
        print(f"Testing with image: {test_image}")
        
        # Add photo to database
        photo_id = db.add_photo(test_image, Path(test_image).name, None)
        print(f"✓ Added photo to database (ID: {photo_id})")
        
        # Process faces
        count = processor.process_faces(limit=1)
        print(f"✓ Processed {count} photos")
        
        # Verify results
        stats = db.get_statistics()
        print(f"✓ Found {stats['total_faces']} faces in the photo")
        
        if stats['total_faces'] == 0:
            print("❌ FAIL: No faces detected")
            return False
        
        # Verify face encodings are 512-dimensional (ArcFace)
        with db.get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT encoding FROM faces LIMIT 1")
            encoding_blob = cursor.fetchone()[0]
            encoding_size = len(encoding_blob)
            expected_size = 512 * 8  # 512 floats * 8 bytes per float
            
            print(f"✓ Encoding size: {encoding_size} bytes (expected: {expected_size})")
            
            if encoding_size != expected_size:
                print(f"❌ FAIL: Wrong encoding size (expected {expected_size}, got {encoding_size})")
                return False
        
        print("\n✅ PASS: Face detection working correctly")
        return True
        
    except Exception as e:
        print(f"\n❌ FAIL: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_face_matching():
    """Test 2: Face matching with DeepFace"""
    print("\n" + "="*60)
    print("Test 2: DeepFace Face Matching")
    print("="*60)
    
    try:
        db = DatabaseManager(":memory:", verbose=0)
        processor = FaceProcessor(db, verbose=1)
        
        # Test with multiple images
        test_images = [
            "demo_photos/2019-11-22_0011.jpg",
            "demo_photos/2019-11-22_0012.jpg"
        ]
        
        # Check if test images exist
        available_images = [img for img in test_images if os.path.exists(img)]
        
        if len(available_images) < 2:
            print(f"⚠️  Only {len(available_images)} test images available")
            print("   Skipping face matching test (need at least 2 images)")
            return True  # Skip but don't fail
        
        print(f"Testing with {len(available_images)} images")
        
        # Add photos to database
        for img in available_images:
            photo_id = db.add_photo(img, Path(img).name, None)
            print(f"✓ Added {Path(img).name} (ID: {photo_id})")
        
        # Process all faces
        count = processor.process_faces(limit=10)
        print(f"✓ Processed {count} photos")
        
        # Get statistics
        stats = db.get_statistics()
        print(f"✓ Found {stats['total_faces']} total faces")
        
        if stats['total_faces'] < 2:
            print("⚠️  Not enough faces for matching test")
            return True  # Skip but don't fail
        
        # Find similar faces
        faces = db.get_all_face_encodings()
        if len(faces) >= 2:
            face_id = faces[0][0]
            print(f"✓ Testing similarity for face ID {face_id}")
            
            matches = processor.find_similar_faces(face_id, tolerance=0.4)
            print(f"✓ Found {len(matches)} similar faces (within tolerance)")
            
            # Display match details
            if matches:
                for i, match in enumerate(matches[:3], 1):  # Show top 3 matches
                    confidence_pct = (1 - match['distance']) * 100
                    print(f"   Match {i}: Face {match['face_id']}, Confidence: {confidence_pct:.1f}%")
        
        print("\n✅ PASS: Face matching working correctly")
        return True
        
    except Exception as e:
        print(f"\n❌ FAIL: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_deepface_metadata():
    """Test 3: DeepFace metadata storage and retrieval"""
    print("\n" + "="*60)
    print("Test 3: DeepFace Metadata Storage")
    print("="*60)
    
    try:
        db = DatabaseManager(":memory:", verbose=0)
        processor = FaceProcessor(db, verbose=1)
        
        # Test with a sample image
        test_image = "demo_photos/2019-11-22_0011.jpg"
        if not os.path.exists(test_image):
            print(f"⚠️  Test image not found: {test_image}")
            return True  # Skip but don't fail
        
        # Add photo and process
        photo_id = db.add_photo(test_image, Path(test_image).name, None)
        processor.process_faces(limit=1)
        
        # Query face metadata
        with db.get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT face_confidence, quality_score, detector_backend, model_name
                FROM faces
                LIMIT 1
            """)
            result = cursor.fetchone()
            
            if not result:
                print("❌ FAIL: No face metadata found")
                return False
            
            face_conf, quality, detector, model = result
            
            print(f"✓ Face Confidence: {face_conf}")
            print(f"✓ Quality Score: {quality}")
            print(f"✓ Detector Backend: {detector}")
            print(f"✓ Model Name: {model}")
            
            # Verify metadata is present
            if detector is None:
                print("❌ FAIL: Detector backend not stored")
                return False
            
            if model is None:
                print("❌ FAIL: Model name not stored")
                return False
            
            # Verify detector matches configuration
            if detector != DEEPFACE_DETECTOR_BACKEND:
                print(f"⚠️  Warning: Detector mismatch (expected {DEEPFACE_DETECTOR_BACKEND}, got {detector})")
            
            # Verify model matches configuration
            if model != DEEPFACE_MODEL_NAME:
                print(f"⚠️  Warning: Model mismatch (expected {DEEPFACE_MODEL_NAME}, got {model})")
        
        print("\n✅ PASS: DeepFace metadata stored correctly")
        return True
        
    except Exception as e:
        print(f"\n❌ FAIL: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_configuration():
    """Test 4: FaceProcessor configuration with different backends"""
    print("\n" + "="*60)
    print("Test 4: FaceProcessor Configuration")
    print("="*60)
    
    try:
        db = DatabaseManager(":memory:", verbose=0)
        
        # Test default configuration
        processor_default = FaceProcessor(db, verbose=0)
        print(f"✓ Default detector: {processor_default.detector_backend}")
        print(f"✓ Default model: {processor_default.model_name}")
        
        if processor_default.detector_backend != DEEPFACE_DETECTOR_BACKEND:
            print(f"❌ FAIL: Default detector mismatch")
            return False
        
        if processor_default.model_name != DEEPFACE_MODEL_NAME:
            print(f"❌ FAIL: Default model mismatch")
            return False
        
        # Test custom configuration
        custom_configs = [
            ('mtcnn', 'Facenet512'),
            ('opencv', 'VGG-Face'),
            ('ssd', 'ArcFace'),
        ]
        
        for detector, model in custom_configs:
            processor = FaceProcessor(db, verbose=0, 
                                    detector_backend=detector,
                                    model_name=model)
            print(f"✓ Custom config: {detector}/{model}")
            
            if processor.detector_backend != detector:
                print(f"❌ FAIL: Custom detector not applied")
                return False
            
            if processor.model_name != model:
                print(f"❌ FAIL: Custom model not applied")
                return False
        
        print("\n✅ PASS: FaceProcessor configuration working correctly")
        return True
        
    except Exception as e:
        print(f"\n❌ FAIL: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_cosine_similarity():
    """Test 5: Cosine similarity calculation"""
    print("\n" + "="*60)
    print("Test 5: Cosine Similarity Calculation")
    print("="*60)
    
    try:
        import numpy as np
        
        db = DatabaseManager(":memory:", verbose=0)
        processor = FaceProcessor(db, verbose=0)
        
        # Test with identical encodings
        encoding1 = np.random.rand(512).astype(np.float64)
        encoding2 = encoding1.copy()
        
        distance = processor._calculate_cosine_similarity(encoding1, encoding2)
        print(f"✓ Identical encodings distance: {distance:.6f}")
        
        if distance > 0.01:  # Should be very close to 0
            print(f"❌ FAIL: Identical encodings should have distance near 0")
            return False
        
        # Test with different encodings
        encoding3 = np.random.rand(512).astype(np.float64)
        distance2 = processor._calculate_cosine_similarity(encoding1, encoding3)
        print(f"✓ Different encodings distance: {distance2:.6f}")
        
        if distance2 < 0.1:  # Should be significantly different
            print(f"⚠️  Warning: Random encodings have low distance (might be coincidence)")
        
        # Test with mismatched lengths
        encoding4 = np.random.rand(128).astype(np.float64)
        distance3 = processor._calculate_cosine_similarity(encoding1, encoding4)
        print(f"✓ Mismatched lengths distance: {distance3:.6f}")
        
        if distance3 != 2.0:  # Should return max distance
            print(f"❌ FAIL: Mismatched lengths should return 2.0")
            return False
        
        print("\n✅ PASS: Cosine similarity calculation working correctly")
        return True
        
    except Exception as e:
        print(f"\n❌ FAIL: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_database_schema():
    """Test 6: Database schema validation"""
    print("\n" + "="*60)
    print("Test 6: Database Schema Validation")
    print("="*60)
    
    try:
        db = DatabaseManager(":memory:", verbose=0)
        
        # Check if new DeepFace columns exist
        with db.get_db_connection() as conn:
            cursor = conn.cursor()
            
            # Get faces table schema
            cursor.execute("PRAGMA table_info(faces)")
            columns = {row[1]: row[2] for row in cursor.fetchall()}
            
            print("✓ Faces table columns:")
            for col_name in columns:
                print(f"  - {col_name}")
            
            # Verify DeepFace columns
            required_columns = {
                'detector_backend': 'TEXT',
                'model_name': 'TEXT',
                'face_confidence': 'REAL'
            }
            
            for col, dtype in required_columns.items():
                if col not in columns:
                    print(f"❌ FAIL: Missing column '{col}' in faces table")
                    return False
                print(f"✓ Column '{col}' exists with type {columns[col]}")
            
            # Check person_encodings table
            cursor.execute("PRAGMA table_info(person_encodings)")
            pe_columns = {row[1]: row[2] for row in cursor.fetchall()}
            
            print("\n✓ Person_encodings table columns:")
            for col_name in pe_columns:
                print(f"  - {col_name}")
            
            # Verify DeepFace columns in person_encodings
            pe_required = {
                'detector_backend': 'TEXT',
                'model_name': 'TEXT',
            }
            
            for col, dtype in pe_required.items():
                if col not in pe_columns:
                    print(f"❌ FAIL: Missing column '{col}' in person_encodings table")
                    return False
                print(f"✓ Column '{col}' exists in person_encodings")
        
        print("\n✅ PASS: Database schema is correct")
        return True
        
    except Exception as e:
        print(f"\n❌ FAIL: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_face_location_format():
    """Test 7: Face location format validation"""
    print("\n" + "="*60)
    print("Test 7: Face Location Format")
    print("="*60)
    
    try:
        import ast
        
        db = DatabaseManager(":memory:", verbose=0)
        processor = FaceProcessor(db, verbose=1)
        
        # Test with a sample image
        test_image = "demo_photos/2019-11-22_0011.jpg"
        if not os.path.exists(test_image):
            print(f"⚠️  Test image not found: {test_image}")
            return True  # Skip but don't fail
        
        # Add photo and process
        photo_id = db.add_photo(test_image, Path(test_image).name, None)
        processor.process_faces(limit=1)
        
        # Check face location format
        with db.get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT location FROM faces LIMIT 1")
            result = cursor.fetchone()
            
            if not result:
                print("⚠️  No faces found")
                return True
            
            location_str = result[0]
            print(f"✓ Raw location: {location_str}")
            
            # Parse location
            try:
                location = ast.literal_eval(location_str)
                print(f"✓ Parsed location: {location}")
                
                # Check if it's DeepFace format (dict with x, y, w, h)
                if isinstance(location, dict):
                    required_keys = ['x', 'y', 'w', 'h']
                    for key in required_keys:
                        if key not in location:
                            print(f"❌ FAIL: Missing key '{key}' in location dict")
                            return False
                    print("✓ Location is in DeepFace dict format {x, y, w, h}")
                else:
                    print(f"❌ FAIL: Location is not a dict, got {type(location)}")
                    return False
                
            except Exception as e:
                print(f"❌ FAIL: Could not parse location: {e}")
                return False
        
        print("\n✅ PASS: Face location format is correct")
        return True
        
    except Exception as e:
        print(f"\n❌ FAIL: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_performance_benchmark():
    """Test 8: Performance benchmarking"""
    print("\n" + "="*60)
    print("Test 8: Performance Benchmark")
    print("="*60)
    
    try:
        import time
        
        db = DatabaseManager(":memory:", verbose=0)
        processor = FaceProcessor(db, verbose=0)
        
        # Test with multiple images
        test_images = [
            "demo_photos/2019-11-22_0011.jpg",
            "demo_photos/2019-11-22_0012.jpg",
            "demo_photos/2019-11-22_0015.jpg",
        ]
        
        available_images = [img for img in test_images if os.path.exists(img)]
        
        if not available_images:
            print("⚠️  No test images available")
            return True  # Skip but don't fail
        
        print(f"Testing with {len(available_images)} images")
        
        # Add photos to database
        for img in available_images:
            db.add_photo(img, Path(img).name, None)
        
        # Benchmark face detection
        start_time = time.time()
        count = processor.process_faces(limit=len(available_images))
        detection_time = time.time() - start_time
        
        print(f"✓ Processed {count} photos in {detection_time:.2f}s")
        print(f"✓ Average time per photo: {detection_time/max(count, 1):.2f}s")
        
        # Get statistics
        stats = db.get_statistics()
        total_faces = stats['total_faces']
        print(f"✓ Found {total_faces} total faces")
        
        if total_faces > 0:
            print(f"✓ Average time per face: {detection_time/total_faces:.2f}s")
        
        # Benchmark similarity calculation
        if total_faces >= 2:
            faces = db.get_all_face_encodings()
            face_id = faces[0][0]
            
            start_time = time.time()
            matches = processor.find_similar_faces(face_id, tolerance=0.4)
            matching_time = time.time() - start_time
            
            print(f"✓ Similarity search completed in {matching_time:.2f}s")
            print(f"✓ Found {len(matches)} matches")
        
        print("\n✅ PASS: Performance benchmark completed")
        return True
        
    except Exception as e:
        print(f"\n❌ FAIL: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_adaptive_tolerance():
    """Test 9: Adaptive tolerance calculation"""
    print("\n" + "="*60)
    print("Test 9: Adaptive Tolerance")
    print("="*60)
    
    try:
        db = DatabaseManager(":memory:", verbose=0)
        processor = FaceProcessor(db, verbose=0)
        
        # Test with different quality scores
        base_tolerance = 0.4
        
        test_cases = [
            (0.1, "Low quality"),
            (0.5, "Medium quality"),
            (0.9, "High quality"),
        ]
        
        print(f"Base tolerance: {base_tolerance}")
        
        for quality, desc in test_cases:
            tolerance = processor._calculate_adaptive_tolerance(base_tolerance, quality)
            print(f"✓ {desc} ({quality:.1f}): tolerance = {tolerance:.3f}")
            
            # Verify tolerance is within bounds
            if tolerance < 0.2 or tolerance > 0.6:
                print(f"❌ FAIL: Tolerance {tolerance} out of bounds [0.2, 0.6]")
                return False
        
        # Test with match confidence
        tolerance_with_conf = processor._calculate_adaptive_tolerance(
            base_tolerance, 0.7, match_confidence=0.8
        )
        print(f"✓ With match confidence: tolerance = {tolerance_with_conf:.3f}")
        
        print("\n✅ PASS: Adaptive tolerance working correctly")
        return True
        
    except Exception as e:
        print(f"\n❌ FAIL: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_multiple_detectors():
    """Test 10: Multiple detector backends"""
    print("\n" + "="*60)
    print("Test 10: Multiple Detector Backends")
    print("="*60)
    
    try:
        # Test different detector backends
        detectors = ['opencv', 'ssd']  # Skip retinaface and mtcnn for speed
        test_image = "demo_photos/2019-11-22_0011.jpg"
        
        if not os.path.exists(test_image):
            print("⚠️  Test image not found")
            return True  # Skip but don't fail
        
        results = {}
        
        for detector in detectors:
            print(f"\n Testing with {detector} detector:")
            
            try:
                db = DatabaseManager(":memory:", verbose=0)
                processor = FaceProcessor(db, verbose=0, 
                                        detector_backend=detector,
                                        model_name='ArcFace')
                
                photo_id = db.add_photo(test_image, Path(test_image).name, None)
                count = processor.process_faces(limit=1)
                
                stats = db.get_statistics()
                faces_found = stats['total_faces']
                
                results[detector] = faces_found
                print(f"✓ {detector}: Found {faces_found} faces")
                
            except Exception as e:
                print(f"⚠️  {detector} failed: {e}")
                results[detector] = 0
        
        # Verify at least one detector worked
        if sum(results.values()) == 0:
            print("\n❌ FAIL: No detectors found any faces")
            return False
        
        print("\n✅ PASS: Multiple detectors tested")
        return True
        
    except Exception as e:
        print(f"\n❌ FAIL: {e}")
        import traceback
        traceback.print_exc()
        return False


def run_all_tests():
    """Run all DeepFace integration tests"""
    print("\n" + "="*70)
    print("DEEPFACE INTEGRATION TEST SUITE - PHASE 6")
    print("="*70)
    print()
    print("Testing complete DeepFace integration in PunimTag")
    print("This comprehensive test suite validates all aspects of the migration")
    print()
    
    tests = [
        ("Face Detection", test_face_detection),
        ("Face Matching", test_face_matching),
        ("Metadata Storage", test_deepface_metadata),
        ("Configuration", test_configuration),
        ("Cosine Similarity", test_cosine_similarity),
        ("Database Schema", test_database_schema),
        ("Face Location Format", test_face_location_format),
        ("Performance Benchmark", test_performance_benchmark),
        ("Adaptive Tolerance", test_adaptive_tolerance),
        ("Multiple Detectors", test_multiple_detectors),
    ]
    
    results = []
    for test_name, test_func in tests:
        try:
            result = test_func()
            results.append((test_name, result))
        except Exception as e:
            print(f"\n❌ Test '{test_name}' crashed: {e}")
            import traceback
            traceback.print_exc()
            results.append((test_name, False))
    
    # Print summary
    print("\n" + "="*70)
    print("TEST SUMMARY")
    print("="*70)
    
    passed = 0
    failed = 0
    for test_name, result in results:
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"{status}: {test_name}")
        if result:
            passed += 1
        else:
            failed += 1
    
    print("="*70)
    print(f"Tests passed: {passed}/{len(tests)}")
    print(f"Tests failed: {failed}/{len(tests)}")
    print("="*70)
    
    if failed == 0:
        print("\n🎉 ALL TESTS PASSED! DeepFace integration is working correctly!")
        return 0
    else:
        print(f"\n⚠️  {failed} test(s) failed. Please review the errors above.")
        return 1


if __name__ == "__main__":
    sys.exit(run_all_tests())

