#!/opt/homebrew/bin/python3.11
"""
Store Traced Faces - Pipeline integration for face trace + position data

Flow:
1. Reads face.json output from face_processor.py
2. Runs face_tracker.py to assign trace_id per face (IoU + embedding)
3. Inserts traced faces into face_detections table with trace_id and position (x,y,w,h)

Usage:
    python store_traced_faces.py --file-uuid <uuid> [--face-json <path>]

TKG Export:
    trace_id + position (x,y,w,h) per frame enables spatial-temporal graph construction.
    Each trace is a temporal entity; position tracks movement across frames.
"""

import sys
import os
import json
import argparse
import numpy as np
import psycopg2
import psycopg2.extras
from datetime import datetime

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

# Config
DB_URL = os.environ.get("DATABASE_URL", "postgresql://accusys@localhost:5432/momentry")
SCHEMA = os.environ.get("MOMENTRY_DB_SCHEMA", "dev")
OUTPUT_DIR = os.environ.get("MOMENTRY_OUTPUT_DIR", "/Users/accusys/momentry/output_dev")


def get_conn():
    return psycopg2.connect(DB_URL)


def merge_traces_within_cuts(face_data: dict, cut_scenes: list) -> dict:
    """Merge traces within the same cut if they have similar embeddings (same person re-appeared)."""
    frames = face_data.get("frames", {})
    if not frames:
        return face_data
    
    # Map each frame to its scene/cut number
    frame_to_scene = {}
    for s in cut_scenes:
        for f in range(s["start_frame"], s["end_frame"] + 1):
            frame_to_scene[f] = s["scene_number"]
    
    # Collect per-trace data: scene numbers, embeddings, face positions
    trace_frames = defaultdict(list)
    trace_embeddings = defaultdict(list)
    trace_poses = {}
    
    for fnum_str, frm_data in frames.items():
        fnum = int(fnum_str)
        for face in frm_data.get("faces", []):
            tid = face.get("trace_id")
            if tid is None:
                continue
            trace_frames[tid].append(fnum)
            emb = face.get("embedding")
            if emb is not None:
                trace_embeddings[tid].append(emb)
            if tid not in trace_poses:
                trace_poses[tid] = (face.get("x", 0), face.get("y", 0),
                                    face.get("width", 0), face.get("height", 0))
    
    if len(trace_embeddings) < 2:
        return face_data
    
    # Compute centroid per trace
    trace_centroids = {}
    for tid, embs in trace_embeddings.items():
        centroid = np.mean(embs, axis=0)
        norm = np.linalg.norm(centroid)
        trace_centroids[tid] = centroid / norm if norm > 0 else centroid
    
    # Determine which scene each trace belongs to (majority of frames)
    trace_scene = {}
    for tid, fns in trace_frames.items():
        scene_votes = defaultdict(int)
        for fn in fns:
            scene = frame_to_scene.get(fn, -1)
            scene_votes[scene] += 1
        trace_scene[tid] = max(scene_votes, key=scene_votes.get) if scene_votes else -1
    
    # Within each scene, merge traces with similar centroids
    scene_traces = defaultdict(list)
    for tid, scene in trace_scene.items():
        if scene >= 0 and tid in trace_centroids:
            scene_traces[scene].append(tid)
    
    merged = 0
    next_new_id = max(trace_frames.keys()) + 1 if trace_frames else 0
    SIMILARITY_THRESHOLD = 0.75
    
    for scene, tids in scene_traces.items():
        if len(tids) < 2:
            continue
        used = set()
        for i in range(len(tids)):
            if tids[i] in used:
                continue
            keep_tid = tids[i]
            for j in range(i + 1, len(tids)):
                if tids[j] in used:
                    continue
                sim = float(np.dot(trace_centroids[tids[i]], trace_centroids[tids[j]]))
                if sim >= SIMILARITY_THRESHOLD:
                    # Merge tids[j] into keep_tid
                    for fnum_str, frm_data in frames.items():
                        for face in frm_data.get("faces", []):
                            if face.get("trace_id") == tids[j]:
                                face["trace_id"] = keep_tid
                    used.add(tids[j])
                    merged += 1
    
    # If any merges happened, rebuild trace metadata
    if merged > 0:
        # Rebuild traces dict
        new_traces = {}
        new_trace_frames = defaultdict(list)
        for fnum_str, frm_data in frames.items():
            fnum = int(fnum_str)
            for face in frm_data.get("faces", []):
                tid = face.get("trace_id")
                if tid is not None:
                    new_trace_frames[tid].append({
                        "frame": fnum,
                        "face_index": 0,
                        "bbox": {"x": face.get("x", 0), "y": face.get("y", 0),
                                 "width": face.get("width", 0), "height": face.get("height", 0)},
                        "confidence": face.get("confidence", 0.0),
                    })
        
        for tid, path in new_trace_frames.items():
            if len(path) >= 1:
                frames_sorted = sorted(set(p["frame"] for p in path))
                new_traces[str(tid)] = {
                    "trace_id": tid,
                    "start_frame": frames_sorted[0],
                    "end_frame": frames_sorted[-1],
                    "duration_frames": frames_sorted[-1] - frames_sorted[0] + 1,
                    "duration_seconds": (frames_sorted[-1] - frames_sorted[0]) / face_data.get("metadata", {}).get("fps", 25.0),
                    "total_appearances": len(path),
                    "path": path,
                }
        
        face_data["traces"] = new_traces
        face_data["metadata"]["trace_stats"] = {
            "total_traces": len(new_traces),
            "active_traces": len(new_traces),
            "long_traces": len([t for t in new_traces.values() if t["duration_frames"] >= 2]),
        }
        print(f"[TRACE] Post-merge: {merged} traces merged, {len(new_traces)} total traces")
    
    return face_data


