import streamlit as st
import os
from PIL import Image
from huggingface_hub import HfApi

# Constants
KNOWN_FACES_DIR = "known_faces"  # Directory to save user images
REPO_NAME = "face_and_emotion_detection"
REPO_ID = f"LovnishVerma/{REPO_NAME}"

# Ensure the directory exists
os.makedirs(KNOWN_FACES_DIR, exist_ok=True)

# Retrieve Hugging Face token from environment variable
hf_token = os.getenv("upload")  # Replace with your actual Hugging Face token
if not hf_token:
    st.error("Hugging Face token not found. Please set the environment variable.")
    st.stop()

# Initialize Hugging Face API
api = HfApi()

def save_image_to_hugging_face(image, name, roll_no):
    """ Saves the image locally to the KNOWN_FACES_DIR and uploads it to Hugging Face. """
    # Ensure the directory exists
    if not os.path.exists(KNOWN_FACES_DIR):
        os.makedirs(KNOWN_FACES_DIR)

    # Construct the local file path
    filename = f"{name}_{roll_no}.jpg"
    local_path = os.path.join(KNOWN_FACES_DIR, filename)
    
    try:
        # Convert image to RGB if necessary
        if image.mode != "RGB":
            image = image.convert("RGB")
        
        # Save the image to the known_faces directory
        image.save(local_path)
        
        # Upload the saved file to Hugging Face
        api.upload_file(
            path_or_fileobj=local_path,
            path_in_repo=filename,
            repo_id=REPO_ID,
            repo_type="space",
            token=hf_token,
        )
        st.success(f"Image saved to {KNOWN_FACES_DIR} and uploaded to Hugging Face as {filename}.")
    except Exception as e:
        st.error(f"Error saving or uploading image: {e}")
    
    return local_path

# Streamlit UI
st.title("Student Registration with Hugging Face Image Upload")

# Input fields for student details
name = st.text_input("Enter your name")
roll_no = st.text_input("Enter your roll number")

# Choose input method for the image (webcam or file upload)
capture_mode = st.radio("Choose an option to upload your image", ["Use Webcam", "Upload File"])

# Handle webcam capture or file upload
if capture_mode == "Use Webcam":
    picture = st.camera_input("Take a picture")  # Capture image using webcam
elif capture_mode == "Upload File":
    picture = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])  # Upload image from file system

# Save data and process image on button click
if st.button("Register"):
    if not name or not roll_no:
        st.error("Please fill in both name and roll number.")
    elif not picture:
        st.error("Please upload or capture an image.")
    else:
        try:
            # Open the image based on capture mode
            if capture_mode == "Use Webcam" and picture:
                image = Image.open(picture)
            elif capture_mode == "Upload File" and picture:
                image = Image.open(picture)
            
            # Save the image locally and upload it to Hugging Face
            image_path = save_image_to_hugging_face(image, name, roll_no)
        except Exception as e:
            st.error(f"An error occurred: {e}")