import gradio as gr

# Custom CSS: smaller arrow, pushed down, video enlarged
mycss = """
.arrow {
  display: flex;
  align-items: center;
  justify-content: center;
  margin: 80px auto;
}
.arrow-tail {
  width: 150px;
  height: 20px;
  background-color: black;
}
.arrow-head {
  width: 0;
  height: 0;
  border-top: 30px solid transparent;
  border-bottom: 30px solid transparent;
  border-left: 60px solid black;
}
@media (prefers-color-scheme: dark) {
  .arrow-tail { background-color: white; }
  .arrow-head { border-left-color: white; }
}
/* ensure the video container has enough space */
#generated-video {
  min-height: 360px;
}
/* enlarge the video */
#generated-video video {
  width: 640px !important;
  height: auto !important;
}
"""

# Arrow HTML
myhtml = """
<div class="arrow">
  <div class="arrow-tail"></div>
  <div class="arrow-head"></div>
</div>
"""

def generate_media(echo_idx, ecg_idx):
    echo_path = f"resources/examples/ef{echo_idx}.png"
    ecg_path  = f"resources/examples/ecg{ecg_idx}.png"
    video_path = f"resources/examples/ef{echo_idx}.mp4"
    return echo_path, ecg_path, video_path

with gr.Blocks(css=mycss) as demo:
    gr.Markdown("""
    # ECHOPULSE: ECG-Controlled Echocardiogram Video Generation  
    Enter the example index for your ECHO image and ECG image below, then click **Generate Video** to load both images and play the corresponding video.
    """)

    # Two input fields for the example index
    echo_idx_input = gr.Textbox(label="Patient's Past ECHO Example Index (i)", placeholder="e.g. 0000")
    ecg_idx_input  = gr.Textbox(label="Patient's ECG Example Index (i)",  placeholder="e.g. 0000")

    with gr.Row():
        # Column 1: ECHO + ECG images
        with gr.Column():
            echo_image = gr.Image(interactive=False, label="ECHO Image")
            ecg_image  = gr.Image(interactive=False, label="ECG Image")
        # Column 2: arrow + button
        with gr.Column():
            gr.HTML(myhtml)
            gen_button = gr.Button("Generate Video")
        # Column 3: video output
        with gr.Column():
            video_out = gr.Video(interactive=False,
                                 autoplay=True,
                                 label="Generated Video",
                                 elem_id="generated-video")

    # Wire up the button to load images & video based on the indices
    gen_button.click(
        fn=generate_media,
        inputs=[echo_idx_input, ecg_idx_input],
        outputs=[echo_image, ecg_image, video_out]
    )

    # Delay display of the <video> element by 2s after it's loaded
    gr.HTML("""
<script>
const wrapper = document.getElementById('generated-video');
const observer = new MutationObserver(mutations => {
  mutations.forEach(mutation => {
    mutation.addedNodes.forEach(node => {
      if (node.tagName === 'VIDEO') {
        node.style.display = 'none';
        node.addEventListener('loadeddata', () => {
          setTimeout(() => {
            node.style.display = 'block';
            if (node.paused) node.play();
          }, 2000);
        });
      }
    });
  });
});
observer.observe(wrapper, { childList: true });
</script>
""")

demo.launch(share=False)