import gradio as gr
import os
import cv2
import tempfile
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import spaces

# CONFIGURATION
REPO_ID = "RohanSardar/yolo11s_drone_tracker"
FILENAME = "best.pt"

# Download model from HuggingFace Hub
print(f"Downloading model from {REPO_ID}...")
try:
    model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
    model = YOLO(model_path)
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")
    print("Ensure REPO_ID is set correctly above.")
    model = None

try:
    bytetrack_path = hf_hub_download(repo_id=REPO_ID, filename="bytetrack.yaml")
    print("Custom bytetrack.yaml loaded from HF!")
except Exception:
    bytetrack_path = None
    print("No custom bytetrack.yaml in HF repo, using default.")

# INFERENCE LOGIC
@spaces.GPU
def process_video(input_video, conf_threshold):
    if not model:
        return None, "Model not loaded. Please configure REPO_ID in app.py"
    if not input_video:
        return None, "Please upload a video."

    # Setup output path
    output_dir = tempfile.mkdtemp()
    out_path = os.path.join(output_dir, "tracked_output.mp4")

    # Read input video properties
    cap = cv2.VideoCapture(input_video)
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    # We use mp4v; Gradio automatically handles standard mp4 browser playback conversions on the backend
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    writer = cv2.VideoWriter(out_path, fourcc, fps, (w, h))

    # Use custom tracker if available, else built-in
    tracker_config = "bytetrack.yaml"
    if bytetrack_path and os.path.exists(bytetrack_path):
        tracker_config = bytetrack_path
    elif os.path.exists("configs/bytetrack.yaml"):
        tracker_config = "configs/bytetrack.yaml"
    elif os.path.exists("bytetrack.yaml"):
        tracker_config = "bytetrack.yaml"

    try:
        # Run YOLO ByteTrack
        results = model.track(
            source=input_video, 
            tracker=tracker_config,
            conf=conf_threshold,
            stream=True, 
            verbose=False
        )

        for r in results:
            # r.plot() returns the BGR image array with boxes drawn
            annotated_frame = r.plot()
            writer.write(annotated_frame)
            
    except Exception as e:
        cap.release()
        writer.release()
        return None, f"Error processing video: {str(e)}"
        
    cap.release()
    writer.release()

    # Convert to standard H264 for guaranteed browser compatibility on HF Spaces
    final_output = os.path.join(output_dir, "final_h264.mp4")
    exit_code = os.system(f"ffmpeg -y -i {out_path} -vcodec libx264 -acodec aac {final_output} 2> NUL")

    if exit_code == 0 and os.path.exists(final_output):
        return final_output, final_output
    
    print("Note: FFmpeg not found or failed. Returning uncompressed mp4v file.")
    return out_path, out_path

# GRADIO INTERFACE
with gr.Blocks() as app:
    gr.Markdown(
        """
        # Drone Tracker 
        Upload a video to detect and track small drones. 
        Powered by fine-tuned YOLOv11 and ByteTrack.
        """
    )
    
    with gr.Row():
        with gr.Column():
            video_input = gr.Video(label="Upload Video")
            conf_slider = gr.Slider(minimum=0.1, maximum=0.9, value=0.25, step=0.05, label="Confidence Threshold")
            submit_btn = gr.Button("Start Tracking", variant="primary")
            
        with gr.Column():
            video_output = gr.Video(label="Tracked Output", interactive=False)
            download_file = gr.File(label="Download File")
            error_text = gr.Markdown()

    submit_btn.click(
        fn=process_video,
        inputs=[video_input, conf_slider],
        outputs=[video_output, download_file]
    )

if __name__ == "__main__":
    app.launch(theme=gr.themes.Soft())