#!/opt/homebrew/bin/python3.11
"""
Face Clustering Processor V3 - Single-stage trace clustering (Stage 2 merge removed)

Strategy:
1. Load face embeddings from Qdrant _faces collection
2. Group by trace_id, compute weighted average embedding per trace
3. Single-stage AgglomerativeClustering (no Stage 2 merge)
4. Assign person_id to all faces in each trace
5. Output face_cluster.json with auto speaker binding

Output format:
  {"status", "file_uuid", "clusters": [{cluster_id, face_count, representative_face}], "frames": [{frame, timestamp, faces: [{face_id, cluster_id, confidence}]}]}

Changes from V2:
- Added argparse for CLI arguments
- Added Redis progress reporting
- Added status + file_uuid to output
- Added auto_bind_speakers() integration

Changes from previous (broken) version:
- REMOVED Stage 2 merge logic (was incorrectly merging different people)
"""

import argparse
import json
import numpy as np
import os
import sys
import psycopg2
from collections import defaultdict
from sklearn.cluster import AgglomerativeClustering

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from redis_publisher import RedisPublisher

UUID = os.getenv("UUID", "quick_preview")
OUTPUT_DIR = os.getenv("MOMENTRY_OUTPUT_DIR", "./output")
OUTPUT_JSON_PATH = os.path.join(OUTPUT_DIR, f"{UUID}.face_clustered.json")
ASRX_JSON_PATH = os.path.join(OUTPUT_DIR, f"{UUID}.asrx.json")
DB_URL = os.getenv("DATABASE_URL", "postgresql://accusys@localhost:5432/momentry")

CLUSTER_THRESHOLD = 0.35


def write_empty_output(status: str, file_uuid: str, output_path: str):
    output_data = {
        "status": status,
        "file_uuid": file_uuid,
        "clusters": [],
        "frames": []
    }
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(output_data, f, indent=2, ensure_ascii=False)
    print(f"Wrote empty output ({status}) to {output_path}")


def load_qdrant_faces(uuid: str, publisher: RedisPublisher) -> list:
    """Load all face points (payload + vector) from Qdrant _faces for a file."""
    publisher.progress("face_cluster", 0, 100, "Loading embeddings from Qdrant")
    try:
        import requests
        qdrant_url = os.environ.get("QDRANT_URL", "http://localhost:6333")
        qdrant_api_key = os.environ.get("QDRANT_API_KEY", "Test3200Test3200Test3200")
        headers = {}
        if qdrant_api_key:
            headers["api-key"] = qdrant_api_key

        all_points = []
        offset = None
        while True:
            body = {
                "limit": 10000,
                "with_payload": True,
                "with_vector": True,
                "filter": {
                    "must": [
                        {"key": "file_uuid", "match": {"value": uuid}}
                    ]
                }
            }
            if offset:
                body["offset"] = offset

            resp = requests.post(
                f"{qdrant_url}/collections/_faces/points/scroll",
                json=body,
                headers=headers,
                timeout=60
            )
            if resp.status_code != 200:
                print(f"Qdrant scroll error: {resp.status_code}")
                break

            data = resp.json()
            batch = data.get("result", {}).get("points", [])
            all_points.extend(batch)
            next_offset = data.get("result", {}).get("next_page_offset")
            if not next_offset:
                break
            offset = next_offset

        print(f"Loaded {len(all_points)} points from Qdrant _faces")
        return all_points
    except Exception as e:
        print(f"Failed to load embeddings from Qdrant: {e}")
        return []


def cluster_by_trace(points: list, min_faces_per_trace: int = 1, confidence_threshold: float = 0.5) -> tuple:
    """
    Cluster faces by trace_id:
    1. Filter low-confidence faces
    2. Group by trace_id
    3. Compute weighted average embedding per trace
    """
    high_conf_points = [
        p for p in points
        if p.get("payload", {}).get("confidence", 0) >= confidence_threshold
    ]
    print(f"[FILTER] {len(high_conf_points)}/{len(points)} faces pass confidence >= {confidence_threshold}")

    traces = defaultdict(list)
    for point in high_conf_points:
        payload = point.get("payload", {})
        trace_id = payload.get("trace_id")
        if trace_id is not None and trace_id >= 0:
            traces[trace_id].append(point)

    print(f"[TRACE] Found {len(traces)} unique traces")

    trace_embeddings = {}
    for trace_id, trace_points in traces.items():
        if len(trace_points) < min_faces_per_trace:
            continue

        embeddings = []
        weights = []
        for p in trace_points:
            embeddings.append(p["vector"])
            weights.append(p.get("payload", {}).get("confidence", 1.0))

        embeddings = np.array(embeddings)
        weights = np.array(weights)
        weights = weights / weights.sum()

        avg_embedding = np.average(embeddings, axis=0, weights=weights)
        avg_embedding = avg_embedding / np.linalg.norm(avg_embedding)

        trace_embeddings[trace_id] = {
            "embedding": avg_embedding,
            "face_count": len(trace_points),
            "avg_confidence": float(np.mean(weights)),
            "frames": [p.get("payload", {}).get("frame") for p in trace_points],
            "points": trace_points
        }

    print(f"[TRACE] {len(trace_embeddings)} traces with >= {min_faces_per_trace} faces")
    return trace_embeddings, traces


