"""Browser-based live demo: face detection + tracking + per-session persistent ID.

Deploy target: Hugging Face Spaces (Gradio SDK, Python 3.10, free CPU Basic).
This is a demo of the pipeline in ../app, not the NUC/production app itself -
no RTSP, no SQLite, nothing saved. Each visitor gets their own tracker and
FAISS gallery for the lifetime of their browser tab (held in gr.State);
nothing is stored server-side or shared between visitors.
"""
import os
from typing import Dict, List, Optional, Tuple

import numpy as np
import gradio as gr
import supervision as sv

from face_engine import FaceEngine
from identity_store import IdentityStore

DET_SIZE = int(os.environ.get("DET_SIZE", "320"))
MODEL_NAME = os.environ.get("MODEL_NAME", "buffalo_s")
ONNX_THREADS = int(os.environ.get("ONNX_THREADS", "2"))
MATCH_THRESHOLD = float(os.environ.get("MATCH_THRESHOLD", "0.38"))

box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()

# One shared FaceEngine process-wide (loading the ONNX models per visitor
# would be far too slow); it holds no per-visitor state, so this is safe.
face_engine = FaceEngine(
    det_size=(DET_SIZE, DET_SIZE),
    ctx_id=-1,
    model_name=MODEL_NAME,
    num_threads=ONNX_THREADS,
)

SessionState = Tuple[sv.ByteTrack, IdentityStore, Dict[int, int]]


def build_detections(faces: list) -> sv.Detections:
    if not faces:
        return sv.Detections.empty()

    xyxy = np.array([face.bbox for face in faces], dtype=np.float32)
    confidence = np.array([face.det_score for face in faces], dtype=np.float32)
    embeddings = np.array([face.normed_embedding for face in faces], dtype=np.float32)
    class_id = np.zeros(len(faces), dtype=int)

    return sv.Detections(
        xyxy=xyxy,
        confidence=confidence,
        class_id=class_id,
        data={"embedding": embeddings},
    )


def process_frame(
    frame_rgb: Optional[np.ndarray], state: Optional[SessionState]
) -> Tuple[Optional[np.ndarray], SessionState]:
    if state is None:
        state = (sv.ByteTrack(), IdentityStore(embedding_dim=512, match_threshold=MATCH_THRESHOLD), {})
    tracker, identity_store, track_to_person = state

    if frame_rgb is None:
        return None, state

    # gradio's numpy Image type is RGB; our pipeline (InsightFace, supervision) expects BGR.
    frame_bgr = frame_rgb[:, :, ::-1]

    faces = face_engine.detect(frame_bgr)
    detections = build_detections(faces)
    detections = tracker.update_with_detections(detections)

    labels: List[str] = []
    embeddings = detections.data.get("embedding") if detections.data else None
    for i in range(len(detections)):
        track_id = int(detections.tracker_id[i])
        if embeddings is not None:
            person_id = identity_store.match_or_register(embeddings[i])
            track_to_person[track_id] = person_id
        person_id = track_to_person.get(track_id, track_id)
        labels.append(f"ID {person_id}")

    annotated_bgr = box_annotator.annotate(scene=frame_bgr.copy(), detections=detections)
    annotated_bgr = label_annotator.annotate(scene=annotated_bgr, detections=detections, labels=labels)

    return annotated_bgr[:, :, ::-1], (tracker, identity_store, track_to_person)


with gr.Blocks(title="Face Tracker Demo") as demo:
    gr.Markdown(
        "# Real-time multi-person face identification & tracking\n"
        "Allow webcam access below. Each detected face gets a box and a "
        "persistent `ID n` for as long as this tab stays open. Nothing is "
        "recorded, stored, or sent anywhere beyond this page - it all resets "
        "when you reload."
    )
    state = gr.State(None)
    with gr.Row():
        webcam = gr.Image(sources=["webcam"], streaming=True, label="Your webcam", mirror_webcam=True)
        output = gr.Image(label="Detected + tracked", streaming=True)
    webcam.stream(fn=process_frame, inputs=[webcam, state], outputs=[output, state])

if __name__ == "__main__":
    demo.queue().launch()