import os
from flask import Flask, request, jsonify
import face_recognition
import pickle
from PIL import Image
import numpy as np

app = Flask(__name__)

UPLOAD_FOLDER = 'student_roll/uploads/'
FACE_DATA_FOLDER = 'student_roll/face_data/'

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(FACE_DATA_FOLDER, exist_ok=True)

def convert_image_to_8bit_rgb(image_path):
    with Image.open(image_path) as img:
        if img.mode != 'RGB':
            img = img.convert('RGB')
        img = img.convert('RGB')  # Ensure conversion to 8-bit per channel
        img_array = np.array(img)
    return img_array

@app.route('/upload', methods=['POST'])
def upload_image():
    if 'image' not in request.files:
        return jsonify({"error": "No image part in the request"}), 400
    
    file = request.files['image']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400
    
    student_roll = request.form.get('student_roll')
    if not student_roll:
        return jsonify({"error": "Student roll number is required"}), 400
    
    file_path = os.path.join(UPLOAD_FOLDER, f"{student_roll}.jpg")
    file.save(file_path)
    
    try:
        # Convert image to 8-bit RGB if necessary
        image = convert_image_to_8bit_rgb(file_path)
        face_encodings = face_recognition.face_encodings(image)
        
        if not face_encodings:
            return jsonify({"error": "No face found in the image"}), 400
        
        face_data_path = os.path.join(FACE_DATA_FOLDER, f"{student_roll}.pkl")
        with open(face_data_path, 'wb') as f:
            pickle.dump(face_encodings[0], f)
        
        return jsonify({"message": "Image and face data uploaded successfully"}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/recognize_student', methods=['POST'])
def recognize_student():
    if 'image' not in request.files:
        return jsonify({"error": "No image part in the request"}), 400
    
    file = request.files['image']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400
    
    file_path = os.path.join(UPLOAD_FOLDER, "temp.jpg")
    file.save(file_path)
    
    try:
        # Convert image to 8-bit RGB if necessary
        image = convert_image_to_8bit_rgb(file_path)
        face_encodings = face_recognition.face_encodings(image)
        
        if not face_encodings:
            return jsonify({"error": "No face found in the image"}), 400
        
        uploaded_face_encoding = face_encodings[0]
        
        # Load all stored face data
        for face_data_file in os.listdir(FACE_DATA_FOLDER):
            with open(os.path.join(FACE_DATA_FOLDER, face_data_file), 'rb') as f:
                known_face_encoding = pickle.load(f)
                matches = face_recognition.compare_faces([known_face_encoding], uploaded_face_encoding)
                
                if matches[0]:
                    student_roll = os.path.splitext(face_data_file)[0]
                    return jsonify({"student_roll": student_roll}), 200
        
        return jsonify({"error": "No matching student found"}), 404
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True, host="0.0.0.0", port=5005)