import os
import cv2
import streamlit as st
from datetime import datetime
from tempfile import NamedTemporaryFile
from ultralytics import YOLO
import time

# --- Streamlit Page Configuration (MUST BE FIRST STREAMLIT COMMAND) ---
st.set_page_config(layout="wide", page_title="Weapon Detection System")

# --- Initialize YOLO models ---
@st.cache_resource
def load_models():
    try:
        weapon_model = YOLO("best.pt")
        effect_model = YOLO("yolov8n.pt")
        print("YOLO Models loaded successfully.")
        return weapon_model, effect_model
    except Exception as e:
        st.error(f"Error loading YOLO models: {e}")
        print(f"Error loading YOLO models: {e}")
        return None, None

weapon_model, effect_model = load_models()
if weapon_model is None or effect_model is None:
    st.error("CRITICAL ERROR: YOLO models could not be loaded. The application cannot continue. Check model paths.")
    st.stop()

# --- COCO Labels, Output Folder, Sample Videos ---
coco_labels = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
    "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
    "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
    "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
    "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
    "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
    "potted plant", "bed", "dining table", "toilet", "TV", "laptop", "mouse", "remote", "keyboard",
    "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
    "scissors", "teddy bear", "hair drier", "toothbrush"
]
output_folder = "detected_clips"
os.makedirs(output_folder, exist_ok=True)
SAMPLE_VIDEOS = {
    "CriminalThreatens (Example)": "sample/CriminalThreatens.mp4", # Replace with your actual file name
    "Person with a Gun (Example)": "sample/Shutter.mp4",       # Replace with your actual file name
    "Person with Knife (Example)": "sample/WithGuns.mp4",
    "GroupPortrait Of Three People (Example)": "sample/GroupPortrait.mp4",
    "Sample Photos (Example)" : "sample/Samplephotos.mp4"
}

# --- Custom CSS ---
st.markdown("""
<style>
    .log-container .stTextArea textarea { color: #FFA500 !important; font-family: 'Consolas', 'Monaco', monospace !important; }
    .latest-log-info .stAlert { font-size: 0.95rem; }
</style>
""", unsafe_allow_html=True)

# --- Session State Variables ---
if "log" not in st.session_state: st.session_state.log = ""
if "stop_camera" not in st.session_state: st.session_state.stop_camera = False
if "is_processing" not in st.session_state: st.session_state.is_processing = False
if "current_video_source" not in st.session_state: st.session_state.current_video_source = None
if "current_video_path" not in st.session_state: st.session_state.current_video_path = None


# --- MODIFIED Helper Function: write_log ---
def write_log(message, placeholder_to_update=None, is_error=False,
              log_level_for_ui="INFO", # Renamed from log_level_for_db
              _db_video_source=None, _db_clip_path=None): # Prefixed unused DB params with _
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    ui_prefix = "⚡"
    # log_level_for_ui will determine the prefix mostly
    actual_ui_log_level = str(log_level_for_ui).upper()

    if is_error or "error" in message.lower():
        ui_prefix = "🔴"; actual_ui_log_level = "ERROR"
    elif "detected" in message.lower():
        ui_prefix = "⚠️"; actual_ui_log_level = "DETECTION"
    elif "recording started" in message.lower(): ui_prefix = "🔴"
    elif "recording stopped" in message.lower() or "finalized" in message.lower(): ui_prefix = "💾"
    elif "initiated" in message.lower() or "selected" in message.lower() or "uploaded" in message.lower():
        ui_prefix = "🚦"; actual_ui_log_level = "SYSTEM_EVENT"
    elif "stop" in message.lower() and "user" in message.lower():
         ui_prefix = "⏹️"; actual_ui_log_level = "USER_ACTION"
    elif "opened successfully" in message.lower(): ui_prefix = "✅"
    elif "loop finished" in message.lower(): ui_prefix = "🏁"


    new_session_log_entry = f"[{timestamp}] {ui_prefix} {message}"
    st.session_state.log = new_session_log_entry + "\n" + st.session_state.log
    max_log_lines = 100 # Keep UI log concise
    log_lines = st.session_state.log.splitlines()
    if len(log_lines) > max_log_lines: st.session_state.log = "\n".join(log_lines[:max_log_lines])

    if placeholder_to_update:
        if actual_ui_log_level == "ERROR": placeholder_to_update.error(f"🔴 Latest: {message}")
        else: placeholder_to_update.info(f"{ui_prefix} Latest: {message}")

    # --- DATABASE LOGGING PART REMOVED ---
    # final_db_video_source = _db_video_source
    # if final_db_video_source is None and st.session_state.get("current_video_source"):
    #     current_src_state = st.session_state.current_video_source
    #     if current_src_state == "webcam": final_db_video_source = "webcam"
    #     elif current_src_state in ["file", "sample"] and st.session_state.get("current_video_path"):
    #         final_db_video_source = os.path.basename(st.session_state.current_video_path)
    #     elif isinstance(current_src_state, str): final_db_video_source = current_src_state
    #
    # write_log_to_db(level=actual_db_log_level, message=message,
    #                 video_source=final_db_video_source, clip_path=_db_clip_path)
    # --- END OF REMOVED DATABASE LOGGING PART ---

    # Console logging for debugging (optional, can be removed if not needed)
    print(f"LOG: [{timestamp}] {actual_ui_log_level} - {message}")
    if _db_video_source: print(f"  Video Source: {_db_video_source}")
    if _db_clip_path: print(f"  Clip Path: {_db_clip_path}")


