import streamlit as st
import cv2
import tempfile
import torch
import numpy as np
from ultralytics import YOLO
from deep_sort_realtime.deepsort_tracker import DeepSort
import warnings
import os
import platform

# Suppress ScriptRunContext warnings from threads
warnings.filterwarnings("ignore", message=".*missing ScriptRunContext.*")

# Check if running in headless environment
IS_HEADLESS = platform.system() == 'Linux'

# Initialize YOLO model
@st.cache_resource
def load_yolo_model():
    try:
        return YOLO("best.pt")
    except Exception as e:
        st.error(f"Error loading YOLO model: {str(e)}")
        return None

# Main app
st.title("📦 Inventory Management")

# Settings
st.sidebar.header("Settings")
CONF_THRESHOLD = st.sidebar.slider(
    "Confidence Threshold",
    min_value=0.0,
    max_value=1.0,
    value=0.4,
    help="Higher values mean more confident detections but might miss objects"
)

FRAME_SKIP = st.sidebar.slider(
    "Frame Skip",
    min_value=0,
    max_value=10,
    value=2,
    help="Process every Nth frame (higher values = faster processing but may miss objects)"
)

# Load YOLO model
yolo_model = load_yolo_model()

if yolo_model is None:
    st.error("Failed to load YOLO model. Please check if the model file exists.")
    st.stop()

# File uploader
uploaded_file = st.sidebar.file_uploader(
    "Upload Video", 
    type=["mp4", "avi", "mov"],
    help="Supported formats: MP4, AVI, MOV"
)

if uploaded_file is not None:
    try:
        tfile = tempfile.NamedTemporaryFile(delete=False)
        tfile.write(uploaded_file.read())
        video_path = tfile.name

        if st.sidebar.button("Start Processing"):
            tracker = DeepSort(
                embedder="mobilenet",
                embedder_gpu=torch.cuda.is_available(),
                max_age=30  # Increase max_age for longer tracking retention
            )

            cap = cv2.VideoCapture(video_path)
            if not cap.isOpened():
                st.error("Error opening video file")
                st.stop()

            # Get video properties for processing
            fps = cap.get(cv2.CAP_PROP_FPS)
            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            
            counted_objects = set()
            frame_placeholder = st.empty()
            status_text = st.sidebar.empty()
            progress_bar = st.progress(0)
            
            # Counter for frame skipping
            frame_counter = 0
            
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                
                frame_counter += 1
                current_position = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
                progress = current_position / frame_count
                progress_bar.progress(progress)
                
                # Skip frames based on user setting
                if FRAME_SKIP > 0 and frame_counter % (FRAME_SKIP + 1) != 0:
                    continue

                try:
                    # Resize frame for faster processing (if needed)
                    # h, w = frame.shape[:2]
                    # if w > 1280:  # Only resize if the frame is large
                    #     frame = cv2.resize(frame, (1280, int(h * 1280 / w)))
                    
                    results = yolo_model(frame, verbose=False)  # Turn off verbose output for speed
                    detections = []
                    for result in results:
                        for box in result.boxes.data.tolist():
                            x1, y1, x2, y2, score, class_id = box
                            if score > CONF_THRESHOLD:
                                detections.append([[x1, y1, x2 - x1, y2 - y1], score, int(class_id)])

                    tracks = tracker.update_tracks(detections, frame=frame)
                    for track in tracks:
                        if not track.is_confirmed():
                            continue
                        track_id = track.track_id
                        ltrb = track.to_ltrb()
                        x1, y1, x2, y2 = map(int, ltrb)
                        counted_objects.add(track_id)
                        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                        cv2.putText(frame, f"ID: {track_id}", (x1, y1-10), 
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

                    cv2.putText(frame, f"Total Objects: {len(counted_objects)}", 
                                (20, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
                    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    frame_placeholder.image(frame_rgb, channels="RGB", use_container_width=True)
                    status_text.info(f"Processing... Current count: {len(counted_objects)}")
                    
                    # Remove sleep to maximize performance
                    # time.sleep(0.01)

                except Exception as e:
                    st.error(f"Error processing frame: {str(e)}")
                    continue

            cap.release()
            progress_bar.progress(1.0)
            st.sidebar.success(f"Final count: {len(counted_objects)} objects")
            st.balloons()

    except Exception as e:
        st.error(f"Error processing video: {str(e)}")
    finally:
        if 'tfile' in locals():
            tfile.close()