import insightface
from insightface.app import FaceAnalysis

import torch
import torch.nn.functional as F

from PIL import Image
import cv2
import os
import streamlit as st
from glob import glob
import pandas as pd
import numpy as np
from datetime import datetime

# Constants
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
IMAGE_SHAPE = 640
data_path = 'employees'
webcam_path = 'captured_image.jpg'
attendance_db = 'attendance.csv'

st.title("👁️‍🗨️  Face Recognition Based Attendance System")

@st.cache_resource
def load_face_app():
    app = FaceAnalysis(name="buffalo_l")
    app.prepare(ctx_id=-1, det_size=(IMAGE_SHAPE, IMAGE_SHAPE))
    return app

app = load_face_app()

# ✅ Function to mark attendance
def mark_attendance(name):
    import pytz
    ist = pytz.timezone('Asia/Kolkata')
    now = datetime.now(ist)
    date = now.strftime("%Y-%m-%d")
    time = now.strftime("%H:%M:%S")

    if not os.path.exists(attendance_db):
        df = pd.DataFrame(columns=["Name", "Date", "Time"])
        df.to_csv(attendance_db, index=False)

    try:
        df = pd.read_csv(attendance_db)
        if df.empty or not all(col in df.columns for col in ["Name", "Date", "Time"]):
            df = pd.DataFrame(columns=["Name", "Date", "Time"])
            df.to_csv(attendance_db, index=False)
    except pd.errors.EmptyDataError:
        df = pd.DataFrame(columns=["Name", "Date", "Time"])
        df.to_csv(attendance_db, index=False)

    if not ((df['Name'] == name) & (df['Date'] == date)).any():
        new_entry = pd.DataFrame([[name, date, time]], columns=["Name", "Date", "Time"])
        df = pd.concat([df, new_entry], ignore_index=True)
        df.to_csv(attendance_db, index=False)
        return f"✅ Attendance marked for {name} at {time} on {date} (IST)"
    else:
        return f"ℹ️ Attendance already marked for {name} today ({date})"

# 🔍 Face Matching Function
def prod_function(app, prod_path, webcam_path):
    webcam_img = Image.open(webcam_path)
    np_webcam = np.array(webcam_img)
    cv2_webcam = cv2.cvtColor(np_webcam, cv2.COLOR_RGB2BGR)

    webcam_emb = app.get(cv2_webcam, max_num=1)
    if not webcam_emb:
        return None
    webcam_emb = torch.from_numpy(webcam_emb[0].embedding)

    similarity_score = []
    for path in prod_path:
        img = cv2.imread(path)
        face_embedding = app.get(img, max_num=1)
        if not face_embedding:
            similarity_score.append(torch.tensor(-1.0))
            continue
        face_embedding = torch.from_numpy(face_embedding[0].embedding)
        similarity_score.append(F.cosine_similarity(face_embedding, webcam_emb, dim=0))

    return torch.stack(similarity_score)

# 📷 MARK ATTENDANCE TAB
def mark_attendance_tab():
    enable = st.checkbox("Enable camera")
    picture = st.camera_input("Take a picture", disabled=not enable)

    if picture is not None:
        with open(webcam_path, "wb") as f:
            f.write(picture.getbuffer())

        image_paths = glob(os.path.join(data_path, "*.jpg"))

        with st.spinner("Matching face..."):
            prediction = prod_function(app, image_paths, webcam_path)

            if prediction is None or len(prediction) == 0:
                st.error("❌ No face detected in the captured image.")
            else:
                match_idx = torch.argmax(prediction)

                if prediction[match_idx] >= 0.6:
                    matched_name = os.path.basename(image_paths[match_idx]).split('.')[0]
                    display_name = matched_name.replace("_", " ").title()
                    st.success(f"✅ Welcome: {display_name}")

                    # Use session state to prevent duplicate marking when toggling checkbox
                    if "attendance_marked" not in st.session_state or st.session_state.get("last_name") != display_name:
                        attendance_message = mark_attendance(display_name)
                        st.info(attendance_message)
                        st.session_state.attendance_marked = True
                        st.session_state.last_name = display_name
                    else:
                        st.info("ℹ️ Attendance already checked for this session.")


                    # 👇 Show similarity scores if user wants
                    if st.checkbox("Show similarity scores"):
                        st.write("🔍 Similarity scores:", prediction)
                else:
                    st.warning("⚠️ Match not found")


# 📑 ATTENDANCE HISTORY TAB
def attendance_history_tab():
    if os.path.exists(attendance_db):
        df = pd.read_csv(attendance_db)
        if not df.empty:
             st.dataframe(df)
        else:
            st.info("No attendance records found.")
    else:
            st.info("No attendance database found.")

# 🚀 TABS
tabs = st.tabs(["Mark Attendance", "Attendance History"])

with tabs[0]:
    mark_attendance_tab()

with tabs[1]:
    attendance_history_tab()