import cv2
import os
from datetime import datetime
import time
import argparse
import subprocess
import dlib
from imutils import face_utils
import pickle
from cryptography.fernet import Fernet
from dotenv import load_dotenv


#if "ENCODINGS_KEY" not in os.environ:
#    print("Please set the ENCODINGS_KEY environment variable with your encryption key.")
#    exit(1)

#key = os.environ.get("ENCODINGS_KEY").encode()
load_dotenv()
key = os.getenv("ENCODINGS_KEY")
fernet = Fernet(key)

def encrypt_data(data):
    serialized = pickle.dumps(data)
    encrypted = fernet.encrypt(serialized)
    return encrypted

def decrypt_data(encrypted_data):
    decrypted = fernet.decrypt(encrypted_data)
    data = pickle.loads(decrypted)
    return data

detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")

def create_folder(name):
    dataset_folder = "dataset"
    if not os.path.exists(dataset_folder):
        os.makedirs(dataset_folder)
    person_folder = os.path.join(dataset_folder, name)
    if not os.path.exists(person_folder):
        os.makedirs(person_folder)
    return person_folder

def is_blinking(landmarks):
    left_eye = landmarks[36:42]
    right_eye = landmarks[42:48]
    def eye_aspect_ratio(eye):
        A = cv2.norm(eye[1] - eye[5])
        B = cv2.norm(eye[2] - eye[4])
        C = cv2.norm(eye[0] - eye[3])
        return (A + B) / (2.0 * C)
    left_eye_ratio = eye_aspect_ratio(left_eye)
    right_eye_ratio = eye_aspect_ratio(right_eye)
    ear = (left_eye_ratio + right_eye_ratio) / 2.0
    return ear < 0.5

def capture_photos(name, auto=False):
    folder = create_folder(name)
    cap = cv2.VideoCapture(2)

    if not cap.isOpened():
        print("Error: Could not open camera.")
        return

    time.sleep(2)
    photo_count = 0

    if not auto:
        print(f"Press ENTER to start taking 10 photos for {name}. Press 'q' to quit.")
        input()
    else:
        print(f"Starting photo capture for {name} automatically.")

    while photo_count < 10:
        ret, frame = cap.read()
        if not ret:
            print("Error: Could not read frame.")
            break

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = detector(gray, 0)

        if len(faces) == 0:
            print("No face detected. Please align properly.")
            time.sleep(1)
            continue

        for face in faces:
            landmarks = predictor(gray, face)
            landmarks = face_utils.shape_to_np(landmarks)

            if len(landmarks) == 68 and is_blinking(landmarks):
                photo_count += 1
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"{name}_{timestamp}_{photo_count}.jpg"
                filepath = os.path.join(folder, filename)
                cv2.imwrite(filepath, frame)
                print(f"Photo {photo_count}/10 saved: {filepath}")
            else:
                print("No blink detected. Try again.")

        time.sleep(1)

    cap.release()
    print(f"Photo capture completed. {photo_count} photos saved for {name}.")

    print("[INFO] Running model training script with encryption...")
    venv_python = "./biometric_system/bin/python3.10"
    subprocess.run([venv_python, "model_training.py"])
    print("[INFO] Model training complete. Exiting...")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Capture photos for a person.")
    parser.add_argument('person_name', type=str, help="The name of the person to photograph")
    parser.add_argument('--auto', action="store_true", help="Run automatically without waiting for keyboard input")
    args = parser.parse_args()
    capture_photos(args.person_name, auto=args.auto)