from flask import Flask, request, jsonify
import os
import face_recognition
import pickle
from werkzeug.utils import secure_filename

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads/'
app.config['STUDENT_DATA_FILE'] = 'students_data.pkl'

if not os.path.exists(app.config['UPLOAD_FOLDER']):
    os.makedirs(app.config['UPLOAD_FOLDER'])

# Load or initialize student data
if os.path.exists(app.config['STUDENT_DATA_FILE']):
    with open(app.config['STUDENT_DATA_FILE'], 'rb') as f:
        students_db = pickle.load(f)
else:
    students_db = {}

def save_student_data():
    with open(app.config['STUDENT_DATA_FILE'], 'wb') as f:
        pickle.dump(students_db, f)

@app.route('/upload', methods=['POST'])
def upload_image():
    if 'image' not in request.files or 'student_id' not in request.form:
        return jsonify({'error': 'Image or student_id not provided'}), 400

    image = request.files['image']
    student_id = request.form['student_id']
    filename = secure_filename(f"{student_id}_{image.filename}")
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    image.save(filepath)

    # Load the image file into a numpy array
    image = face_recognition.load_image_file(filepath)
    # Get the face encoding for the image
    face_encodings = face_recognition.face_encodings(image)

    if len(face_encodings) > 0:
        # Assuming the first face found in the image is the correct one
        students_db[student_id] = face_encodings[0]
        save_student_data()
        return jsonify({'message': 'Image uploaded and student registered successfully', 'student_id': student_id}), 200
    else:
        return jsonify({'error': 'No faces found in the image'}), 400

@app.route('/recognize_image', methods=['POST'])
def recognize_image():
    if 'image' not in request.files:
        return jsonify({'error': 'Image not provided'}), 400

    image = request.files['image']
    # Load the uploaded image file into a numpy array
    unknown_image = face_recognition.load_image_file(image)
    # Get the face encodings for the uploaded image
    unknown_encodings = face_recognition.face_encodings(unknown_image)

    if len(unknown_encodings) > 0:
        unknown_encoding = unknown_encodings[0]
        best_match_id = None
        best_match_score = None

        for student_id, known_encoding in students_db.items():
            # Calculate the distance between the known encoding and the uploaded image encoding
            distance = face_recognition.face_distance([known_encoding], unknown_encoding)[0]
            # Convert distance to accuracy score (the closer the distance, the higher the accuracy)
            accuracy_score = (1 - distance) * 100

            if best_match_score is None or accuracy_score > best_match_score:
                best_match_id = student_id
                best_match_score = accuracy_score

        if best_match_id:
            return jsonify({'student_id': best_match_id, 'accuracy_score': best_match_score}), 200

    return jsonify({'error': 'Student not recognized'}), 404

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