import cv2
import numpy as np
import os
import datetime
import csv
import requests

# Load OpenCV Face Recognition and Haar Cascade
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")

try:
    recognizer = cv2.face.LBPHFaceRecognizer_create(radius=2, neighbors=10, grid_x=8, grid_y=8, threshold=90)
except AttributeError:
    print("Error: OpenCV contrib package is missing.")
    print("Install with: pip install opencv-contrib-python")
    exit()

# File Paths
API_ENDPOINT = "https://community.scrubbed.net:8080/api/checkin"
FACE_DIR = "registered_faces"
LOG_FILE = "attendance_log.csv"
MODEL_FILE = "face_recognizer.yml"
USER_ID_FILE = "last_user_id.txt"

def save_last_user_id(user_id):
    with open(USER_ID_FILE, "w") as file:
        file.write(str(user_id))

def get_last_user_id():
    if os.path.exists(USER_ID_FILE):
        with open(USER_ID_FILE, "r") as file:
            return file.read().strip()
    return None

def preprocess_face(face_img):
    """Applies preprocessing to improve recognition accuracy."""
    face_img = cv2.equalizeHist(face_img)  # Normalize brightness
    face_img = cv2.GaussianBlur(face_img, (3, 3), 0)  # Reduce noise
    return face_img

def collect_training_data(user_id):
    """Captures face images for training."""
    if not user_id.isdigit():
        print("Error: User ID must be a number.")
        return

    user_folder = os.path.join(FACE_DIR, user_id)
    os.makedirs(user_folder, exist_ok=True)

    cap = cv2.VideoCapture(0)
    count = 0

    while count < 30:  # Capture 30 images per user for better accuracy
        ret, frame = cap.read()
        if not ret:
            break

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(50, 50))

        for (x, y, w, h) in faces:
            face_img = gray[y:y + h, x:x + w]
            face_img = preprocess_face(face_img)
            resized_face = cv2.resize(face_img, (100, 100))

            # Save original and flipped image to increase training data
            cv2.imwrite(os.path.join(user_folder, f"{count}.jpg"), resized_face)
            cv2.imwrite(os.path.join(user_folder, f"{count}_flipped.jpg"), cv2.flip(resized_face, 1))
            
            count += 1

        cv2.imshow("Face Registration", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()
    save_last_user_id(user_id)
    print(f"Training data collected for user {user_id}.")

def train_face_recognizer():
    """Trains LBPH face recognizer."""
    faces = []
    labels = []

    for user_id in os.listdir(FACE_DIR):
        if not user_id.isdigit():
            continue  # Skip non-user folders
        user_folder = os.path.join(FACE_DIR, user_id)

        for filename in os.listdir(user_folder):
            img_path = os.path.join(user_folder, filename)
            img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)

            if img is None:
                continue  # Skip unreadable files

            img = preprocess_face(cv2.resize(img, (100, 100)))  # Apply preprocessing
            faces.append(img)
            labels.append(int(user_id))

    if not faces:
        print("No training data found.")
        return

    recognizer.train(faces, np.array(labels))
    recognizer.save(MODEL_FILE)
    print("Training complete.")

def recognize_user():
    """Recognizes a user using the trained model."""
    if not os.path.exists(MODEL_FILE):
        print("Error: Model file not found. Train the model first.")
        return

    recognizer.read(MODEL_FILE)
    cap = cv2.VideoCapture(0)

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(50, 50))

        for (x, y, w, h) in faces:
            face_img = preprocess_face(gray[y:y + h, x:x + w])
            resized_face = cv2.resize(face_img, (100, 100))

            user_id, confidence = recognizer.predict(resized_face)

            if confidence < 90:  # Adjusted threshold
                text = f"EMP #: {user_id} (Confidence: {round(confidence, 2)})"
                color = (0, 255, 0)
                log_attendance(user_id)
            else:
                text = "UNRECOGNIZED FACE (Re-register?)"
                color = (0, 0, 255)

            cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
            cv2.putText(frame, text, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2, cv2.LINE_AA)

        cv2.imshow("Face Recognition", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

def log_attendance(user_id):
    """Logs attendance."""
    now = datetime.datetime.now()
    timestamp = now.strftime("%Y-%m-%d %H:%M:%S")

    # Ensure CSV file exists
    if not os.path.exists(LOG_FILE):
        with open(LOG_FILE, "w", newline="") as file:
            writer = csv.writer(file)
            writer.writerow(["Timestamp", "User ID", "Status"])

    with open(LOG_FILE, "a", newline="") as file:
        writer = csv.writer(file)
        writer.writerow([timestamp, user_id, "Checkin"])
        print(f"✅ Check-in logged for User {user_id}")

    send_post_request(user_id)

def send_post_request(user_id):
    """Sends a POST request to the API."""
    payload = {"user_id": user_id}
    headers = {"Content-Type": "application/json"}

    try:
        response = requests.post(API_ENDPOINT, json=payload, headers=headers)
        if response.status_code == 200:
            print(f"✅ Successfully sent attendance for User {user_id} to API.")
        else:
            print(f"⚠️ Failed to send attendance. Status: {response.status_code}, Response: {response.text}")
    except requests.exceptions.RequestException as e:
        print(f"❌ Error sending attendance: {e}")

if __name__ == "__main__":
    while True:
        choice = input("1: Register User\n2: Train Model\n3: Recognize Face\n4: Exit\nChoose an option: ")
        if choice == "1":
            user_id = input("Enter user ID: ")
            collect_training_data(user_id)
        elif choice == "2":
            train_face_recognizer()
        elif choice == "3":
            recognize_user()
        elif choice == "4":
            break
        else:
            print("Invalid choice. Try again.")