#!/usr/bin/env python3
"""
Test face learning (registration) and recognition
This demonstrates the system's ability to learn new faces
"""

import requests
import json
import os
from PIL import Image
import io

BASE_URL = "http://localhost:3002"
API_KEY = "muser_243c6725b09f43e29f319a648645b992_1774874668_f224a6d2"
VIDEO_UUID = "384b0ff44aaaa1f1"


def extract_face_from_frame(frame_path, face_bbox):
    """Extract a face from frame image"""
    if not os.path.exists(frame_path):
        print(f"⚠️ Frame not found: {frame_path}")
        return None

    try:
        img = Image.open(frame_path)

        # Extract face region (x, y, width, height)
        x, y, w, h = face_bbox
        face_img = img.crop((x, y, x + w, y + h))

        # Convert to bytes
        img_byte_arr = io.BytesIO()
        face_img.save(img_byte_arr, format="JPEG")
        img_byte_arr.seek(0)

        return img_byte_arr
    except Exception as e:
        print(f"❌ Error extracting face: {e}")
        return None


def register_face_from_detection(frame_number, face_index, person_name):
    """Register a face from detected face in video"""
    print(
        f"\n📝 Registering {person_name} from frame {frame_number}, face {face_index}..."
    )

    # First, get face detection details
    import psycopg2

    try:
        conn = psycopg2.connect(
            host="localhost",
            port=5432,
            database="momentry",
            user="accusys",
            password="accusys",
        )
        cursor = conn.cursor()

        cursor.execute(
            """
            SELECT x, y, width, height, confidence, attributes::text
            FROM face_detections 
            WHERE video_uuid = %s AND frame_number = %s
            ORDER BY confidence DESC
            LIMIT 1 OFFSET %s
        """,
            (VIDEO_UUID, frame_number, face_index),
        )

        result = cursor.fetchone()

        if not result:
            print(f"❌ No face found at frame {frame_number}, index {face_index}")
            return False

        x, y, width, height, confidence, attributes_json = result

        # Parse attributes
        attributes = json.loads(attributes_json) if attributes_json else {}

        print(f"  Face bbox: ({x}, {y}, {width}, {height})")
        print(f"  Confidence: {confidence:.3f}")
        print(f"  Attributes: {attributes}")

        # Register via API
        headers = {"X-API-Key": API_KEY, "Content-Type": "application/json"}

        payload = {
            "video_uuid": VIDEO_UUID,
            "frame_number": frame_number,
            "face_index": face_index,
            "person_name": person_name,
            "metadata": {
                "gender": attributes.get("gender", "unknown"),
                "age": attributes.get("age", 0),
                "confidence": float(confidence),
                "source": "Charade (1963)",
                "bbox": [x, y, width, height],
            },
        }

        response = requests.post(
            f"{BASE_URL}/api/v1/face/register",
            headers=headers,
            json=payload,
            timeout=30,
        )

        print(f"  API Status: {response.status_code}")

        if response.status_code == 200:
            data = response.json()
            print(f"  ✅ Registered! Face ID: {data.get('face_id')}")
            print(f"     Embedding: {len(data.get('embedding', []))} dimensions")
            return True
        else:
            print(f"  ❌ Registration failed: {response.text}")
            return False

    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback

        traceback.print_exc()
        return False
    finally:
        if "cursor" in locals():
            cursor.close()
        if "conn" in locals():
            conn.close()


def test_face_recognition(test_image_path):
    """Test face recognition with a test image"""
    print(f"\n🔍 Testing recognition with: {test_image_path}")

    if not os.path.exists(test_image_path):
        print(f"❌ Test image not found: {test_image_path}")
        return False

    headers = {"X-API-Key": API_KEY}

    files = {"image": open(test_image_path, "rb")}

    try:
        response = requests.post(
            f"{BASE_URL}/api/v1/face/recognize",
            headers=headers,
            files=files,
            timeout=30,
        )

        print(f"Status: {response.status_code}")

        if response.status_code == 200:
            data = response.json()
            print("✅ Success!")
            print(f"Total faces detected: {data.get('total_faces', 0)}")

            matches = data.get("matches", [])
            if matches:
                print("\n🎯 Recognition results:")
                for i, match in enumerate(matches):
                    print(f"  {i + 1}. {match.get('person_name', 'Unknown')}")
                    print(f"     Confidence: {match.get('confidence', 0):.3f}")
                    print(f"     Distance: {match.get('distance', 0):.3f}")

                    metadata = match.get("metadata", {})
                    if metadata:
                        print(f"     Gender: {metadata.get('gender', 'unknown')}")
                        print(f"     Age: {metadata.get('age', 'unknown')}")
                    print()
            else:
                print("No matches found")

            return True
        else:
            print(f"❌ Recognition failed: {response.text}")
            return False

    except Exception as e:
        print(f"❌ Error: {e}")
        return False
    finally:
        if "files" in locals():
            files["image"].close()