def run_face_tracker(face_json_path: str, traced_json_path: str) -> str:
    """Run face_tracker.py on face.json, returns path to face_traced.json"""
    from face_tracker import track_faces

    with open(face_json_path) as f:
        face_data = json.load(f)

    # V2.0 uses list format (FaceResult), convert to dict for face_tracker
    if isinstance(face_data.get("frames"), list):
        frames_dict = {}
        for frame in face_data["frames"]:
            fnum = str(frame["frame"])
            faces = []
            for f in frame.get("faces", []):
                bbox = f.get("bbox", f)
                face = {
                    "x": bbox.get("x", f.get("x", 0)),
                    "y": bbox.get("y", f.get("y", 0)),
                    "width": bbox.get("width", f.get("width", 0)),
                    "height": bbox.get("height", f.get("height", 0)),
                    "confidence": f.get("confidence", 0.0),
                }
                if "landmarks" in f:
                    face["landmarks"] = f["landmarks"]
                if "embedding" in f:
                    face["embedding"] = f["embedding"]
                faces.append(face)
            frames_dict[fnum] = {
                "frame_number": frame["frame"],
                "time_seconds": frame.get("timestamp", 0),
                "faces": faces,
            }
        face_data["frames"] = frames_dict
        # Preserve metadata (fps needed by face_tracker)
        if "metadata" not in face_data:
            face_data["metadata"] = {
                "fps": face_data.get("fps", 30.0),
                "total_frames": face_data.get("frame_count", 0),
            }

    print(f"[TRACE] Processing {len(face_data.get('frames', {}))} frames")
    
    # Load embeddings from DB for the face tracker
    file_uuid = face_json_path.split("/")[-1].replace(".face.json", "").replace("_traced.json", "")
    try:
        conn = get_conn()
        cur = conn.cursor()
        cur.execute(f"""
            SELECT frame_number, x, y, width, height, embedding
            FROM {SCHEMA}.face_detections
            WHERE file_uuid = %s AND embedding IS NOT NULL
        """, (file_uuid,))
        emb_rows = cur.fetchall()
        conn.close()
        # Build lookup: frame_number → list of (bbox, embedding)
        emb_map = {}
        for fn, x, y, w, h, emb in emb_rows:
            emb_map.setdefault(fn, []).append(((x, y, w, h), emb))
        print(f"[TRACE] Loaded {len(emb_rows)} embeddings from DB")
        
        # Attach embeddings to face data
        attached = 0
        for fnum_str, frm_data in face_data.get("frames", {}).items():
            fnum = int(fnum_str)
            for face in frm_data.get("faces", []):
                x, y, w, h = face.get("x", 0), face.get("y", 0), face.get("width", 0), face.get("height", 0)
                candidates = emb_map.get(fnum, [])
                # Find matching embedding by bbox proximity
                for (ex, ey, ew, eh), emb in candidates:
                    if abs(x - ex) < 10 and abs(y - ey) < 10 and abs(w - ew) < 10 and abs(h - eh) < 10:
                        face["embedding"] = emb
                        attached += 1
                        break
        print(f"[TRACE] Attached {attached} embeddings to faces")
    except Exception as e:
        print(f"[TRACE] WARNING: Could not load embeddings: {e}")
    
    # Load cut boundaries from cut.json (same directory as face.json)
    cut_boundaries = None
    cut_scenes = None
    cuts_path = face_json_path.replace("_traced.json", ".cut.json").replace(".face.json", ".cut.json")
    if os.path.exists(cuts_path):
        with open(cuts_path) as f:
            cuts = json.load(f)
        cut_scenes = cuts.get("scenes", [])
        cut_boundaries = {s["start_frame"] for s in cut_scenes if s["start_frame"] > 0}
        print(f"[TRACE] Loaded {len(cut_boundaries)} cut boundaries")
    
    face_data = track_faces(face_data, use_embedding=True, cut_boundaries=cut_boundaries)
    
    # Merge traces within same cut (same person re-appearing after occlusion/pose change)
    if cut_scenes and len(cut_scenes) > 0:
        face_data = merge_traces_within_cuts(face_data, cut_scenes)
    
    metadata = face_data.get("metadata", {})
    metadata["tracking_method"] = "iou_embedding"
    metadata["tracked_at"] = datetime.now().isoformat()
    face_data["metadata"] = metadata

    with open(traced_json_path, "w") as f:
        json.dump(face_data, f, indent=2, ensure_ascii=False)

    trace_count = len(face_data.get("traces", {}))
    print(f"[TRACE] Completed: {trace_count} traces -> {traced_json_path}")
    return traced_json_path


