import gradio as gr
import cv2
import pandas as pd
import numpy as np
from ultralytics import YOLO
import tempfile
import time
import os

# Fix Ultralytics path (HF Spaces)
os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"

# 🔥 Replace with best.pt for real pothole detection
model = YOLO("yolov8n.pt")

detections_log = []

def process_video(video, conf_threshold):
    global detections_log
    detections_log = []

    if video is None:
        return None, None, "Upload video first", pd.DataFrame()

    try:
        video_path = video["name"] if isinstance(video, dict) else str(video)

        if not os.path.exists(video_path):
            return None, None, "File not found", pd.DataFrame()

        cap = cv2.VideoCapture(video_path)

        if not cap.isOpened():
            return None, None, "Cannot open video", pd.DataFrame()

        # 🎥 FIX FPS (avoid fast video)
        fps = cap.get(cv2.CAP_PROP_FPS)
        if fps <= 0 or fps > 60:
            fps = 20

        output_fps = max(10, fps * 0.6)

        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) or 640
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) or 480

        output_path = os.path.join(tempfile.gettempdir(), f"out_{int(time.time())}.mp4")

        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, output_fps, (width, height))

        if not out.isOpened():
            return None, None, "VideoWriter failed", pd.DataFrame()

        pothole_count = 0
        frame_count = 0

        # ⚡ Frame skipping for smoother output
        frame_skip = 2
        frame_id = 0

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

            frame_id += 1

            if frame_id % frame_skip != 0:
                continue

            frame = cv2.resize(frame, (width, height))

            # 🔧 Smooth frame
            frame = cv2.GaussianBlur(frame, (3, 3), 0)

            results = model(frame)[0]

            for box in results.boxes:
                conf = float(box.conf[0])

                # 🎯 Better filtering
                if conf < max(conf_threshold, 0.4):
                    continue

                x1, y1, x2, y2 = map(int, box.xyxy[0])

                # 🔴 RED BOX
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 4)

                label = f"Pothole {conf:.2f}"

                # Label background
                (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
                cv2.rectangle(frame, (x1, y1 - h - 10), (x1 + w, y1), (0, 0, 255), -1)

                # White text
                cv2.putText(frame, label,
                            (x1, y1 - 5),
                            cv2.FONT_HERSHEY_SIMPLEX,
                            0.6,
                            (255, 255, 255),
                            2)

                pothole_count += 1

                detections_log.append({
                    "Frame": frame_count,
                    "Confidence": round(conf, 2),
                    "Time": time.strftime("%H:%M:%S")
                })

            out.write(frame)
            frame_count += 1

        cap.release()
        out.release()
        time.sleep(1)

        if not os.path.exists(output_path) or os.path.getsize(output_path) < 1000:
            return None, None, "Output video failed", pd.DataFrame()

        df = pd.DataFrame(detections_log)

        msg = f"Detections: {pothole_count} | Frames: {frame_count}"

        return video_path, output_path, msg, df

    except Exception as e:
        return None, None, str(e), pd.DataFrame()


# ---------------- UI ---------------- #
with gr.Blocks() as demo:

    gr.Markdown("# 🚧 Smart Pothole Detection (Improved)")

    with gr.Row():

        # LEFT PANEL
        with gr.Column(scale=1):
            video_input = gr.Video(label="Upload Video")
            conf = gr.Slider(0, 1, value=0.4, step=0.05, label="Confidence")
            btn = gr.Button("Run Detection")

        # RIGHT PANEL (BIG VIDEOS)
        with gr.Column(scale=3):
            with gr.Row():
                before = gr.Video(label="Original")
                after = gr.Video(label="Detected")

            result = gr.Textbox(label="Summary")
            table = gr.Dataframe(label="Detection Log")

    btn.click(
        process_video,
        inputs=[video_input, conf],
        outputs=[before, after, result, table]
    )


# HF Spaces launch
if __name__ == "__main__":
    demo.launch()