import os
import shutil
import spaces
import gradio as gr

from handlers import frame_handler_yolo as fh
from handlers import video_handler as vh

model_path = "yolov8n.pt"  # YOLOv8 model path


@spaces.GPU(duration=300)
def process_video(video_file):
    """
    Processes the uploaded video file by extracting key frames, cropping them, and generating a processed video.
    """
    status_message = "Processing started..."

    # Define output directories
    output_folder = "output_data"
    all_frames_folder = os.path.join(output_folder, "all_frames")
    key_frames_folder = os.path.join(output_folder, "key_frames")
    nonkey_frames_folder = os.path.join(output_folder, "nonkey_frames")
    cropped_frames_folder = os.path.join(output_folder, "cropped_frames")
    processed_video_path = os.path.join(output_folder, "processed_video.mp4")

    print("Calling process_video function: Output folder:", output_folder)

    # Clear output directory before processing
    if os.path.exists(output_folder):
        shutil.rmtree(output_folder)
    os.makedirs(output_folder, exist_ok=True)

    # Save uploaded video temporarily
    video_path = os.path.join(output_folder, "input_video.mp4")
    with open(video_file.name, "rb") as vf:
        with open(video_path, "wb") as f:
            f.write(vf.read())

    status_message = "Extracting frames..."
    yield status_message, None

    # Step 1: Extract all frames
    vh.extract_all_frames(video_path, all_frames_folder)

    status_message = "Detecting key frames..."
    yield status_message, None

    # Step 2: Extract key frames
    original_fps = 30
    fh.extract_key_frames(all_frames_folder, key_frames_folder, original_fps, model_path)

    status_message = "Cropping key frames..."
    yield status_message, None

    # Step 3: Crop key frames based on object detection
    target_resolution = (360, 640)  # Output resolution (9:16)
    fh.crop_preserve_key_objects(key_frames_folder, cropped_frames_folder, model_path, target_resolution)

    status_message = "Generating final video..."
    yield status_message, None

    # Step 4: Generate short video
    target_resolution = (360, 640)  # Output resolution (9:16)
    target_frame_rate = 30
    vh.create_video_from_frames(cropped_frames_folder, processed_video_path, target_frame_rate, target_resolution)

    status_message = "Processing complete!"
    yield status_message, processed_video_path


# Gradio Blocks UI
with gr.Blocks() as demo:
    gr.Markdown("## Generate short video for your football match")
    gr.Markdown("Upload a video file. The app will extract key frames, crop them to fix 9:16 aspect ratio, "
                "and generate a short video.")

    with gr.Row():
        with gr.Column():
            video_input = gr.File(label="Upload Video", type="filepath", file_types=["video"], file_count="single")
        with gr.Column():
            process_button = gr.Button("Process Video", variant="primary")
            status_output = gr.Textbox(label="Status", interactive=False)
        with gr.Column():
            video_output = gr.Video(label="Processed Video", width=360, height=640)

    process_button.click(process_video, inputs=video_input, outputs=[status_output, video_output])

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