import gradio as gr
import cv2
import tempfile
import os
import smtplib
from email.mime.text import MIMEText
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

# -------- LOAD MODEL FROM HUGGING FACE --------
model_path = hf_hub_download(
    repo_id="Ultralytics/YOLOv8",
    filename="yolov8n.pt"
)

model = YOLO(model_path)

# -------- EMAIL ALERT (SECURE) --------
def send_alert():
    sender = os.getenv("EMAIL_USER")
    password = os.getenv("EMAIL_PASS")
    receiver = os.getenv("EMAIL_TO")

    if not sender or not password or not receiver:
        print("Email credentials not set.")
        return

    msg = MIMEText("⚠️ Alert: Person detected!")
    msg["Subject"] = "YOLO Surveillance Alert"
    msg["From"] = sender
    msg["To"] = receiver

    try:
        server = smtplib.SMTP("smtp.gmail.com", 587)
        server.starttls()
        server.login(sender, password)
        server.sendmail(sender, receiver, msg.as_string())
        server.quit()
    except Exception as e:
        print("Email error:", e)

# -------- IMAGE DETECTION --------
def detect_image(image):
    results = model(image)
    annotated = results[0].plot()

    for box in results[0].boxes:
        cls = int(box.cls[0])
        if model.names[cls] == "person":
            send_alert()
            break

    return annotated

# -------- VIDEO DETECTION + TRACKING --------
def detect_video(video):
    cap = cv2.VideoCapture(video)

    temp_out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    output_path = temp_out.name

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    fps = int(cap.get(cv2.CAP_PROP_FPS)) or 20
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    out = cv2.VideoWriter(output_path, fourcc, fps, (w, h))

    alert_sent = False

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

        results = model.track(frame, persist=True)
        annotated = results[0].plot()

        if not alert_sent:
            for box in results[0].boxes:
                cls = int(box.cls[0])
                if model.names[cls] == "person":
                    send_alert()
                    alert_sent = True
                    break

        out.write(annotated)

    cap.release()
    out.release()

    return output_path

# -------- WEBCAM (SAFE VERSION FOR SPACES) --------
def webcam_detect(image):
    results = model(image)
    annotated = results[0].plot()

    for box in results[0].boxes:
        cls = int(box.cls[0])
        if model.names[cls] == "person":
            send_alert()
            break

    return annotated

# -------- UI --------
with gr.Blocks() as demo:
    gr.Markdown("# 🔐 Smart Surveillance System (YOLO + Alerts)")

    # IMAGE TAB
    with gr.Tab("📷 Image"):
        img_in = gr.Image(type="numpy")
        img_out = gr.Image()
        btn1 = gr.Button("Run Detection")
        btn1.click(detect_image, inputs=img_in, outputs=img_out)

    # VIDEO TAB
    with gr.Tab("🎥 Video"):
        vid_in = gr.Video()
        vid_out = gr.Video()
        btn2 = gr.Button("Run Detection + Tracking")
        btn2.click(detect_video, inputs=vid_in, outputs=vid_out)

    # WEBCAM TAB (Stable Version)
    with gr.Tab("📡 Webcam"):
        cam_in = gr.Image(sources=["webcam"], type="numpy")
        cam_out = gr.Image()
        btn3 = gr.Button("Capture & Detect")
        btn3.click(webcam_detect, inputs=cam_in, outputs=cam_out)

demo.launch()