#!/usr/bin/env python3
"""
端到端人臉識別測試
測試完整的人臉識別流程：註冊 -> 識別 -> 搜索
"""

import os
import sys
import numpy as np
import cv2

# 添加項目根目錄到 Python 路徑
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def create_test_image_with_faces():
    """創建帶有人臉的測試圖像"""
    print("創建測試圖像...")

    # 創建一個簡單的測試圖像（640x480）
    img = np.zeros((480, 640, 3), dtype=np.uint8)

    # 添加一些"人臉"（簡單的橢圓形）
    # 人臉1
    cv2.ellipse(img, (200, 200), (80, 100), 0, 0, 360, (255, 200, 150), -1)
    cv2.ellipse(img, (170, 170), (15, 10), 0, 0, 360, (0, 0, 0), -1)  # 左眼
    cv2.ellipse(img, (230, 170), (15, 10), 0, 0, 360, (0, 0, 0), -1)  # 右眼
    cv2.ellipse(img, (200, 230), (30, 15), 0, 0, 360, (0, 0, 0), -1)  # 嘴巴

    # 人臉2
    cv2.ellipse(img, (450, 300), (70, 90), 0, 0, 360, (200, 220, 180), -1)
    cv2.ellipse(img, (420, 270), (12, 8), 0, 0, 360, (0, 0, 0), -1)  # 左眼
    cv2.ellipse(img, (480, 270), (12, 8), 0, 0, 360, (0, 0, 0), -1)  # 右眼
    cv2.ellipse(img, (450, 330), (25, 12), 0, 0, 360, (0, 0, 0), -1)  # 嘴巴

    # 保存測試圖像
    test_image_path = "/tmp/test_face_image.jpg"
    cv2.imwrite(test_image_path, img)
    print(f"✅ 測試圖像已保存到: {test_image_path}")

    return test_image_path, img


def test_face_registration():
    """測試人臉註冊"""
    print("\n=== 測試人臉註冊 ===")

    try:
        from scripts.face_registration import FaceRegistration

        # 創建測試圖像
        image_path, img = create_test_image_with_faces()

        # 初始化註冊器
        print("初始化人臉註冊器...")
        registration = FaceRegistration()

        # 加載模型
        print("加載模型...")
        registration.load_models(use_mps=False)

        # 註冊人臉
        print("註冊人臉...")
        result = registration.register_face(
            image_path=image_path,
            name="Test Person 1",
            metadata={"source": "test", "age": 30, "gender": "male"},
        )

        if result["success"]:
            print("✅ 人臉註冊成功")
            print(f"  - Face ID: {result.get('face_id')}")
            print(f"  - 嵌入向量維度: {len(result.get('embedding', []))}")
            print(f"  - 屬性: {result.get('attributes', {})}")

            # 保存嵌入向量供後續測試使用
            embedding = result.get("embedding", [])
            if embedding:
                np.save("/tmp/test_face_embedding.npy", embedding)
                print("✅ 嵌入向量已保存")

            return True, result
        else:
            print(f"❌ 人臉註冊失敗: {result.get('message', 'Unknown error')}")
            return False, None

    except Exception as e:
        print(f"❌ 人臉註冊測試失敗: {e}")
        import traceback

        traceback.print_exc()
        return False, None