def single_stage_clustering(trace_embeddings: dict) -> dict:
    """
    Single-stage AgglomerativeClustering on trace means.
    No Stage 2 merge (the problematic logic has been removed).
    """
    if not trace_embeddings:
        return {}

    trace_ids = list(trace_embeddings.keys())
    n_traces = len(trace_ids)

    if n_traces == 0:
        return {}

    if n_traces == 1:
        return {trace_ids[0]: 0}

    embeddings = np.array([trace_embeddings[tid]["embedding"] for tid in trace_ids])

    print(f"[CLUSTER] Clustering {n_traces} traces (threshold={CLUSTER_THRESHOLD})...")
    clustering = AgglomerativeClustering(
        n_clusters=None,
        distance_threshold=CLUSTER_THRESHOLD,
        metric="cosine",
        linkage="average"
    )
    labels = clustering.fit_predict(embeddings)

    unique_labels = set(labels)
    print(f"[CLUSTER] Detected {len(unique_labels)} unique persons")

    trace_to_person = {}
    for i, trace_id in enumerate(trace_ids):
        trace_to_person[trace_id] = labels[i]

    return trace_to_person


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("cluster_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 or start is None or end is None:
            continue
        candidates = [f for f in face_spans if f.get("ts") is not None and 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}")


def main():
    global OUTPUT_JSON_PATH, UUID

    parser = argparse.ArgumentParser(description="Face Clustering Processor V3")
    parser.add_argument("video_path", help="Path to video file")
    parser.add_argument("output_path", help="Path to output JSON file")
    parser.add_argument("--uuid", help="Video UUID (optional, overrides env var)")
    parser.add_argument("--force", action="store_true", help="Overwrite existing output")
    parser.add_argument("--frames", type=str, default=None, help=argparse.SUPPRESS)
    args = parser.parse_args()

    OUTPUT_JSON_PATH = args.output_path
    output_dir = os.path.dirname(args.output_path)
    if args.uuid:
        uuid = args.uuid
    else:
        uuid = os.path.basename(args.video_path).rsplit(".", 1)[0]
    UUID = uuid

    face_json_path = os.path.join(output_dir, f"{uuid}.face.json")
    if not os.path.exists(face_json_path):
        face_json_path = os.path.join(output_dir, uuid, f"{uuid}.face.json")

    publisher = RedisPublisher(uuid)
    publisher.info("face_cluster", "Face clustering started")

    if not os.path.exists(face_json_path):
        print("Face JSON not found.")
        write_empty_output("no_face_json", uuid, OUTPUT_JSON_PATH)
        publisher.complete("face_cluster", "No face.json found")
        return

    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 (no faces).")
        write_empty_output("no_faces", uuid, OUTPUT_JSON_PATH)
        publisher.complete("face_cluster", "No faces to cluster")
        return

    fps_value = face_data.get("fps", 23.98)

    points = load_qdrant_faces(uuid, publisher)
    if not points:
        print("No embeddings found in Qdrant.")
        write_empty_output("no_embeddings", uuid, OUTPUT_JSON_PATH)
        publisher.complete("face_cluster", "No embeddings in Qdrant")
        return

    publisher.progress("face_cluster", 10, 100, f"Loaded {len(points)} Qdrant points")

    trace_embeddings, traces = cluster_by_trace(
        points,
        min_faces_per_trace=1,
        confidence_threshold=0.5
    )

    if not trace_embeddings:
        print("No valid traces found.")
        write_empty_output("no_embeddings", uuid, OUTPUT_JSON_PATH)
        publisher.complete("face_cluster", "No valid traces")
        return

    publisher.progress("face_cluster", 30, 100, f"Found {len(trace_embeddings)} traces")

    trace_to_person = single_stage_clustering(trace_embeddings)
    if not trace_to_person:
        print("Clustering produced no results.")
        write_empty_output("no_embeddings", uuid, OUTPUT_JSON_PATH)
        publisher.complete("face_cluster", "Clustering empty")
        return

    n_persons = len(set(trace_to_person.values()))
    print(f"Clustering result: {n_persons} persons ({len(trace_to_person)} traces)")
    publisher.progress("face_cluster", 50, 100, f"Found {n_persons} persons")

    qdrant_by_frame_trace = {}
    for point in points:
        payload = point.get("payload", {})
        frame = payload.get("frame")
        trace_id = payload.get("trace_id", -1)
        bbox = payload.get("bbox", {})
        if frame is not None and trace_id >= 0:
            cx = bbox.get("x", 0) + bbox.get("width", 0) // 2
            cy = bbox.get("y", 0) + bbox.get("height", 0) // 2
            key = (frame, cx, cy)
            qdrant_by_frame_trace[key] = trace_id

    matched_count = 0
    for frame_idx, frame_obj in enumerate(frames_list):
        frame_num = frame_obj.get("frame", frame_idx)
        for face in frame_obj.get("faces", []):
            x = face.get("x", 0)
            y = face.get("y", 0)
            w = face.get("width", 0)
            h = face.get("height", 0)
            cx = x + w // 2
            cy = y + h // 2
            key = (frame_num, cx, cy)

            trace_id = qdrant_by_frame_trace.get(key)
            if trace_id is None:
                for (qf, qx, qy), tid in qdrant_by_frame_trace.items():
                    if qf == frame_num and abs(qx - cx) < 50 and abs(qy - cy) < 50:
                        trace_id = tid
                        break

            if trace_id is not None and trace_id in trace_to_person:
                person_label = trace_to_person[trace_id]
                face["person_id"] = f"Person_{person_label}"
                matched_count += 1

    print(f"   Assigned person_id to {matched_count} faces")

    if matched_count == 0:
        print("No faces matched to any trace - check Qdrant data integrity")
        publisher.progress("face_cluster", 100, 100, "No trace matches")
        write_empty_output("no_faces", uuid, OUTPUT_JSON_PATH)
        return

    publisher.progress("face_cluster", 70, 100, f"Matched {matched_count} faces")

    person_face_count = defaultdict(int)
    person_best_face = {}
    for frame_idx, frame_obj in enumerate(frames_list):
        for face in frame_obj.get("faces", []):
            person_id = face.get("person_id")
            if not person_id:
                continue
            person_face_count[person_id] += 1
            confidence = face.get("confidence", 0.9)
            if person_id not in person_best_face or confidence > person_best_face[person_id]["confidence"]:
                person_best_face[person_id] = {
                    "face_id": f"face_{frame_idx}_{frame_idx}",
                    "confidence": confidence,
                    "frame": frame_obj.get("frame", frame_idx),
                    "bbox": {
                        "x": face.get("x"),
                        "y": face.get("y"),
                        "width": face.get("width"),
                        "height": face.get("height")
                    }
                }

    person_labels_sorted = sorted(person_face_count.keys(), key=lambda p: -person_face_count[p])
    clusters = []
    for person_id in person_labels_sorted:
        cluster_entry = {
            "cluster_id": person_id,
            "face_count": person_face_count[person_id],
            "representative_face": person_best_face.get(person_id)
        }
        clusters.append(cluster_entry)

    output_frames = []
    for frame_idx, frame_obj in enumerate(frames_list):
        timestamp = frame_obj.get("timestamp", frame_obj.get("frame", 0) / fps_value if fps_value > 0 else 0)
        output_faces = []
        for face in frame_obj.get("faces", []):
            person_id = face.get("person_id")
            if person_id:
                output_faces.append({
                    "face_id": f"face_{frame_idx}_{frame_idx}",
                    "cluster_id": person_id,
                    "confidence": face.get("confidence", 0.9),
                })
        if output_faces:
            output_frames.append({
                "frame": frame_obj.get("frame", frame_idx),
                "timestamp": timestamp,
                "faces": output_faces,
            })

    output_data = {
        "status": "has_faces",
        "file_uuid": UUID,
        "clusters": clusters,
        "frames": output_frames,
    }
    with open(OUTPUT_JSON_PATH, "w", encoding="utf-8") as f:
        json.dump(output_data, f, indent=2, ensure_ascii=False)
    print(f"Saved clustered data to {OUTPUT_JSON_PATH}")
    publisher.progress("face_cluster", 90, 100, f"{len(clusters)} clusters")

    auto_bind_speakers()

    publisher.complete("face_cluster", f"{len(clusters)} clusters")


if __name__ == "__main__":
    main()