import os
import sys
import argparse
from deepface import DeepFace
import cv2
import json
from db import get_connection_from_env, ensure_table, save_embedding

CASCADE_PATH = "src/face.xml"
EMBED_DIR = "embeddings"


def ensure_dir(d):
    if not os.path.exists(d):
        os.makedirs(d)


def crop_faces(image_path):
    img = cv2.imread(image_path)
    if img is None:
        return []
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    face_cascade = cv2.CascadeClassifier(CASCADE_PATH)
    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
    crops = []
    for i, (x, y, w, h) in enumerate(faces):
        crop = img[y:y+h, x:x+w]
        crops.append(crop)
    return crops


def enroll(name, image_path, save_db=True):
    ensure_dir(EMBED_DIR)
    crops = crop_faces(image_path)
    if len(crops) == 0:
        print("No faces found to enroll")
        return
    # Use first face crop for embedding
    crop = crops[0]
    tmp_path = f"/tmp/{name}_crop.jpg"
    cv2.imwrite(tmp_path, crop)
    print("Computing embedding for", name)
    # DeepFace.represent returns embedding vector or dict depending on params
    embedding = DeepFace.represent(tmp_path, model_name='Facenet', enforce_detection=True)
    # Save embedding to JSON file
    out_file = os.path.join(EMBED_DIR, f"{name}.json")
    with open(out_file, "w") as f:
        json.dump(embedding, f)
    print("Saved embedding to", out_file)

    # Optionally save to DB
    if save_db:
        try:
            conn = get_connection_from_env()
            ensure_table(conn)
            # make embedding JSON-serializable: try to extract vector
            vec = None
            if isinstance(embedding, dict) and 'embedding' in embedding:
                vec = embedding['embedding']
            elif isinstance(embedding, list) and len(embedding) > 0 and isinstance(embedding[0], dict) and 'embedding' in embedding[0]:
                vec = embedding[0]['embedding']
            else:
                try:
                    # try convert to list
                    vec = list(embedding)
                except Exception:
                    vec = None
            if vec is not None:
                save_embedding(conn, name, json.dumps(vec))
                print("Saved embedding to database for", name)
            else:
                print("Could not parse embedding to save to DB")
        except Exception as e:
            print("DB save failed:", e)


def parse_args():
    parser = argparse.ArgumentParser(description="Enroll a face and save embedding")
    parser.add_argument("name", help="Person name")
    parser.add_argument("image", help="Path to image file")
    parser.add_argument(
        "--save-db",
        action="store_true",
        default=True,
        help="Save embedding to MySQL (default: true)",
    )
    parser.add_argument(
        "--no-db",
        action="store_false",
        dest="save_db",
        help="Do not save embedding to MySQL",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    enroll(args.name, args.image, save_db=args.save_db)