def list_registered_faces():
    """List all registered faces"""
    print("\n📋 Listing registered faces...")

    headers = {"X-API-Key": API_KEY}

    try:
        response = requests.get(
            f"{BASE_URL}/api/v1/face/list", headers=headers, timeout=10
        )

        print(f"Status: {response.status_code}")

        if response.status_code == 200:
            data = response.json()
            faces = data.get("faces", [])
            print(f"Total registered faces: {len(faces)}")

            for face in faces:
                print(f"\n  👤 {face.get('name')}")
                print(f"     ID: {face.get('face_id')}")
                print(f"     Created: {face.get('created_at')}")
                print(f"     Active: {face.get('is_active', False)}")

                metadata = face.get("metadata", {})
                if metadata:
                    print(f"     Gender: {metadata.get('gender', 'unknown')}")
                    print(f"     Age: {metadata.get('age', 'unknown')}")

            return True
        else:
            print(f"❌ Failed: {response.text}")
            return False

    except Exception as e:
        print(f"❌ Error: {e}")
        return False


def main():
    print("=" * 70)
    print("🧠 Face Learning and Recognition Test")
    print("=" * 70)

    print(f"\n🎬 Source video: {VIDEO_UUID}")
    print("   Charade (1963) - Audrey Hepburn and Cary Grant")

    # Step 1: List current registered faces
    print("\n" + "-" * 50)
    print("STEP 1: Check current registered faces")
    print("-" * 50)
    list_registered_faces()

    # Step 2: Register sample faces
    print("\n" + "-" * 50)
    print("STEP 2: Register faces from video analysis")
    print("-" * 50)

    # Register faces from our analysis
    # Frame 19778 has 3 faces (based on our analysis)
    faces_to_learn = [
        {
            "frame": 19778,
            "index": 0,
            "name": "Audrey_Hepburn",
            "description": "Main actress in Charade",
        },
        {
            "frame": 19778,
            "index": 1,
            "name": "Cary_Grant",
            "description": "Main actor in Charade",
        },
        {
            "frame": 17980,
            "index": 0,
            "name": "Walter_Mathau",
            "description": "Supporting actor in Charade",
        },
    ]

    learned_count = 0
    for face in faces_to_learn:
        if register_face_from_detection(face["frame"], face["index"], face["name"]):
            learned_count += 1

    print(f"\n📊 Learned {learned_count}/{len(faces_to_learn)} faces")

    # Step 3: List registered faces again
    print("\n" + "-" * 50)
    print("STEP 3: Verify registered faces")
    print("-" * 50)
    list_registered_faces()

    # Step 4: Test recognition
    print("\n" + "-" * 50)
    print("STEP 4: Test face recognition")
    print("-" * 50)

    # Test with the female faces frame we extracted earlier
    test_images = [
        "/tmp/female_faces/female_faces_frame_19778.jpg",
        "/tmp/face_analysis_results/384b0ff44aaaa1f1_frame_19778.jpg",
    ]

    for test_image in test_images:
        if os.path.exists(test_image):
            test_face_recognition(test_image)
            break

    # Step 5: Test with a different frame
    print("\n" + "-" * 50)
    print("STEP 5: Test with another frame")
    print("-" * 50)

    # Try frame 17980 (has 2 females)
    test_image_2 = "/tmp/face_analysis_results/384b0ff44aaaa1f1_frame_17980.jpg"
    if os.path.exists(test_image_2):
        test_face_recognition(test_image_2)

    print("\n" + "=" * 70)
    print("✅ Face Learning Test Completed!")
    print("=" * 70)

    # Summary
    print("\n📋 SUMMARY:")
    print(f"• API Server: {BASE_URL}")
    print(f"• Video analyzed: {VIDEO_UUID}")
    print("• Faces detected: 78 (from analysis)")
    print(f"• Faces registered: {learned_count}")
    print("• System ready for learning new faces!")

    print("\n💡 NEXT STEPS:")
    print("1. Upload new photos to recognize registered faces")
    print("2. Register more faces from other videos")
    print("3. Build a face database for your organization")
    print("4. Integrate with your applications via API")


if __name__ == "__main__":
    main()