import streamlit as st
import cv2
import tempfile
import numpy as np
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

# Force the page to wide layout for better visibility
st.set_page_config(page_title="VirtuGuard AI Security Suite 2026", layout="wide")

# 2. Add your logo
# This renders it in the upper-left corner and sidebar
st.logo("logo.png", size="medium")

# Optional: Add a larger version or specific text to the sidebar
with st.sidebar:
    st.image("logo.png", width=150)
    st.markdown("VirtuGuard AI Security Suite 2026")
    
# --- Load Models ---
@st.cache_resource
def load_models():
    # 1. Advanced YOLO26 for Person (0) and Vehicles (2,3,5,7)
    main = YOLO('yolo26n.pt') 
    
    # 2. Corrected Firearm/Threat model path
    # 'weights/best.pt' is the correct location for this repo
    threat_path = hf_hub_download(
        repo_id="Subh775/Threat-Detection-YOLOv8n", 
        filename="weights/best.pt"
    )
    threat = YOLO(threat_path)
    return main, threat

main_model, threat_model = load_models()

st.title("🛡️ VirtuGuard Advanced AI Security Suite (YOLO26)")
st.info("Detecting People, Vehicles, and Multi-Threats (Guns & Knives)")

uploaded_file = st.file_uploader("Upload Image or Video", type=["mp4", "mov", "jpg", "png"])

if uploaded_file:
    is_video = uploaded_file.type.startswith('video')
    
    if is_video:
        # Create a temp file to store the uploaded video
        tfile = tempfile.NamedTemporaryFile(delete=False) 
        tfile.write(uploaded_file.read())
        cap = cv2.VideoCapture(tfile.name)
        st_frame = st.empty()

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret: break

            # Inference
            # Main model: Person=0, Car=2, Motorcycle=3, Bus=5, Truck=7
            res_main = main_model(frame, classes=[0, 2, 3, 5, 7], verbose=False)[0]
            # Threat model: 1=Gun, 2=Explosive, 3=Grenade, 4=Knife
            res_threat = threat_model(frame, conf=0.5, verbose=False)[0]

            # Combine visualizations
            annotated = res_main.plot()
            annotated = res_threat.plot(img=annotated)

            # Display
            st_frame.image(annotated, channels="BGR")
        cap.release()
    else:
        # Static Image Processing
        img = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)
        r1 = main_model(img, classes=[0, 2, 3, 5, 7])[0]
        r2 = threat_model(img, conf=0.5)[0]
        
        # Merge results onto one image
        final_img = r2.plot(img=r1.plot())
        st.image(final_img, channels="BGR", use_container_width=True)
        
        # Display Summary
        st.subheader("Detections Summary")
        col1, col2 = st.columns(2)
        with col1:
            st.write("**Base Detections:**")
            for box in r1.boxes:
                label = main_model.names[int(box.cls[0])]
                st.write(f"✅ {label.capitalize()}: {box.conf[0]:.2%}")
        with col2:
            st.write("**Security Threats:**")
            for box in r2.boxes:
                # Map threat model class IDs to names
                threat_names = {1: "Gun", 2: "Explosive", 3: "Grenade", 4: "Knife"}
                label = threat_names.get(int(box.cls[0]), "Weapon")
                st.write(f"⚠️ {label.upper()}: {box.conf[0]:.2%}")