import cv2
import os
import time

# --- Configuration ---
USER_ID = "12345"  # Hardcoded user ID for non-interactive use
IMAGE_COUNT = 20   # Number of images to capture for registration

# Create a folder for storing registered faces if it doesn't exist
if not os.path.exists("registered_faces"):
    os.makedirs("registered_faces")

# Initialize the face detector
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

# Initialize webcam
cap = cv2.VideoCapture(0)

# Check if the camera opened successfully
if not cap.isOpened():
    print("Error: Could not open video stream.")
    exit()

# Function to handle registration
def register_user(user_id):
    # Create folder for the user if not exists
    user_folder = f"registered_faces/{user_id}"
    if not os.path.exists(user_folder):
        os.makedirs(user_folder)

    count = 0
    print(f"Starting registration for user {user_id}. Please look at the camera...")

    start_time = time.time()
    while count < IMAGE_COUNT:
        ret, frame = cap.read()
        if not ret:
            print("Failed to grab frame!")
            break

        # Convert to grayscale for face detection
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # Detect faces in the image
        faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))

        if len(faces) > 0:
            # Assuming one face for registration
            (x, y, w, h) = faces[0]

            # Save the face image as a .jpg file in the user folder
            count += 1
            face_img = frame[y:y + h, x:x + w]
            image_path = f"{user_folder}/face_{count}.jpg"
            cv2.imwrite(image_path, face_img)
            print(f"Saved {image_path}")

        # Add a small delay to avoid capturing images too quickly
        time.sleep(0.2)

        # Timeout after 30 seconds
        if time.time() - start_time > 30:
            print("Registration timed out.")
            break

    if count >= IMAGE_COUNT:
        print(f"Registration complete for user {user_id}.")
    else:
        print(f"Registration incomplete. Captured {count}/{IMAGE_COUNT} images.")


# Main execution
if __name__ == "__main__":
    register_user(USER_ID)

# Release the webcam
cap.release()
print("Webcam released.")