import os
import hashlib
from PIL import Image
import gradio as gr

# Directories and files
trainimage_path = "TrainingImage"
hash_store_file = os.path.join(trainimage_path, "image_hashes.txt")
os.makedirs(trainimage_path, exist_ok=True)

# Load stored hashes in the form {enrollment: hash}
def load_hashes():
    hashes = {}
    if os.path.exists(hash_store_file):
        with open(hash_store_file, "r") as f:
            for line in f:
                parts = line.strip().split(",")
                if len(parts) == 2:
                    enrollment, img_hash = parts
                    hashes[enrollment] = img_hash
    return hashes

# Save a new hash entry
def save_hash(enrollment, img_hash):
    with open(hash_store_file, "a") as f:
        f.write(f"{enrollment},{img_hash}\n")

# Create a consistent hash from image
def image_hash(image):
    image = image.resize((100, 100)).convert("L")  # Resize for consistency
    return hashlib.md5(image.tobytes()).hexdigest()

# Main registration logic
def register_student(image, enrollment, name):
    try:
        if not enrollment or not name:
            return "Enrollment ID and Name are required."
        if not enrollment.isdigit():
            return "Enrollment ID must be numeric."
        if not isinstance(image, Image.Image):
            return "Invalid image. Please upload a valid image file."

        # Generate hash of uploaded image
        new_hash = image_hash(image)

        # Load existing image hashes
        existing_hashes = load_hashes()

        # Check if the same image is already registered with the same ID
        if enrollment in existing_hashes:
            existing_hash = existing_hashes[enrollment]
            if new_hash == existing_hash:
                return f"Image already registered with the same image and Enrollment ID {enrollment}."
            else:
                return f"Enrollment ID {enrollment} is already used with a different image!"

        # If new enrollment, save image and record hash
        image_path = os.path.join(trainimage_path, f"User.{enrollment}.jpg")
        image.save(image_path)
        save_hash(enrollment, new_hash)

        return f"Student {name} registered successfully with Enrollment ID {enrollment}!"

    except Exception as e:
        return f"Error registering student: {str(e)}"

# Gradio UI
def create_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## 🧑‍🎓 Student Registration with Duplicate & ID Validation")

        enrollment_input = gr.Textbox(label="Enrollment ID", placeholder="Enter numeric ID")
        name_input = gr.Textbox(label="Name", placeholder="Enter student name")
        image_input = gr.Image(label="Upload Image", type="pil")
        output_text = gr.Textbox(label="Result")

        register_button = gr.Button("Register Student")
        register_button.click(register_student, inputs=[image_input, enrollment_input, name_input], outputs=output_text)

    return demo

demo = create_interface()
demo.launch()