import gradio as gr
from spatialmedia import metadata_utils
import os
import shutil

def console_print(message):
    print(message)

def inject_360_metadata(input_video, stereo_mode, spatial_audio, crop):
    if input_video is None:
        return None

    base_dir = "processed_videos"
    os.makedirs(base_dir, exist_ok=True)
    input_video_path = os.path.join(base_dir, "input_video.mp4")
    output_video_path = os.path.join(base_dir, "output_video.mp4")
    
    with open(input_video_path, 'wb') as f:
        f.write(input_video)

    metadata = metadata_utils.Metadata()
    metadata.video = metadata_utils.generate_spherical_xml(stereo_mode, crop)
    
    if spatial_audio:
        parsed_metadata = metadata_utils.parse_metadata(input_video_path, console_print)
        if not metadata.audio:
            spatial_audio_description = metadata_utils.get_spatial_audio_description(
                parsed_metadata.num_audio_channels)
            if spatial_audio_description.is_supported:
                metadata.audio = metadata_utils.get_spatial_audio_metadata(
                    spatial_audio_description.order,
                    spatial_audio_description.has_head_locked_stereo)
            else:
                raise ValueError("Audio has %d channel(s) and isn't a supported spatial audio format." % (parsed_metadata.num_audio_channels))

    if metadata.video:
        metadata_utils.inject_metadata(input_video_path, output_video_path, metadata, console_print)
        return output_video_path
    else:
        raise ValueError("Failed to generate metadata.")

def update_output(output_file_path):
    if output_file_path and os.path.exists(output_file_path):
        return output_file_path
    else:
        return None

def main():
    with gr.Blocks(theme='JohnSmith9982/small_and_pretty') as demo:
        gr.Markdown("This space adds appropriate metadata to 360° equirectangular videos so they can be recognized as such. The audio checkbox option is only viable if your video has spatial audio in ambiX ACN/SN3D with head-locked stereo")
        with gr.Row():
            video_input = gr.File(label="Select video file", type="binary")
            stereo_dropdown = gr.Dropdown(choices=["none", "top-bottom", "left-right"], label="Stereo Mode")
            spatial_audio_checkbox = gr.Checkbox(label="Spatial Audio")
            crop_input = gr.Textbox(label="Crop Region (w:h:f_w:f_h:x:y)- Optional")
            submit_btn = gr.Button("Inject 360° Metadata")
        
        output_file = gr.File(label="Download Injected Video", type="filepath", visible=True)
        
        submit_btn.click(
            fn=inject_360_metadata,
            inputs=[video_input, stereo_dropdown, spatial_audio_checkbox, crop_input],
            outputs=output_file,
            postprocess=update_output
        )

    demo.launch(share=True)

if __name__ == "__main__":
    main()