import os
import gradio as gr
import cv2
import shutil
from services.under_construction.combined_detection import run_combined_detection

TEMP_DIR = "temp_frames"
VIDEO_PATH = "uploaded_video.mp4"

def extract_frames(video_path, output_dir, interval=30):
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir, exist_ok=True)

    cap = cv2.VideoCapture(video_path)
    frame_count = 0
    saved = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        if frame_count % interval == 0:
            frame_path = os.path.join(output_dir, f"frame_{saved:05d}.jpg")
            cv2.imwrite(frame_path, frame)
            saved += 1
        frame_count += 1
    cap.release()
    return output_dir

def process_video(video_file):
    # Save uploaded video
    with open(VIDEO_PATH, "wb") as f:
        f.write(video_file.read())

    # Extract frames
    frame_dir = extract_frames(VIDEO_PATH, TEMP_DIR, interval=30)

    # Run combined detection
    results, annotated = run_combined_detection(frame_dir)

    return results, annotated

demo = gr.Interface(
    fn=process_video,
    inputs=gr.File(label="Upload Drone Video (.mp4)", file_types=[".mp4"]),
    outputs=[
        gr.Textbox(label="Detection Results"),
        gr.Gallery(label="Detected Frames").style(grid=3)
    ],
    title="NHAI Combined Detection: Earthwork + Culvert + Bridge Pier",
    description="Upload a drone video. This app detects earthwork, culverts, and bridge piers using YOLOv8 models.",
)

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