def test_face_recognition():
    """測試人臉識別"""
    print("\n=== 測試人臉識別 ===")

    try:
        from scripts.face_recognition_processor import FaceRecognitionProcessor

        # 創建測試圖像
        image_path, img = create_test_image_with_faces()

        # 初始化處理器
        print("初始化人臉識別處理器...")
        processor = FaceRecognitionProcessor(
            enable_recognition=True, enable_tracking=True, enable_clustering=True
        )

        # 加載模型
        print("加載模型...")
        processor.load_models(use_mps=False)

        # 讀取圖像
        print("讀取測試圖像...")
        image = cv2.imread(image_path)
        if image is None:
            print("❌ 無法讀取測試圖像")
            return False

        # 檢測人臉
        print("檢測人臉...")
        detections = processor.detect_faces(image)

        print(f"✅ 檢測到 {len(detections)} 個人臉")

        if len(detections) > 0:
            for i, detection in enumerate(detections):
                print(f"\n人臉 {i + 1}:")
                print(
                    f"  - 位置: x={detection['x']}, y={detection['y']}, width={detection['width']}, height={detection['height']}"
                )
                print(f"  - 置信度: {detection['confidence']:.4f}")

                if "embedding" in detection and detection["embedding"] is not None:
                    embedding = detection["embedding"]
                    if hasattr(embedding, "shape"):
                        print(f"  - 嵌入向量維度: {embedding.shape}")
                    else:
                        print(f"  - 嵌入向量長度: {len(embedding)}")

                if "attributes" in detection:
                    attrs = detection["attributes"]
                    print(f"  - 屬性: {attrs}")

        # 測試人臉追蹤（模擬多幀）
        print("\n測試人臉追蹤...")
        # 創建模擬幀數據
        frames = [
            {"frame_id": 1, "faces": detections},
            {"frame_id": 2, "faces": detections},  # 簡單重複使用相同的檢測
        ]
        tracked_frames = processor.track_faces(frames)
        print(f"✅ 追蹤完成，處理了 {len(tracked_frames)} 幀")

        # 測試人臉聚類
        print("\n測試人臉聚類...")
        # 創建模擬多幀檢測數據
        all_frames = []
        for i in range(3):  # 模擬3幀
            frame_data = {
                "frame_id": i + 1,
                "faces": detections if i == 0 else [],  # 只在第一幀有檢測
            }
            all_frames.append(frame_data)

        clusters = processor.cluster_faces(all_frames)
        print(f"✅ 創建 {len(clusters)} 個聚類")

        return True

    except Exception as e:
        print(f"❌ 人臉識別測試失敗: {e}")
        import traceback

        traceback.print_exc()
        return False


def test_database_operations():
    """測試數據庫操作"""
    print("\n=== 測試數據庫操作 ===")

    try:
        import psycopg2
        from psycopg2.extras import Json

        # 連接數據庫
        conn = psycopg2.connect(
            host="localhost",
            port=5432,
            database="momentry",
            user="accusys",
            password="accusys",
        )

        cursor = conn.cursor()

        # 測試1: 插入人臉檢測記錄
        print("測試插入人臉檢測記錄...")
        cursor.execute(
            """
            INSERT INTO face_detections 
            (video_uuid, frame_number, timestamp_secs, face_id, x, y, width, height, confidence, attributes)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING id;
        """,
            (
                "test_video_001",
                1,
                0.0,
                "test_face_001",
                100,
                100,
                200,
                200,
                0.95,
                Json({"age": 30, "gender": "male", "test": True}),
            ),
        )

        detection_id = cursor.fetchone()[0]
        print(f"✅ 插入人臉檢測記錄成功，ID: {detection_id}")

        # 測試2: 查詢人臉檢測記錄
        print("\n測試查詢人臉檢測記錄...")
        cursor.execute(
            """
            SELECT id, video_uuid, frame_number, face_id, confidence, attributes
            FROM face_detections 
            WHERE id = %s;
        """,
            (detection_id,),
        )

        result = cursor.fetchone()
        print("✅ 查詢結果:")
        print(f"  - ID: {result[0]}")
        print(f"  - 視頻UUID: {result[1]}")
        print(f"  - 幀號: {result[2]}")
        print(f"  - 人臉ID: {result[3]}")
        print(f"  - 置信度: {result[4]}")
        print(f"  - 屬性: {result[5]}")

        # 測試3: 測試向量搜索函數
        print("\n測試向量搜索函數...")

        # 創建一個測試嵌入向量
        test_embedding = np.random.randn(512).tolist()

        # 首先插入一個帶有嵌入向量的人臉身份
        cursor.execute(
            """
            SELECT find_or_create_face_identity(
                'test_search_001',
                'Search Test Person',
                %s::vector,
                '{"age": 25, "gender": "female", "test": true}'::jsonb,
                '{"source": "search_test"}'::jsonb
            );
        """,
            (test_embedding,),
        )

        identity_id = cursor.fetchone()[0]
        print(f"✅ 創建測試人臉身份，ID: {identity_id}")

        # 測試搜索相似人臉
        cursor.execute(
            """
            SELECT * FROM find_similar_faces(
                %s::vector,
                0.5,  -- similarity_threshold
                5     -- limit_count
            );
        """,
            (test_embedding,),
        )

        similar_faces = cursor.fetchall()
        print(f"✅ 找到 {len(similar_faces)} 個相似人臉")

        for face in similar_faces:
            print(f"  - {face[0]}: {face[1]} (相似度: {face[2]:.4f})")

        # 清理測試數據
        print("\n清理測試數據...")
        cursor.execute(
            "DELETE FROM face_detections WHERE video_uuid = 'test_video_001';"
        )
        cursor.execute("DELETE FROM face_identities WHERE face_id LIKE 'test_%';")
        conn.commit()
        print("✅ 測試數據清理完成")

        cursor.close()
        conn.close()

        return True

    except Exception as e:
        print(f"❌ 數據庫操作測試失敗: {e}")
        import traceback

        traceback.print_exc()
        return False


