import streamlit as st
import cv2
import os
import numpy as np
from keras.models import load_model
from PIL import Image
import sqlite3
import requests
from io import BytesIO

# Constants
DB_FILE = "students.db"
KNOWN_FACES_DIR = "known_faces"
EMOTION_MODEL_FILE = "CNN_Model_acc_75.h5"
HUGGING_FACE_TOKEN = os.getenv("HUGGING_FACE_TOKEN")
HUGGING_FACE_REPO = "username/repo_name"

# Emotion Labels
EMOTION_LABELS = ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]

# Create directories
os.makedirs(KNOWN_FACES_DIR, exist_ok=True)

# Load models
try:
    emotion_model = load_model(EMOTION_MODEL_FILE)
except Exception as e:
    st.error(f"Error loading emotion model: {e}")

# Database Functions
def create_table():
    with sqlite3.connect(DB_FILE) as conn:
        cursor = conn.cursor()
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS students (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                roll_number TEXT NOT NULL UNIQUE,
                image_path TEXT NOT NULL
            )
        """)
        conn.commit()

def insert_student(name, roll_number, image_path):
    try:
        with sqlite3.connect(DB_FILE) as conn:
            cursor = conn.cursor()
            cursor.execute("INSERT INTO students (name, roll_number, image_path) VALUES (?, ?, ?)",
                           (name, roll_number, image_path))
            conn.commit()
    except sqlite3.IntegrityError:
        st.warning("Roll number already exists!")

# Hugging Face Functions
def upload_to_hugging_face(file_path, file_name):
    if not HUGGING_FACE_TOKEN:
        st.error("Hugging Face token not found.")
        return

    url = f"https://huggingface.co/api/repos/{HUGGING_FACE_REPO}/uploads/{file_name}"
    headers = {"Authorization": f"Bearer {HUGGING_FACE_TOKEN}"}

    try:
        with open(file_path, "rb") as file:
            response = requests.post(url, headers=headers, files={"file": file})
        if response.status_code == 200:
            st.success(f"Uploaded {file_name} to Hugging Face!")
        else:
            st.error(f"Failed to upload: {response.content}")
    except Exception as e:
        st.error(f"Error uploading file: {e}")

# Image Processing Functions
def detect_faces_and_emotions(image):
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
    faces = face_cascade.detectMultiScale(gray_image, scaleFactor=1.3, minNeighbors=5)

    for (x, y, w, h) in faces:
        face = gray_image[y:y+h, x:x+w]
        resized_face = cv2.resize(face, (48, 48))
        normalized_face = resized_face / 255.0
        reshaped_face = np.reshape(normalized_face, (1, 48, 48, 1))
        emotion_prediction = emotion_model.predict(reshaped_face)
        emotion_label = np.argmax(emotion_prediction)
        return EMOTION_LABELS[emotion_label]

    return None

# UI Design
st.title("Student Registration and Emotion Detection")
create_table()

menu = ["Register Student", "Face Recognition and Emotion Detection"]
choice = st.sidebar.selectbox("Menu", menu)

if choice == "Register Student":
    st.subheader("Register a New Student")

    with st.form("register_form"):
        name = st.text_input("Name")
        roll_number = st.text_input("Roll Number")
        image_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
        submitted = st.form_submit_button("Register")

        if submitted:
            if name and roll_number and image_file:
                try:
                    img = Image.open(image_file)
                    img_path = os.path.join(KNOWN_FACES_DIR, f"{roll_number}.png")
                    img.save(img_path)
                    insert_student(name, roll_number, img_path)
                    upload_to_hugging_face(img_path, f"{roll_number}.png")
                    st.success("Student Registered Successfully!")
                except Exception as e:
                    st.error(f"Error: {e}")
            else:
                st.warning("Please fill in all fields and upload an image.")

elif choice == "Face Recognition and Emotion Detection":
    st.subheader("Recognize Faces and Detect Emotions")

    action = st.radio("Choose Action", ["Upload Image", "Use Webcam"])

    if action == "Upload Image":
        uploaded_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
        if uploaded_file:
            try:
                img = Image.open(uploaded_file)
                img_array = np.array(img)
                emotion_label = detect_faces_and_emotions(img_array)
                if emotion_label is not None:
                    st.success(f"Emotion Detected: {emotion_label}")
                else:
                    st.warning("No face detected.")
            except Exception as e:
                st.error(f"Error: {e}")

    elif action == "Use Webcam":
        st.info("Use the camera input widget to capture an image.")
        camera_image = st.camera_input("Take a picture")
        if camera_image:
            try:
                img = Image.open(camera_image)
                img_array = np.array(img)
                emotion_label = detect_faces_and_emotions(img_array)
                if emotion_label is not None:
                    st.success(f"Emotion Detected: {emotion_label}")
                else:
                    st.warning("No face detected.")
            except Exception as e:
                st.error(f"Error: {e}")