#!/opt/homebrew/bin/python3.11
"""
Sync Face & Speaker IDs to Chunks
將 face.json 和 asrx.json 中的機器 ID 聚合寫入 chunks 表的 face_ids 和 speaker_ids 欄位。
"""

import sys
import json
import os
import argparse
import psycopg2
from psycopg2.extras import execute_values

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


def get_db_connection():
    db_url = os.getenv("DATABASE_URL", "postgresql://accusys@localhost:5432/momentry")
    return psycopg2.connect(db_url)


def load_json_safe(path):
    if not os.path.exists(path):
        return None
    with open(path, "r") as f:
        return json.load(f)


def sync_video(conn, uuid: str, output_dir: str):
    print(f"Syncing video: {uuid}")

    # 1. 加載 JSON 數據
    face_data = load_json_safe(os.path.join(output_dir, f"{uuid}.face.json"))
    asrx_data = load_json_safe(os.path.join(output_dir, f"{uuid}.asrx.json"))

    if not face_data and not asrx_data:
        print(f"  No face or asrx JSON found for {uuid}")
        return

    # 2. 獲取該視頻的所有 chunks
    cur = conn.cursor()
    cur.execute(
        "SELECT id, start_frame, end_frame, fps FROM chunks WHERE uuid = %s", (uuid,)
    )
    chunks = cur.fetchall()

    if not chunks:
        print(f"  No chunks found for {uuid}")
        return

    print(f"  Found {len(chunks)} chunks to process.")

    updates = []

    # 3. 遍歷 chunks 並匹配 ID
    for chunk_id, start_frame, end_frame, fps in chunks:
        start_sec = start_frame / fps
        end_sec = end_frame / fps

        face_ids = []
        speaker_ids = []

        # 匹配 Face IDs
        if face_data and "frames" in face_data:
            for frame in face_data["frames"]:
                # 簡單判斷：如果幀的時間在 chunk 範圍內
                # 注意：這裡假設 frame 有 timestamp 欄位
                ts = frame.get("timestamp", 0.0)
                if start_sec <= ts < end_sec:
                    for face in frame.get("faces", []):
                        fid = face.get("face_id") or face.get("id")
                        if fid and fid not in face_ids:
                            face_ids.append(fid)

        # 匹配 Speaker IDs (ASRX)
        if asrx_data and "segments" in asrx_data:
            for seg in asrx_data["segments"]:
                # 判斷時間重疊
                seg_start = seg.get("start", 0.0)
                seg_end = seg.get("end", 0.0)

                # 重疊條件: seg_start < end_sec AND seg_end > start_sec
                if seg_start < end_sec and seg_end > start_sec:
                    sid = seg.get("speaker_id")
                    if sid and sid not in speaker_ids:
                        speaker_ids.append(sid)

        if face_ids or speaker_ids:
            updates.append((face_ids, speaker_ids, chunk_id))

    # 4. 批量更新
    if updates:
        query = """
            UPDATE chunks 
            SET face_ids = %s, speaker_ids = %s
            WHERE id = %s
        """
        # psycopg2 會自動將 Python list 轉為 PostgreSQL array
        execute_values(
            cur,
            "UPDATE chunks SET face_ids = data.face_ids, speaker_ids = data.speaker_ids FROM (VALUES %s) AS data(face_ids, speaker_ids, id) WHERE chunks.id = data.id",
            updates,
            template=None,
        )
        # 注意：execute_values 對於非簡單 INSERT 有時語法複雜，這裡改用循環或小批次以確保穩定
        # 或者使用簡單的 executemany
        for f_ids, s_ids, cid in updates:
            cur.execute(
                "UPDATE chunks SET face_ids = %s, speaker_ids = %s WHERE id = %s",
                (f_ids, s_ids, cid),
            )

        conn.commit()
        print(f"  Successfully updated {len(updates)} chunks.")
    else:
        print("  No matches found to update.")

    cur.close()


def main():
    parser = argparse.ArgumentParser(description="Sync Face/Speaker IDs to Chunks")
    parser.add_argument("--uuid", help="Specific video UUID to sync")
    parser.add_argument("--all", action="store_true", help="Sync all videos in DB")
    parser.add_argument(
        "--output-dir", default="./output", help="Path to JSON output directory"
    )

    args = parser.parse_args()

    conn = get_db_connection()
    try:
        cur = conn.cursor()

        if args.uuid:
            sync_video(conn, args.uuid, args.output_dir)
        elif args.all:
            cur.execute("SELECT DISTINCT uuid FROM chunks")
            uuids = [row[0] for row in cur.fetchall()]
            print(f"Found {len(uuids)} videos to sync.")
            for uuid in uuids:
                sync_video(conn, uuid, args.output_dir)
        else:
            print("Please specify --uuid <UUID> or --all")

        cur.close()
    finally:
        conn.close()


if __name__ == "__main__":
    main()