# --- UI Definition ---
with st.sidebar:
    st.title("Controls & Log")
    st.markdown("Upload video, use webcam, or select sample.")
    st.markdown("---")
    uploaded_file_sidebar = st.file_uploader("Upload Video", type=["mp4", "avi", "mov"], key="uploader",
                                           disabled=st.session_state.is_processing)
    if st.button("📸 Start Webcam", key="start_cam", disabled=st.session_state.is_processing, use_container_width=True):
        if not st.session_state.is_processing:
            st.session_state.is_processing = True; st.session_state.stop_camera = False
            st.session_state.current_video_source = "webcam"; st.session_state.current_video_path = None
            write_log("Webcam feed initiated...", log_level_for_ui="SYSTEM_EVENT")
            st.rerun()
    st.markdown("---")
    st.subheader("Sample Videos")
    sample_options = ["None"] + list(SAMPLE_VIDEOS.keys())
    selected_sample = st.selectbox("Select sample:", options=sample_options, index=0, key="sample_select",
                                   disabled=st.session_state.is_processing)
    if selected_sample != "None":
        sample_p = SAMPLE_VIDEOS[selected_sample]
        if not os.path.exists(sample_p): st.warning(f"Sample missing: {sample_p}")
        elif st.button(f"Process: {selected_sample}", key=f"proc_samp_{selected_sample.replace(' ','_')}",
                       disabled=st.session_state.is_processing, use_container_width=True):
            if not st.session_state.is_processing:
                st.session_state.is_processing = True; st.session_state.stop_camera = False
                st.session_state.current_video_source = "sample"; st.session_state.current_video_path = sample_p
                write_log(f"Sample '{selected_sample}' selected.", log_level_for_ui="SYSTEM_EVENT")
                st.rerun()
    st.markdown("---")
    st.subheader(" Session Log")
    st.markdown('<div class="log-container">', unsafe_allow_html=True)
    st.text_area("LogView", value=st.session_state.log, height=250, disabled=True, key="session_log_area", label_visibility="collapsed")
    st.markdown('</div>', unsafe_allow_html=True)

# Main Area
st.title("Weapon & Object Detection System ")
st.markdown("Real-time weapon and object detection. Detected weapon clips are saved.")
st.markdown("---")
frame_placeholder = st.empty()
st.markdown('<div class="latest-log-info">', unsafe_allow_html=True)
status_placeholder = st.empty()
st.markdown('</div>', unsafe_allow_html=True)

