import gradio as gr

import altair as alt
from vega_datasets import data

import re

CLEAN_RE = r"<(\/*head|\/*html|\/*body|!DOCTYPE html)>"
SELECT_RE = r"document.getElementById\('vis'\)"
SELECT_RE_TWO = r"vegaEmbed\(\"#vis\", spec, embedOpt\)"
x = r".catch\(error => showError\(el, error\)\)"


def give_chart():
    source = data.cars()
    alt.Chart(source).mark_circle().encode(
        x='Horsepower',
        y='Miles_per_Gallon',
        color='Origin',
    ).interactive().save('chart.html')
    f = open("chart.html", "r").read()
    f = re.sub(CLEAN_RE, "", f, flags=re.IGNORECASE)
    f = re.sub(
        SELECT_RE, "document.getElementsByTagName(\"gradio-app\")[0].shadowRoot.querySelectorAll(\"#vis\")[0]", f, flags=re.IGNORECASE)
    f = re.sub(
        SELECT_RE_TWO, "vegaEmbed(document.getElementsByTagName(\"gradio-app\")[0].shadowRoot.querySelectorAll(\"#vis\")[0], spec, embedOpt)", f, flags=re.IGNORECASE)

    return f


with gr.Blocks(css=".output-html \{ display: flex;justify-content: center;\} ") as demo:
    html = gr.HTML()
    btn = gr.Button("Give me chart")
    btn.click(give_chart, inputs=None, outputs=html,
              _js="""
      (x) => {
        const html_wrapper = document.getElementsByTagName("gradio-app")[0].shadowRoot.querySelectorAll(".output-html")[0]
        const head = document.getElementsByTagName("head")[0];

        const callback = (mutationList, observer) => {
          for (const mutation of mutationList) {
            if (mutation.type === 'childList') {
              observer.disconnect();
              const scripts = Array.from(html_wrapper.querySelectorAll("script"))
                .map(e => [e.type, e.src, e.textContent])
              const wait = [];
              const last = scripts.splice(scripts.length - 1)[0];
              scripts.forEach(([type, src, inner], i, arr) => {
                let r;
                wait.push(new Promise((success) => r = success))
                const s = document.createElement("script")
                if (type) s.type = type;
                if (src) s.src = src;
                s.onload = () => r();
                html_wrapper.appendChild(s);
                s.remove();
              })

              Promise.all(wait).then(() => {
                const s = document.createElement("script")
                s.innerHTML = last[2]
                html_wrapper.appendChild(s);
                
                s.remove();
              })
            }
          }
        };

        const observer = new MutationObserver(callback);
        observer.observe(
            html_wrapper, { attributes: false, childList: true, subtree: true }
        );

      }
    """)

demo.launch()