#!/opt/homebrew/bin/python3.11
"""
Fast Face Clustering Processor (Linear Scan)
職責：針對長片優化，使用線性讀取取代隨機跳轉，大幅提升速度。
"""

import cv2
import json
import numpy as np
import os
import sys
import psycopg2
from collections import defaultdict

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

try:
    from deepface import DeepFace

    HAS_DEEPFACE = True
except ImportError:
    print("❌ DeepFace not found.")
    sys.exit(1)

from sklearn.cluster import AgglomerativeClustering

# 設定
UUID = os.getenv("UUID", "384b0ff44aaaa1f1")
OUTPUT_DIR = os.getenv("MOMENTRY_OUTPUT_DIR", "./output")
VIDEO_PATH = os.path.join(OUTPUT_DIR, UUID, f"{UUID}.mp4")
FACE_JSON_PATH = os.path.join(OUTPUT_DIR, UUID, f"{UUID}.face.json")
OUTPUT_JSON_PATH = os.path.join(OUTPUT_DIR, UUID, f"{UUID}.face_clustered.json")
ASRX_JSON_PATH = os.path.join(OUTPUT_DIR, UUID, f"{UUID}.asrx.json")
DB_URL = os.getenv("DATABASE_URL", "postgresql://accusys@localhost:5432/momentry")


