# -------------------- SAFETY PATCHES --------------------
import os
import asyncio

# ✅ Redirect config folders to writable /tmp locations
os.environ["STREAMLIT_HOME"] = "/tmp/.streamlit"
os.environ["STREAMLIT_CONFIG_DIR"] = "/tmp/.streamlit"
os.environ["YOLO_CONFIG_DIR"] = "/tmp/ultralytics"
os.environ["MPLCONFIGDIR"] = "/tmp/mpl"
os.environ["STREAMLIT_WATCH_FILE_SYSTEM"] = "false"  # Optional: disable file watcher

# ✅ Create these directories if they don't exist
os.makedirs("/tmp/.streamlit", exist_ok=True)
os.makedirs("/tmp/ultralytics", exist_ok=True)
os.makedirs("/tmp/mpl", exist_ok=True)

# ✅ Fix for "no running event loop" in asyncio
try:
    asyncio.get_running_loop()
except RuntimeError:
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

# -------------------- STREAMLIT APP --------------------
import streamlit as st
import cv2
import tempfile
import numpy as np
import time
from ultralytics import YOLO
from deep_sort_realtime.deepsort_tracker import DeepSort
from insightface.app import FaceAnalysis

st.set_page_config(layout="wide")
st.title("🎥 Real-Time Gender Detection & Person Tracking")

# ✅ Load models once with cache
@st.cache_resource
def load_models():
    yolo_model = YOLO("yolov8n.pt")
    tracker = DeepSort(max_age=30, max_cosine_distance=0.4, nn_budget=None)
    face_app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
    face_app.prepare(ctx_id=0)
    return yolo_model, tracker, face_app

video_file = st.file_uploader("Upload a video", type=["mp4", "avi", "mov"])

if video_file:
    # ✅ Save to temp file
    temp_video = tempfile.NamedTemporaryFile(delete=False)
    temp_video.write(video_file.read())

    yolo_model, tracker, face_app = load_models()

    cap = cv2.VideoCapture(temp_video.name)
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    gender_map = {}
    total_counts = {"Male": 0, "Female": 0}
    gender_colors = {"Male": (255, 255, 0), "Female": (255, 0, 255)}

    play = st.checkbox("▶️ Play video", value=True)
    FRAME_SKIP = 1

    frame_window = st.empty()

    frame_id = 0
    while cap.isOpened() and play:
        ret, frame = cap.read()
        if not ret:
            break

        frame_id += 1
        if frame_id % FRAME_SKIP != 0:
            continue

        results = yolo_model(frame)[0]
        detections = []
        person_boxes = {}

        for box in results.boxes:
            cls_id = int(box.cls[0].item())
            label = yolo_model.names[cls_id]
            conf = float(box.conf[0].item())
            x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
            if label == "person":
                detections.append(([x1, y1, x2, y2], conf, label))
                person_boxes[(x1, y1, x2, y2)] = None

        tracks = tracker.update_tracks(detections, frame=frame)

        # Gender detection
        for box in person_boxes:
            x1, y1, x2, y2 = box
            face_crop = frame[y1:y2, x1:x2]
            if face_crop.size == 0:
                continue
            faces = face_app.get(face_crop)
            for face in faces:
                gender = "Male" if face.sex == 1 else "Female"
                person_boxes[box] = gender
                total_counts[gender] += 1

        # Draw detections
        male_count, female_count = 0, 0
        for track in tracks:
            if not track.is_confirmed():
                continue
            track_id = track.track_id
            label = track.get_det_class()
            x1, y1, x2, y2 = map(int, track.to_ltrb())
            gender = None
            for box, g in person_boxes.items():
                bx1, by1, bx2, by2 = box
                if abs(x1 - bx1) < 20 and abs(y1 - by1) < 20:
                    gender = g
                    break
            color = gender_colors.get(gender, (200, 200, 200))
            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
            label_str = f"ID:{track_id}"
            if gender:
                label_str += f" ({gender})"
                if gender == "Male":
                    male_count += 1
                else:
                    female_count += 1
            cv2.putText(frame, label_str, (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

        # Overlay count
        cv2.rectangle(frame, (10, 10), (250, 80), (0, 0, 0), -1)
        cv2.putText(frame, f"👨 Male: {male_count}", (20, 35),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, gender_colors["Male"], 2)
        cv2.putText(frame, f"👩 Female: {female_count}", (20, 65),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, gender_colors["Female"], 2)

        frame_window.image(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), channels="RGB")
        time.sleep(1 / fps)

    cap.release()
    st.success("✅ Video playback finished.")