# --- Core Processing Function ---
def process_video_feed(video_input, status_ph): # video_input is path or 0 for webcam
    cap = None
    ui_source_name = "Webcam"; # db_source_name variable removed or can be reused for console log

    if isinstance(video_input, str):
        if not os.path.exists(video_input):
            msg = f"Video file not found: {video_input}"
            st.error(msg); write_log(msg, status_ph, is_error=True, _db_video_source=video_input)
            st.session_state.is_processing = False; st.session_state.stop_camera = True; return
        cap = cv2.VideoCapture(video_input)
        ui_source_name = os.path.basename(video_input)
        write_log(f"Processing video: {ui_source_name}...", status_ph, _db_video_source=ui_source_name)
    else:
        cap = cv2.VideoCapture(video_input) # 0 for webcam
        write_log("Attempting to start webcam...", status_ph, _db_video_source="webcam")

    if not cap or not cap.isOpened():
        msg = "Error opening video source."
        st.error(msg); write_log(msg, status_ph, is_error=True, _db_video_source=ui_source_name if isinstance(video_input, str) else "webcam")
        st.session_state.is_processing = False; st.session_state.stop_camera = True
        if cap: cap.release(); return

    status_ph.info(f"⏳ Processing {ui_source_name}...")
    write_log("Video source opened.", None, _db_video_source=ui_source_name if isinstance(video_input, str) else "webcam")

    recording = False; out = None; current_clip_path = ""
    fps = cap.get(cv2.CAP_PROP_FPS)
    fw = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)); fh = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    if fps <= 0 or fps > 120: fps = 20.0 # Default FPS

    try:
        while not st.session_state.stop_camera:
            ret, frame = cap.read()
            if not ret:
                write_log("End of video or stream error.", status_ph, _db_video_source=ui_source_name if isinstance(video_input, str) else "webcam")
                break

            weapon_res = weapon_model(frame, verbose=False, half=True)
            effect_res = effect_model(frame, verbose=False, half=True)
            annot_frame = frame.copy()
            weapon_detected_this_frame = False
            detected_weapon_name = "Weapon" # Default, can be model.names[cls]
            now_time = datetime.now()

            for r in weapon_res:
                if r.boxes is not None:
                    for box in r.boxes:
                        cls = int(box.cls.item()); conf = float(box.conf.item())
                        if cls == 0 and conf > 0.5: # Assuming class 0 is weapon
                            weapon_detected_this_frame = True
                            x1,y1,x2,y2 = map(int, box.xyxy[0].tolist())
                            cv2.rectangle(annot_frame, (x1,y1),(x2,y2), (0,0,255),2)
                            cv2.putText(annot_frame, f"{detected_weapon_name} ({conf:.2f})", (x1,y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5,(0,0,255),2)
            for r in effect_res:
                if r.boxes is not None:
                    for box in r.boxes:
                        cls = int(box.cls.item()); conf = float(box.conf.item())
                        if conf > 0.5:
                            obj_name = coco_labels[cls] if cls < len(coco_labels) else f"Obj-{cls}"
                            x1,y1,x2,y2 = map(int, box.xyxy[0].tolist())
                            cv2.rectangle(annot_frame, (x1,y1),(x2,y2), (255,0,0),2)
                            cv2.putText(annot_frame, f"{obj_name} ({conf:.2f})", (x1,y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5,(255,0,0),2)

            if weapon_detected_this_frame:
                if not recording:
                    ts_date = now_time.strftime("%d-%m-%y"); ts_time = now_time.strftime("%H-%M-%S")
                    vid_fname = f"{detected_weapon_name}_{ts_date}_{ts_time}.mp4"
                    current_clip_path = os.path.join(output_folder, vid_fname)
                    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
                    out = cv2.VideoWriter(current_clip_path, fourcc, fps, (fw,fh))
                    if not out.isOpened():
                        msg = f"Error opening VideoWriter: {current_clip_path}"
                        st.error(msg); write_log(msg, status_ph, is_error=True, _db_video_source=ui_source_name if isinstance(video_input, str) else "webcam")
                        st.session_state.stop_camera=True; break
                    recording = True
                    write_log(f"REC Start: {vid_fname}", status_ph, log_level_for_ui="RECORDING_EVENT", _db_clip_path=current_clip_path)
                if recording and out: out.write(frame) # Original frame
                log_ts = now_time.strftime("%I:%M:%S %p")
                write_log(f"{detected_weapon_name} detected @{log_ts}", status_ph, log_level_for_ui="DETECTION", _db_clip_path=current_clip_path if recording else None)
            else: # No weapon
                if recording:
                    recording=False
                    if out: out.release(); out=None
                    write_log(f"REC Stop: {os.path.basename(current_clip_path)}", status_ph, log_level_for_ui="RECORDING_EVENT", _db_clip_path=current_clip_path)
                    current_clip_path = ""
            
            rgb_frame = cv2.cvtColor(annot_frame, cv2.COLOR_BGR2RGB)
            frame_placeholder.image(rgb_frame, channels="RGB", use_container_width=True)
    finally:
        if cap: cap.release()
        if recording and out: # If loop/app stops while recording
            out.release()
            write_log(f"REC Finalized (incomplete?): {os.path.basename(current_clip_path)}", status_ph, log_level_for_ui="RECORDING_EVENT", _db_clip_path=current_clip_path)
        write_log("Processing loop ended.", status_ph, _db_video_source=ui_source_name if isinstance(video_input, str) else "webcam")
        st.session_state.is_processing = False
        st.session_state.stop_camera = True # Ensure stop is true

# --- Main Control Logic ---
if uploaded_file_sidebar and not st.session_state.is_processing:
    if st.session_state.current_video_source != "file_processed_marker":
        st.session_state.is_processing = True; st.session_state.stop_camera = False
        st.session_state.current_video_source = "file"
        with NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_f:
            tmp_f.write(uploaded_file_sidebar.read())
            st.session_state.current_video_path = tmp_f.name
        write_log(f"File '{uploaded_file_sidebar.name}' uploaded.", status_placeholder, log_level_for_ui="SYSTEM_EVENT", _db_video_source=uploaded_file_sidebar.name)
        st.rerun()

if st.session_state.is_processing:
    if st.button("🛑 Stop Processing", key="stop_main_btn", use_container_width=False):
        write_log("User stopped processing.", status_placeholder, log_level_for_ui="USER_ACTION")
        st.session_state.stop_camera = True
        st.session_state.is_processing = False # Crucial: mark as not processing now
        st.session_state.stop_button_pressed_this_run = True # Flag to manage rerun logic
        st.rerun()
else: # Not processing
    if not st.session_state.current_video_source and not uploaded_file_sidebar and selected_sample == "None":
         frame_placeholder.info("✨ Welcome! Upload a video, start webcam, or select a sample. ✨")

if st.session_state.is_processing and not st.session_state.stop_camera:
    frame_placeholder.empty()
    input_src = None
    if st.session_state.current_video_source == "webcam": input_src = 0
    elif st.session_state.current_video_source in ["file", "sample"]: input_src = st.session_state.current_video_path
    
    if input_src is not None:
        process_video_feed(input_src, status_placeholder)
    else:
        write_log("Error: Invalid video source for processing.", status_placeholder, is_error=True)
        st.session_state.is_processing = False; st.session_state.stop_camera = True; st.rerun()

if not st.session_state.is_processing:
    if st.session_state.current_video_source is not None and \
       st.session_state.current_video_source != "file_processed_marker":
        
        if st.session_state.current_video_source == "file" and st.session_state.current_video_path:
            if os.path.exists(st.session_state.current_video_path):
                try:
                    os.unlink(st.session_state.current_video_path)
                    write_log(f"Temp file deleted: {os.path.basename(st.session_state.current_video_path)}", status_placeholder, log_level_for_ui="SYSTEM_EVENT")
                except Exception as e:
                    write_log(f"Error deleting temp file: {e}", status_placeholder, is_error=True)
            st.session_state.current_video_path = None
        
        status_placeholder.success("✅ Processing finished or stopped. Ready for new input.")
        st.session_state.current_video_source = "file_processed_marker"
        
        if not st.session_state.get("stop_button_pressed_this_run", False):
             st.rerun()
    st.session_state.pop("stop_button_pressed_this_run", None)

# Display video at the end (outside of processing logic)
if not st.session_state.is_processing and st.session_state.current_video_source == "file_processed_marker":
    st.video("https://youtu.be/nJ0QmeDjbkE")