def main():
    if not os.path.exists(FACE_JSON_PATH):
        print(f"❌ Face JSON not found: {FACE_JSON_PATH}")
        return

    print(f"⚡ 開始執行快速面孔聚類 (Linear Scan Mode) for {UUID}...")

    # 1. 載入並建立索引 (以 frame number 為 key)
    with open(FACE_JSON_PATH) as f:
        face_data = json.load(f)

    frames_list = face_data.get("frames", [])
    if not frames_list:
        print("❌ No frames in JSON.")
        return

    # 建立 map: frame_index -> faces
    # 注意：JSON 中的 frame 是 int，但也許是 float?
    # face_processor 輸出通常是 int
    faces_map = defaultdict(list)

    # 為了安全，我們也建立 timestamp map 以防萬一，但優先使用 frame number
    print(f"📂 Indexing {len(frames_list)} frames with faces...")
    for frame_obj in frames_list:
        # JSON 中可能是 'frame' (int) 或 'frame_number'
        idx = frame_obj.get("frame") or frame_obj.get("frame_number")
        if idx is not None:
            faces_map[int(idx)].extend(frame_obj.get("faces", []))

    # 如果沒有 frame number 字段，我們只能依靠 timestamp (比較慢)
    if not faces_map:
        print("⚠️ No frame numbers found in JSON. Falling back to timestamp seeking.")
        # 這裡我們可以呼叫舊的邏輯，但為了簡單，我們假設 face_processor 有寫 frame
        # 檢查第一個 frame 的 key
        if frames_list:
            print(f"   Keys: {frames_list[0].keys()}")
        return  # 暫時中斷

    total_faces = sum(len(faces) for faces in faces_map.values())
    print(f"✅ Indexed {len(faces_map)} frames, containing {total_faces} faces.")
    print("🚀 Starting Linear Video Scan...")

    # 2. 線性掃描
    video_path = VIDEO_PATH  # 使用區域變數避免 global 問題
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        # 嘗試找 mov
        alt_path = video_path.replace(".mp4", ".mov")
        if os.path.exists(alt_path):
            video_path = alt_path
            cap = cv2.VideoCapture(video_path)
        else:
            print("❌ Video file not found.")
            return

    embeddings = []
    face_refs = []  # 存儲 (frame_index, face_index_in_list)

    # 為了追蹤進度
    processed_frames = 0
    current_frame = 0

    # 獲取影片總幀數
    total_video_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # 檢查這一幀是否有我們需要處理的臉
        # 使用 round 處理可能的浮點誤差 (雖然 face_processor 應該寫的是 int)
        # 如果 JSON 的 frame 是 0.0, 1.0...
        # 這裡我們直接看 current_frame 是否在 faces_map 中

        # 由於 face_processor 可能跳幀，或者時間戳對齊問題
        # 我們檢查 current_frame 以及 current_frame +/- 1 的容差
        # 但最好的方式是嚴格匹配 frame number

        if current_frame in faces_map:
            faces = faces_map[current_frame]
            for face_idx, face in enumerate(faces):
                try:
                    x, y, w, h = face["x"], face["y"], face["width"], face["height"]
                    margin = 5
                    crop = frame[
                        max(0, y - margin) : y + h + margin,
                        max(0, x - margin) : x + w + margin,
                    ]

                    if crop is not None and crop.size > 0:
                        # 使用 Fast Model: VGG-Face 或 OpenFace 比 ArcFace 快，但 ArcFace 準
                        # 這裡保持 ArcFace 以求準確，但因為是線性讀取，省去了 seek 時間
                        # 為了速度，我們可以每 2 秒只取 1 幀？
                        # 不，我們需要標記所有幀。
                        # DeepFace 提取
                        res = DeepFace.represent(
                            img_path=crop, model_name="ArcFace", enforce_detection=False
                        )
                        if res and "embedding" in res[0]:
                            embeddings.append(res[0]["embedding"])
                            face_refs.append(
                                {"frame_idx": current_frame, "face_idx": face_idx}
                            )
                except Exception:
                    pass

            processed_frames += 1
            if processed_frames % 500 == 0:
                pct = (current_frame / total_video_frames) * 100
                print(
                    f"   📊 Progress: Frame {current_frame}/{total_video_frames} ({pct:.1f}%) | Extracted: {len(embeddings)} embeddings"
                )

        current_frame += 1

    cap.release()

    if not embeddings:
        print("❌ No embeddings extracted.")
        return

    embeddings = np.array(embeddings)
    print(f"✅ Total Embeddings Extracted: {len(embeddings)}")

    # 3. 聚類
    print(f"🧠 Clustering {len(embeddings)} faces...")

    # 優化：KMeans 或 MiniBatchKMeans 對於大數據集更快
    # 但 Agglomerative 對於找任意形狀的簇更好。
    # 25000 個點做層次聚類還是慢。
    # 我們使用 "Sample -> Cluster -> Assign" 策略

    print("   🚀 Using Sampling Strategy for speed...")
    sample_size = 5000
    n_faces = len(embeddings)

    if n_faces > sample_size:
        indices = np.random.choice(n_faces, sample_size, replace=False)
        sample_embeddings = embeddings[indices]
    else:
        sample_embeddings = embeddings
        indices = np.arange(n_faces)

    clustering = AgglomerativeClustering(
        n_clusters=None, distance_threshold=0.45, metric="cosine", linkage="average"
    )
    sample_labels = clustering.fit_predict(sample_embeddings)

    # 計算簇中心
    unique_labels = set(sample_labels)
    centroids = []
    for label in unique_labels:
        mask = sample_labels == label
        centroids.append(np.mean(sample_embeddings[mask], axis=0))
    centroids = np.array(centroids)

    # 分配所有數據
    print("   🏃 Assigning remaining faces to clusters...")
    from sklearn.metrics.pairwise import cosine_distances

    # 批次計算
    all_labels = np.zeros(n_faces, dtype=int)
    batch_size = 10000
    for i in range(0, n_faces, batch_size):
        batch = embeddings[i : i + batch_size]
        dists = cosine_distances(batch, centroids)
        all_labels[i : i + batch_size] = np.argmin(dists, axis=1)

    print(f"   👥 Detected {len(unique_labels)} unique persons.")

    # 4. 生成標籤
    label_to_person = {l: f"Person_{i}" for i, l in enumerate(unique_labels)}

    # 5. 寫回 JSON
    # face_data 是原始結構，我們需要修改它
    # face_data['frames'] 是一個列表
    # 我們需要快速找到對應的 frame

    # 建立 map frame_idx -> frame_object reference
    frame_ref_map = {}
    for f_obj in face_data.get("frames", []):
        idx = f_obj.get("frame") or f_obj.get("frame_number")
        if idx is not None:
            frame_ref_map[int(idx)] = f_obj

    count = 0
    for ref, label in zip(face_refs, all_labels):
        f_idx = ref["frame_idx"]
        face_idx = ref["face_idx"]  # 這是原始 faces list 中的 index

        person_id = label_to_person[label]

        if f_idx in frame_ref_map:
            frame_obj = frame_ref_map[f_idx]
            faces_list = frame_obj.get("faces", [])
            if face_idx < len(faces_list):
                faces_list[face_idx]["person_id"] = person_id
                count += 1

    print(f"   ✅ Tagged {count} faces with Person ID.")

    with open(OUTPUT_JSON_PATH, "w", encoding="utf-8") as f:
        json.dump(face_data, f, indent=2, ensure_ascii=False)
    print(f"✅ Saved clustered data to {OUTPUT_JSON_PATH}")

    # 6. 綁定 Speaker
    auto_bind_speakers()


