import os
import cv2
import torch
import numpy as np
import streamlit as st
import requests

from PIL import Image
from glob import glob
from insightface.app import FaceAnalysis
import torch.nn.functional as F

# Set the device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Global Variables
IMAGE_SHAPE = 640
data_path = 'employees'
webcam_path = 'captured_image.jpg'

# Set Streamlit title
st.title("📸 AI-Driven Multi-Face Attendance System")

# Load employee image paths
image_paths = glob(os.path.join(data_path, '*.jpg'))
employee_names = [os.path.basename(p).split('.')[0] for p in image_paths]

# Initialize Face Analysis
app = FaceAnalysis(name="buffalo_l")  # ArcFace model
app.prepare(ctx_id=0 if torch.cuda.is_available() else -1, det_size=(IMAGE_SHAPE, IMAGE_SHAPE))

# Load all employee embeddings
employee_embeddings = []
for path in image_paths:
    img = cv2.imread(path)
    faces = app.get(img, max_num=1)
    if faces:
        embedding = torch.tensor(faces[0].embedding, dtype=torch.float32)
        employee_embeddings.append(embedding)
    else:
        employee_embeddings.append(None)

# Function to compute cosine similarity and find best match
def find_best_match(face_emb):
    similarities = []
    for emp_emb in employee_embeddings:
        if emp_emb is None:
            similarities.append(torch.tensor(-1.0))
            continue
        score = F.cosine_similarity(emp_emb, face_emb, dim=0)
        similarities.append(score)
    scores = torch.stack(similarities)
    best_idx = torch.argmax(scores)
    return scores[best_idx].item(), best_idx

# Streamlit tabs
about_tab, app_tab = st.tabs(["About the app", "Face Recognition"])

with about_tab:
    st.markdown("""
    # 🎯 Multi-Face AI Attendance System
    - Detects multiple faces in one frame
    - Matches using ArcFace embeddings
    - Marks attendance for each recognized student
    """)

    st.markdown("### 📚 Students Trained in the System:")
    if employee_names:
        for name in employee_names:
            st.markdown(f"- {name}")
    else:
        st.warning("No student images found in the 'employees/' folder.")

with app_tab:
    enable = st.checkbox("Enable camera")
    picture = st.camera_input("Take a picture", disabled=not enable)

    if picture is not None:
        with st.spinner("Analyzing faces..."):
            image_pil = Image.open(picture)
            image_np = np.array(image_pil)
            image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)

            detected_faces = app.get(image_bgr)

            if not detected_faces:
                st.warning("No face detected in the captured image.")
            else:
                marked_names = []

                for face in detected_faces:
                    emb = torch.tensor(face.embedding, dtype=torch.float32)
                    score, idx = find_best_match(emb)

                    if score >= 0.6:
                        name = employee_names[idx]

                        # Draw box and name
                        x1, y1, x2, y2 = [int(i) for i in face.bbox]
                        cv2.rectangle(image_bgr, (x1, y1), (x2, y2), (0, 255, 0), 2)
                        cv2.rectangle(image_bgr, (x1, y2 - 35), (x2, y2), (0, 255, 0), cv2.FILLED)
                        cv2.putText(image_bgr, name, (x1 + 6, y2 - 10), cv2.FONT_HERSHEY_COMPLEX, 0.8, (255, 255, 255), 2)

                        # Mark attendance only once
                        if name not in marked_names:
                            marked_names.append(name)
                            # Send attendance via POST
                            url = "https://attendance-system-25.glitch.me/adds"
                            data = {'rno': 15, 'sname': name, 'sclass': 7}
                            try:
                                response = requests.post(url, data=data)
                                if response.status_code == 200:
                                    st.success(f"✅ Attendance marked for {name}")
                                else:
                                    st.warning(f"⚠️ Attendance failed for {name}")
                            except Exception as e:
                                st.error(f"Request failed for {name}: {e}")
                    else:
                        x1, y1, x2, y2 = [int(i) for i in face.bbox]
                        cv2.rectangle(image_bgr, (x1, y1), (x2, y2), (0, 0, 255), 2)
                        cv2.putText(image_bgr, "Unknown", (x1 + 6, y2 - 10), cv2.FONT_HERSHEY_COMPLEX, 0.8, (255, 255, 255), 2)

                # Show final image
                final_image = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
                st.image(final_image, caption="Detected Faces with Names", use_column_width=True)