import gradio as gr
import prompt_enhancer
import video_merger
from image_utils import extract_frame, get_timestamp_js
from pipeline_manager import generate_video
from ui_modules import (
    CSS, load_template, get_header_html,
    render_column_inputs, render_column_player,
    render_sulphur_vision_suite, render_vip_suite, render_cpu_tools
)

def create_ui():
    with gr.Blocks(delete_cache=(3600, 10800), title="I2V EXTENDED v3.5 VIP") as demo:
        gr.HTML(f"<style>{CSS}</style>" + get_header_html())

        with gr.Row():
            # Column 1: Standard Generation Inputs (scale=4)
            c1 = render_column_inputs()

            # Column 2: Video Player & Output Tools (scale=4)
            c2 = render_column_player()

            # Column 3: Specialized Tool Suite (Stacked Accordions Top-to-Bottom, scale=4)
            with gr.Column(scale=4):
                silver = render_sulphur_vision_suite(
                    input_img=c1['input_image_component']
                )

                vip = render_vip_suite(
                    input_img=c1['input_image_component'],
                    output_video=c2['video_output']
                )

                cpu = render_cpu_tools(
                    input_img=c1['input_image_component']
                )

        # -------------------------------------------------------------
        # GLOBAL ACTION CALLBACKS & CLICK HANDLERS
        # -------------------------------------------------------------

        # Sulphur Vision Enhancer Button Click (No VIP password required!)
        silver['sulphur_enhance_btn'].click(
            fn=prompt_enhancer.call_sulphur_enhancer_api,
            inputs=[
                c1['input_image_component'],
                silver['sulphur_subject_input'],
                silver['sulphur_adegan_input'],
                silver['sulphur_camera_dropdown'],
                silver['sulphur_atmosphere_dropdown'],
                c1['duration_seconds_input'],
                silver['sulphur_mode_dropdown'],
                silver['sulphur_enable_swap_checkbox']
            ],
            outputs=[
                silver['relay_prompt_schedule_input'],
                silver['enable_prompt_relay_checkbox'],
                silver['sulphur_api_status_box'],
                c1['input_image_component']
            ]
        )

        # Main Video Generator Click Handler
        ui_inputs = [
            c1['input_image_component'], c2['last_image_component'], c1['prompt_input'], vip['steps_slider'],
            vip['negative_prompt_input'], c1['duration_seconds_input'],
            vip['guidance_scale_input'], vip['guidance_scale_2_input'], vip['seed_input'], vip['randomize_seed_checkbox'],
            vip['quality_slider'], vip['scheduler_dropdown'], vip['flow_shift_slider'], c1['frame_multi'],
            c1['motion_extension_dropdown'],
            c1['safe_mode_checkbox'],
            c1['custom_lora_url_input'],
            c1['custom_lora_scale_input'],
            silver['enable_prompt_relay_checkbox'],
            silver['relay_prompt_schedule_input'],
            gr.State(None),
            gr.State("Female Faces Only"),
            c2['play_result_video'],
            vip['custom_filename_input'],
            c1['noise_temperature_slider'],
            vip['enable_vip_rife_checkbox'],
            vip['vip_rife_multiplier_radio'],
            vip['vip_rife_mode_dropdown'],
            vip['vip_rife_upscale_checkbox'],
            vip['vip_rife_enhance_face_checkbox'],
            vip['vip_master_password_input'],
            c1['cached_lora_dropdown'],
            c1['extra_gpu_buffer_slider']
        ]
        
        c1['generate_button'].click(
            fn=generate_video, 
            inputs=ui_inputs, 
            outputs=[c2['video_output'], c2['file_output'], vip['seed_input'], c2['gpu_report_box']]
        )

        # Standalone Manual Video Merge Click Handlers
        c2['append_merge_btn'].click(
            fn=video_merger.handle_append_video,
            inputs=[c2['file_output'], c2['merge_queue_state'], vip['custom_filename_input']],
            outputs=[c2['merged_video_output'], c2['merged_file_output'], c2['merge_queue_state'], c2['merge_status_box']]
        )

        c2['clear_merge_btn'].click(
            fn=video_merger.handle_clear_merge_queue,
            inputs=None,
            outputs=[c2['merged_video_output'], c2['merged_file_output'], c2['merge_queue_state'], c2['merge_status_box']]
        )

        # 3-Way Instant Video Frame Extraction Handlers
        c2['grab_frame_btn'].click(
            fn=extract_frame,
            inputs=[c2['video_output'], c2['timestamp_box']],
            outputs=[c1['input_image_component'], c2['timestamp_box']],
            js=get_timestamp_js
        )

        c2['grab_last_frame_btn'].click(
            fn=extract_frame,
            inputs=[c2['video_output'], c2['timestamp_box']],
            outputs=[c2['last_image_component'], c2['timestamp_box']],
            js=get_timestamp_js
        )

        c2['grab_swap_target_btn'].click(
            fn=extract_frame,
            inputs=[c2['video_output'], c2['timestamp_box']],
            outputs=[cpu['tool_target_img'], c2['timestamp_box']],
            js=get_timestamp_js
        )

        gr.HTML(load_template("footer.html"))

    return demo