from flask import Flask, request, jsonify, render_template, send_from_directory, make_response
from flask_restful import Api, Resource
import os
import face_recognition
from werkzeug.utils import secure_filename
from PIL import Image
import numpy as np

app = Flask(__name__)
api = Api(app)

UPLOAD_FOLDER = 'upload'
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# Helper function to save image
def save_image(image, roll_number):
    roll_number_folder = os.path.join(app.config['UPLOAD_FOLDER'], roll_number)
    if not os.path.exists(roll_number_folder):
        os.makedirs(roll_number_folder)
    image_count = len(os.listdir(roll_number_folder)) + 1
    image_path = os.path.join(roll_number_folder, f"{image_count}.jpg")
    image.save(image_path)
    return image_path

# Helper function to load known faces
def load_known_faces():
    known_faces = []
    known_roll_numbers = []
    for roll_number in os.listdir(app.config['UPLOAD_FOLDER']):
        roll_number_folder = os.path.join(app.config['UPLOAD_FOLDER'], roll_number)
        for filename in os.listdir(roll_number_folder):
            image_path = os.path.join(roll_number_folder, filename)
            image = face_recognition.load_image_file(image_path)
            encodings = face_recognition.face_encodings(image)
            if encodings:
                known_faces.append(encodings[0])
                known_roll_numbers.append(roll_number)
    return known_faces, known_roll_numbers

class UploadImage(Resource):
    def post(self):
        roll_number = request.form['roll_number']
        if 'image' not in request.files:
            return jsonify({"error": "No image provided"}), 400
        image = request.files['image']
        filename = secure_filename(image.filename)
        image = Image.open(image)
        image_path = save_image(image, roll_number)
        return jsonify({"message": f"Image saved as {image_path}"}), 200

class RecognizeStudent(Resource):
    def post(self):
        if 'image' not in request.files:
            return jsonify({"error": "No image provided"}), 400
        image = request.files['image']
        filename = secure_filename(image.filename)
        image = face_recognition.load_image_file(image)
        unknown_encodings = face_recognition.face_encodings(image)
        if not unknown_encodings:
            return jsonify({"error": "No faces found in the image"}), 400
        unknown_encoding = unknown_encodings[0]
        
        known_faces, known_roll_numbers = load_known_faces()
        results = face_recognition.compare_faces(known_faces, unknown_encoding)
        
        if True in results:
            matched_index = results.index(True)
            roll_number = known_roll_numbers[matched_index]
            return jsonify({"roll_number": roll_number}), 200
        else:
            return jsonify({"error": "No matching student found"}), 404

api.add_resource(UploadImage, '/upload')
api.add_resource(RecognizeStudent, '/recognize')

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