import gradio as gr
import cv2
import tempfile
import os
from ultralytics import YOLO

model = YOLO("yolov8x.pt")

vehicle_classes = {
    "bicycle": 1,
    "car": 2,
    "motorcycle": 3,
    "bus": 5,
    "truck": 7
}

def detect_vehicles_with_logs(video_input):
    cap = cv2.VideoCapture(video_input)
    if not cap.isOpened():
        return None, "", "❌ Error: Could not open video."

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

    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmpfile:
        out_path = tmpfile.name

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out_vid = cv2.VideoWriter(out_path, fourcc, fps, (width, height))

    logs = ""
    frame_index = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_index += 1

        results = model(frame, verbose=False)[0]
        detected = []

        for box in results.boxes:
            cls_id = int(box.cls.item())
            if cls_id in vehicle_classes.values():
                xyxy = box.xyxy[0].cpu().numpy().astype(int)
                conf = box.conf.item()
                label = model.names[cls_id]
                detected.append(f"{label} ({conf:.2f})")

                cv2.rectangle(frame, tuple(xyxy[:2]), tuple(xyxy[2:]), (0, 255, 0), 2)
                cv2.putText(frame, f"{label} {conf:.2f}", (xyxy[0], xyxy[1] - 5),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

        out_vid.write(frame)
        logs += f"Frame {frame_index}: {', '.join(detected) if detected else 'No vehicles detected'}\n"

    cap.release()
    out_vid.release()

    if not os.path.exists(out_path):
        return None, "", "❌ Error: Processed video not found."

    # Return three outputs:
    # 1) path for video player to display (as video file)
    # 2) path for download button (same video file)
    # 3) logs text
    return out_path, out_path, logs


demo = gr.Interface(
    fn=detect_vehicles_with_logs,
    inputs=gr.Video(label="Upload a video"),
    outputs=[
        gr.Video(label="Processed Video"),  # show inline video
        gr.File(label="Download Processed Video"),  # download button
        gr.Textbox(label="Detection Logs", lines=20, interactive=False)
    ],
    title="YOLOv8x Vehicle Detection with Inline Video and Download",
    description="Upload a video, detect vehicles, see the processed video inline, download it, and view logs."
)

if __name__ == "__main__":
    demo.launch(share=True)