def test_mps_acceleration():
    """測試 MPS 加速"""
    print("\n=== 測試 MPS 加速 ===")

    try:
        import onnxruntime as ort

        available_providers = ort.get_available_providers()
        print(f"可用的執行提供者: {available_providers}")

        if "CoreMLExecutionProvider" in available_providers:
            print("✅ CoreML (MPS) 支援可用")

            # 測試使用 MPS 初始化模型
            from scripts.face_recognition_processor import FaceRecognitionProcessor

            print("測試使用 MPS 初始化模型...")
            processor = FaceRecognitionProcessor()

            try:
                processor.load_models(use_mps=True)
                print("✅ MPS 模型加載成功")

                # 測試推理
                test_image = np.random.randint(0, 255, (640, 480, 3), dtype=np.uint8)
                detections = processor.detect_faces(test_image)
                print(f"✅ MPS 推理完成，檢測到 {len(detections)} 個人臉")

                return True

            except Exception as e:
                print(f"⚠️  MPS 初始化失敗，回退到 CPU: {e}")
                print("嘗試使用 CPU...")
                processor.load_models(use_mps=False)
                print("✅ CPU 模型加載成功")
                return True

        else:
            print("⚠️  CoreML (MPS) 不可用，使用 CPU")
            return True

    except Exception as e:
        print(f"❌ MPS 測試失敗: {e}")
        return False


def main():
    """主測試函數"""
    print("=" * 60)
    print("端到端人臉識別測試")
    print("=" * 60)

    tests = [
        ("人臉註冊", test_face_registration),
        ("人臉識別", test_face_recognition),
        ("數據庫操作", test_database_operations),
        ("MPS 加速", test_mps_acceleration),
    ]

    results = []

    for test_name, test_func in tests:
        try:
            print(f"\n{'=' * 40}")
            print(f"開始測試: {test_name}")
            print(f"{'=' * 40}")

            success = test_func()
            results.append((test_name, success))

            if success:
                print(f"✅ {test_name} 測試通過")
            else:
                print(f"❌ {test_name} 測試失敗")

        except Exception as e:
            print(f"❌ {test_name} 測試異常: {e}")
            import traceback

            traceback.print_exc()
            results.append((test_name, False))

    print("\n" + "=" * 60)
    print("端到端測試結果摘要")
    print("=" * 60)

    passed = 0
    for test_name, success in results:
        status = "✅ 通過" if success else "❌ 失敗"
        print(f"{test_name}: {status}")
        if success:
            passed += 1

    print(f"\n總計: {passed}/{len(results)} 個測試通過")

    if passed == len(results):
        print("\n🎉 所有端到端測試通過！人臉識別系統完全可用。")
        print("\n下一步:")
        print("1. 啟動 Momentry 服務器")
        print("2. 使用 API 端點進行人臉註冊和識別")
        print("3. 測試視頻處理功能")
        return 0
    else:
        print(f"\n⚠️  有 {len(results) - passed} 個測試失敗，請檢查問題。")
        return 1


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