#!/opt/homebrew/bin/python3.11
"""
Extract face embedding from an image using InsightFace + CoreML FaceNet.

Usage:
    python3 scripts/extract_face_embedding.py <image_path>

Output: JSON with "embedding" key (512 floats) or "error" key.
Exit code: 0 on success, 1 on failure.
"""
import json
import os
import sys

# Prefer venv if it exists (has insightface + coremltools installed)
VENV_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "venv")
VENV_SITE = os.path.join(VENV_PATH, "lib", "python3.11", "site-packages")
if os.path.isdir(VENV_SITE):
    sys.path.insert(0, VENV_SITE)

import cv2
import numpy as np

MODELS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "models")
FACENET_PATH = os.path.join(MODELS_DIR, "facenet512.mlpackage")


def extract_embedding(image_path: str):
    import io
    import warnings
    warnings.filterwarnings("ignore")

    # Suppress InsightFace verbose stdout during model loading
    old_stdout = sys.stdout
    sys.stdout = io.StringIO()
    try:
        import insightface
        from insightface.app import FaceAnalysis
        import coremltools as ct

        app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
        app.prepare(ctx_id=0, det_thresh=0.5)
        coreml_model = ct.models.MLModel(FACENET_PATH)
    finally:
        sys.stdout = old_stdout

    img_bytes = open(image_path, "rb").read()
    nparr = np.frombuffer(img_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    if img is None:
        print(json.dumps({"error": "Failed to decode image"}))
        sys.exit(1)

    # Detect faces
    faces = app.get(img)
    if not faces:
        print(json.dumps({"error": "No face detected"}))
        sys.exit(1)

    largest = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
    x1, y1, x2, y2 = [int(v) for v in largest.bbox]
    x1, y1 = max(0, x1), max(0, y1)
    x2, y2 = min(img.shape[1], x2), min(img.shape[0], y2)
    if x2 <= x1 or y2 <= y1:
        print(json.dumps({"error": "Invalid face bbox"}))
        sys.exit(1)

    face_img = img[y1:y2, x1:x2]
    face_img = cv2.resize(face_img, (160, 160))
    normalized = (face_img.astype(np.float32) / 127.5) - 1.0
    normalized = np.transpose(normalized, (2, 0, 1))
    input_array = np.expand_dims(normalized, axis=0)

    result = coreml_model.predict({"input": input_array})
    emb_key = [k for k in result.keys() if k.startswith("var_")][0]
    embedding = result[emb_key].flatten().tolist()
    print(json.dumps({"embedding": embedding}))


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(json.dumps({"error": "Usage: extract_face_embedding.py <image_path>"}))
        sys.exit(1)
    extract_embedding(sys.argv[1])