import os
import cv2
import pandas as pd
import gradio as gr
from ultralytics import YOLO
from datetime import datetime, timedelta
import numpy as np

class AttendanceSystem:
    def __init__(self, model_path, csv_path):
        self.model_path = model_path
        self.csv_path = csv_path
        self.model = None
        self.class_names_list = []
        self.time_list = []
        self.load_previous_attendance()

    def load_model(self):
        """Load the YOLO model"""
        if self.model is None:
            try:
                self.model = YOLO(self.model_path)
                return True, "Model loaded successfully!"
            except Exception as e:
                return False, f"Error loading model: {str(e)}"
        return True, "Model already loaded"

    def load_previous_attendance(self):
        """Load previous attendance data if CSV exists"""
        if os.path.exists(self.csv_path):
            try:
                df = pd.read_csv(self.csv_path)
                if not df.empty:
                    self.class_names_list = df["Class Name"].tolist()
                    self.time_list = df["Time"].tolist()
                    return True, f"Loaded {len(self.class_names_list)} previous attendance records"
            except Exception as e:
                return False, f"Error loading previous attendance: {str(e)}"
        return False, "No previous attendance data found"

    def process_frame(self, frame):
        """Process a single frame and update attendance"""
        if self.model is None:
            success, message = self.load_model()
            if not success:
                return frame, message, [], []

        # Create a copy of the frame to draw on
        display_frame = frame.copy()

        # Store detected names in this frame
        detected_names = []

        # Detect objects
        results = self.model(frame)

        # Check if any results are detected
        if results:
            for result in results:
                boxes = result.boxes
                for box in boxes:
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
                    class_id = int(box.cls[0])
                    confidence = float(box.conf[0])

                    # Get the class name from YOLO class names
                    class_name = self.model.names[class_id]
                    detected_names.append(class_name)

                    # Draw rectangle around detected object
                    cv2.rectangle(display_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(display_frame, f"{class_name}: {confidence:.2f}",
                                (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

                    # Get the current time
                    current_time = datetime.now()
                    current_time_str = current_time.strftime("%Y-%m-%d %H:%M:%S")

                    # Check if the class has been recorded already
                    person_already_recorded = False
                    for idx, name in enumerate(self.class_names_list):
                        if name == class_name:
                            last_recorded_time = datetime.strptime(self.time_list[idx], "%Y-%m-%d %H:%M:%S")
                            # If less than 24 hours have passed since the last recording
                            if (current_time - last_recorded_time) < timedelta(hours=24):
                                person_already_recorded = True
                                break

                    # Record attendance if not already recorded in the last 24 hours
                    if not person_already_recorded:
                        self.class_names_list.append(class_name)
                        self.time_list.append(current_time_str)
                        self.save_attendance()

        # Create attendance list for display
        attendance_data = []
        for name, time_str in zip(self.class_names_list, self.time_list):
            attendance_data.append(f"{name} - {time_str}")

        return display_frame, f"Detected: {', '.join(detected_names) if detected_names else 'None'}", attendance_data, detected_names

    def save_attendance(self):
        """Save attendance data to CSV"""
        if self.class_names_list and self.time_list:
            df = pd.DataFrame({
                "Class Name": self.class_names_list,
                "Time": self.time_list
            })
            try:
                df.to_csv(self.csv_path, index=False)
                return True, "Attendance saved to CSV"
            except Exception as e:
                return False, f"Error saving to CSV: {str(e)}"
        return False, "No attendance data to save"

    def clear_attendance(self):
        """Clear attendance records"""
        self.class_names_list = []
        self.time_list = []
        if os.path.exists(self.csv_path):
            try:
                os.remove(self.csv_path)
                return True, "Attendance records cleared"
            except Exception as e:
                return False, f"Error clearing records: {str(e)}"
        return True, "No records to clear"


# Function for Gradio interface - process images from webcam
def process_webcam_image(image, state):
    if state is None:
        # Default paths - update these to match your Hugging Face deployment
        model_path = "best(attendance).pt"  # Make sure to upload this model to your Space
        csv_path = "attendance_data.csv"
        state = AttendanceSystem(model_path, csv_path)
        # Load model immediately
        state.load_model()

    if image is None:
        return None, "No image received", "", state

    # Process the frame
    processed_frame, message, attendance_data, detected_names = state.process_frame(image)

    # Format attendance as HTML table for better display
    if attendance_data:
        attendance_html = "<table style='width:100%; border-collapse: collapse;'>"
        attendance_html += "<tr><th style='border:1px solid black; padding:8px;'>Name</th><th style='border:1px solid black; padding:8px;'>Time</th></tr>"

        for record in attendance_data:
            name, time_str = record.split(" - ", 1)
            attendance_html += f"<tr><td style='border:1px solid black; padding:8px;'>{name}</td><td style='border:1px solid black; padding:8px;'>{time_str}</td></tr>"

        attendance_html += "</table>"
    else:
        attendance_html = "No attendance records."

    return processed_frame, message, attendance_html, state


def clear_attendance_records(state):
    if state is not None:
        success, message = state.clear_attendance()
        return message, "<table></table>", state
    return "System not initialized", "<table></table>", None


def change_model_path(model_path, csv_path, state):
    if not model_path or not csv_path:
        return "Please provide both paths", state

    state = AttendanceSystem(model_path, csv_path)
    success, message = state.load_model()
    return message, state


# Create Gradio interface that works on Hugging Face
with gr.Blocks(title="Attendance System") as app:
    gr.Markdown("# Automated Attendance System")
    gr.Markdown("This system uses YOLO to detect and record attendance of individuals.")

    with gr.Row():
        with gr.Column(scale=2):
            # Fixed: Use sources instead of source for the webcam component
            input_image = gr.Image(label="Webcam Feed", sources=["webcam"])
            process_button = gr.Button("Process Image")
            output_image = gr.Image(label="Processed Feed")
            status_text = gr.Textbox(label="Status", value="Capture an image and click 'Process Image'")

        with gr.Column(scale=1):
            # Attendance records and controls
            attendance_display = gr.HTML(label="Attendance Records", value="No records yet.")
            clear_button = gr.Button("Clear Attendance Records")

            # Configuration options
            with gr.Accordion("Configuration", open=False):
                model_path_input = gr.Textbox(label="Model Path", value="best(attendance).pt")
                csv_path_input = gr.Textbox(label="CSV Output Path", value="attendance_data.csv")
                update_paths_button = gr.Button("Update Paths")

    # State for storing the attendance system object
    state = gr.State(None)

    # Set up event handlers
    process_button.click(
        process_webcam_image, 
        inputs=[input_image, state],
        outputs=[output_image, status_text, attendance_display, state]
    )

    clear_button.click(
        clear_attendance_records, 
        inputs=[state],
        outputs=[status_text, attendance_display, state]
    )

    update_paths_button.click(
        change_model_path, 
        inputs=[model_path_input, csv_path_input, state],
        outputs=[status_text, state]
    )

# Create a requirements.txt file for Hugging Face
with open('requirements.txt', 'w') as f:
    f.write('opencv-python-headless\npandas\ngradio>=3.32.0\nultralytics')

# Launch the app
if __name__ == "__main__":
    app.launch()  # Remove share=True for Hugging Face deployment