def store_traced_faces(file_uuid: str, traced_json_path: str, schema: str = SCHEMA):
    """Insert traced face detections into face_detections table with trace_id"""
    conn = get_conn()
    cur = conn.cursor()

    with open(traced_json_path) as f:
        data = json.load(f)

    frames = data.get("frames", {})
    total_stored = 0

    for frame_num_str, frame_data in sorted(frames.items(), key=lambda x: int(x[0])):
        frame_num = int(frame_num_str)
        faces = frame_data.get("faces", [])

        for face in faces:
            trace_id = face.get("trace_id")
            if trace_id is None:
                continue

            x = face.get("x", 0)
            y = face.get("y", 0)
            w = face.get("width", 0)
            h = face.get("height", 0)
            confidence = face.get("confidence", 0.0)
            face_id = face.get("face_id")
            attributes = face.get("attributes")
            embedding = face.get("embedding")

            bbox = json.dumps({"x": x, "y": y, "width": w, "height": h})
            embed_vec = embedding if embedding and len(embedding) > 0 else None

            try:
                cur.execute(
                    f"""
                    UPDATE {schema}.face_detections
                    SET trace_id = %s
                    WHERE file_uuid = %s AND frame_number = %s
                      AND x = %s AND y = %s AND width = %s AND height = %s
                    """,
                    (
                        trace_id,
                        file_uuid, frame_num, x, y, w, h,
                    ),
                )
                if cur.rowcount > 0:
                    total_stored += 1
            except Exception as e:
                print(f"[TRACE] Error storing face at frame {frame_num}: {e}")
                conn.rollback()
                continue

    conn.commit()

    # Log trace summary
    cur.execute(
        f"SELECT COUNT(DISTINCT trace_id) FROM {schema}.face_detections WHERE file_uuid = %s AND trace_id IS NOT NULL",
        (file_uuid,),
    )
    db_trace_count = cur.fetchone()[0]

    cur.close()
    conn.close()

    print(f"[TRACE] Stored {total_stored} face detections, {db_trace_count} unique traces in DB")
    return total_stored, db_trace_count


def main():
    parser = argparse.ArgumentParser(description="Store traced faces in DB")
    parser.add_argument("--file-uuid", required=True, help="Video file UUID")

    parser.add_argument("--face-json", help="Path to face.json (default: auto-detect)")

    parser.add_argument("--schema", default=SCHEMA, help="DB schema name")

    parser.add_argument("--uuid", help="UUID for Redis tracking (accepted by executor)")
    args = parser.parse_args()

    face_json = args.face_json or os.path.join(
        OUTPUT_DIR, f"{args.file_uuid}.face.json"
    )
    traced_json = os.path.join(OUTPUT_DIR, f"{args.file_uuid}.face_traced.json")

    if not os.path.exists(face_json):
        print(f"[TRACE] face.json not found: {face_json}", file=sys.stderr)
        sys.exit(1)

    # Step 1: Run face tracker
    run_face_tracker(face_json, traced_json)

    # Step 2: Store in DB with trace_id
    total, traces = store_traced_faces(args.file_uuid, traced_json, args.schema)
    print(f"[TRACE] Done: {total} detections, {traces} traces")


if __name__ == "__main__":
    main()