import os
import sqlite3
import streamlit as st
from PIL import Image
import numpy as np
import cv2
from keras.models import load_model
from datetime import datetime
from huggingface_hub import HfApi

# Constants
KNOWN_FACES_DIR = "known_faces"  # Directory to save user images
DATABASE = "students.db"  # SQLite database file to store student information

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

# Initialize Hugging Face API
hf_token = os.getenv("upload")  # Ensure this is set correctly as a secret in Hugging Face
if not hf_token:
    raise ValueError("Hugging Face token not found. Ensure it's set as a secret in Hugging Face")
api = HfApi()

# Repository Details on Hugging Face
REPO_NAME = "face_and_emotion_detection"  # Replace with your Hugging Face repository name
REPO_ID = "LovnishVerma/" + REPO_NAME  # Replace "LovnishVerma" with your Hugging Face username
REPO_TYPE = "space"  # 'space' type for Streamlit-based projects

# Load emotion detection model
model = load_model('CNN_Model_acc_75.h5')
emotion_labels = ['angry', 'fear', 'happy', 'neutral', 'sad', 'surprise']
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

# Initialize the SQLite database
def initialize_database():
    conn = sqlite3.connect(DATABASE)
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS students (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT NOT NULL,
            roll_no TEXT NOT NULL UNIQUE,
            image_path TEXT NOT NULL,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    """)
    conn.commit()
    conn.close()

# Save student information in the SQLite database
def save_to_database(name, roll_no, image_path):
    conn = sqlite3.connect(DATABASE)
    cursor = conn.cursor()
    try:
        cursor.execute("""
            INSERT INTO students (name, roll_no, image_path)
            VALUES (?, ?, ?)
        """, (name, roll_no, image_path))
        conn.commit()
        st.success("Data saved successfully!")
    except sqlite3.IntegrityError:
        st.error("Roll number already exists!")
    finally:
        conn.close()

# Save the captured image locally in known_faces directory and upload to Hugging Face
def save_image_to_hugging_face(image, name, roll_no):
    # Create a filename based on the student name and roll number
    filename = f"{name}_{roll_no}.jpg"
    local_path = os.path.join(KNOWN_FACES_DIR, filename)
    
    # Save the image locally
    image.save(local_path)
    st.success(f"Image saved locally to {local_path}")
    
    try:
        # Upload the image to Hugging Face repository
        api.upload_file(
            path_or_fileobj=local_path,
            path_in_repo=filename,
            repo_id=REPO_ID,
            repo_type=REPO_TYPE,
            token=hf_token
        )
        st.success(f"Image uploaded to Hugging Face: {filename}")
    except Exception as e:
        st.error(f"Error uploading image to Hugging Face: {e}")
    
    return local_path

# Process each frame for emotion detection
def process_frame(frame):
    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray_frame, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))

    for (x, y, w, h) in faces:
        roi_gray = gray_frame[y:y+h, x:x+w]
        roi_color = frame[y:y+h, x:x+w]

        face_roi = cv2.resize(roi_color, (48, 48))
        face_roi = np.expand_dims(face_roi, axis=0)
        face_roi = face_roi / float(48)
        predictions = model.predict(face_roi)
        emotion = emotion_labels[np.argmax(predictions[0])]

        # Display emotion text on face
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
        cv2.putText(frame, emotion, (x, y+h), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

    return frame

# User Interface for registration
st.title("Student Registration and Attendance")

# 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"])

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"])

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

# Handle image upload or webcam capture
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)

            # Save user data to the database
            save_to_database(name, roll_no, image_path)

            # Detect faces and emotions
            cap = cv2.VideoCapture(0)
            while True:
                ret, frame = cap.read()
                if not ret:
                    break

                frame = process_frame(frame)
                st.image(frame, channels="BGR", use_column_width=True)
                break  # Stop after capturing one frame

            cap.release()

        except Exception as e:
            st.error(f"An error occurred: {e}")

# Display registered students and attendance history
if st.checkbox("Show registered students"):
    conn = sqlite3.connect(DATABASE)
    cursor = conn.cursor()
    cursor.execute("SELECT name, roll_no, image_path, timestamp FROM students")
    rows = cursor.fetchall()
    conn.close()

    st.write("### Registered Students")
    for row in rows:
        name, roll_no, image_path, timestamp = row
        st.write(f"**Name:** {name}, **Roll No:** {roll_no}, **Timestamp:** {timestamp}")
        st.image(image_path, caption=f"{name} ({roll_no})", use_column_width=True)