def auto_bind_speakers():
    if not os.path.exists(OUTPUT_JSON_PATH) or not os.path.exists(ASRX_JSON_PATH):
        print("⚠️ Missing data for speaker binding.")
        return

    with open(OUTPUT_JSON_PATH) as f:
        face_clustered = json.load(f)
    with open(ASRX_JSON_PATH) as f:
        asrx_data = json.load(f)

    print("🔗 Auto-binding Speakers to Persons...")

    face_spans = []
    for frame_obj in face_clustered.get("frames", []):
        ts = frame_obj.get("timestamp")
        for face in frame_obj.get("faces", []):
            person_id = face.get("person_id")
            if person_id and ts is not None:
                face_spans.append({"ts": ts, "person_id": person_id})

    speaker_person_counts = {}

    for seg in asrx_data.get("segments", []):
        start = seg.get("start")
        end = seg.get("end")
        speaker = seg.get("speaker_id")
        if not speaker:
            continue

        candidates = [f for f in face_spans if start <= f["ts"] <= end]
        if candidates:
            person_counts = {}
            for c in candidates:
                pid = c["person_id"]
                person_counts[pid] = person_counts.get(pid, 0) + 1

            if speaker not in speaker_person_counts:
                speaker_person_counts[speaker] = {}

            best_person = max(person_counts, key=person_counts.get)
            speaker_person_counts[speaker][best_person] = (
                speaker_person_counts[speaker].get(best_person, 0) + 1
            )

    try:
        conn = psycopg2.connect(DB_URL)
        cur = conn.cursor()

        for speaker, persons in speaker_person_counts.items():
            if not persons:
                continue
            best_person = max(persons, key=persons.get)
            print(
                f"   🎤 {speaker} is likely {best_person} ({persons[best_person]} votes)"
            )

            cur.execute("SELECT id FROM talents WHERE real_name = %s", (best_person,))
            row = cur.fetchone()

            if row:
                talent_id = row[0]
            else:
                cur.execute(
                    "INSERT INTO talents (real_name) VALUES (%s) RETURNING id",
                    (best_person,),
                )
                talent_id = cur.fetchone()[0]
                print(f"   ✨ Created Talent #{talent_id} ({best_person})")

            cur.execute(
                """
                INSERT INTO identity_bindings (talent_id, binding_type, binding_value, source, confidence)
                VALUES (%s, 'speaker', %s, 'auto_cluster', 0.8)
                ON CONFLICT (binding_type, binding_value) DO UPDATE SET talent_id = EXCLUDED.talent_id
            """,
                (talent_id, speaker),
            )
            print(f"   ✅ Bound {speaker} -> {best_person}")

        conn.commit()
        cur.close()
        conn.close()
    except Exception as e:
        print(f"   ❌ DB Error: {e}")


if __name__ == "__main__":
    main()
