#!/opt/homebrew/bin/python3.11
"""
Update person demographics (Age, Gender) using InsightFace.
This script scans the representative face of each person and updates the DB.
"""

import os
import cv2
import json
import psycopg2
import numpy as np
import insightface
from insightface.app import FaceAnalysis

# Configuration
DB_CONFIG = {"host": "localhost", "user": "accusys", "dbname": "momentry"}
VIDEO_PATH = "output/384b0ff44aaaa1f1/384b0ff44aaaa1f1.mp4"


def get_face_app():
    app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
    app.prepare(ctx_id=0, det_size=(640, 640))
    return app


def get_person_frames(conn):
    """Get one frame timestamp for each person."""
    cur = conn.cursor()
    # Get the first appearance time for each person to save processing time
    cur.execute("""
        SELECT person_id, MIN(start_time) as start_time 
        FROM dev.person_appearances 
        GROUP BY person_id
    """)
    return cur.fetchall()


def get_frame_at_time(video_path, time_sec):
    """Extract a single frame from video."""
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None

    # InsightFace expects BGR (OpenCV default), but we need to seek correctly
    # cv2 CAP_PROP_POS_MSEC is sometimes inaccurate, better to use FPS if known
    # But for a "representative" frame, seeking by ms is usually "good enough"
    cap.set(cv2.CAP_PROP_POS_MSEC, time_sec * 1000)
    ret, frame = cap.read()
    cap.release()

    if ret:
        return frame
    return None


def analyze_face(app, frame):
    """Run InsightFace to get age/gender."""
    faces = app.get(frame)
    if faces:
        # Just take the first face found
        face = faces[0]
        age = int(face.age) if hasattr(face, "age") else None
        gender = (
            "female"
            if (hasattr(face, "gender") and face.gender == 0)
            else ("male" if (hasattr(face, "gender") and face.gender == 1) else None)
        )
        return age, gender
    return None, None


def update_person_db(conn, person_id, age, gender):
    """Update DB."""
    cur = conn.cursor()
    cur.execute(
        """
        UPDATE dev.person_identities 
        SET age = %s, gender = %s 
        WHERE person_id = %s
    """,
        (age, gender, person_id),
    )
    conn.commit()


def main():
    print("=== Person Demographics Updater ===")

    # 1. Init DB
    conn = psycopg2.connect(**DB_CONFIG)

    # 2. Init Model
    print("Loading InsightFace model...")
    app = get_face_app()

    # 3. Get List
    persons = get_person_frames(conn)
    print(f"Found {len(persons)} persons to process.")

    for i, (person_id, start_time) in enumerate(persons):
        print(
            f"Processing {i + 1}/{len(persons)}: {person_id} (Time: {start_time:.1f}s)"
        )

        # 4. Get Frame
        frame = get_frame_at_time(VIDEO_PATH, start_time)

        if frame is not None:
            # 5. Analyze
            age, gender = analyze_face(app, frame)

            if age and gender:
                print(f"  -> Detected: Age {age}, Gender {gender}")
                # 6. Update
                update_person_db(conn, person_id, age, gender)
            else:
                print(f"  -> Face not found or attributes missing in this frame.")
        else:
            print(f"  -> Failed to retrieve frame.")

    print("=== Done ===")
    conn.close()


if __name__ == "__main__":
    main()