templates.py

from enum import Enum
from pathlib import Path


class DemoType(Enum):
    GRADIO = 1
    STREAMLIT = 2


gradio_lite_html_template = Path('templates/gradio-lite/gradio-lite-template.html').read_text()
stlite_html_template = Path('templates/stlite/stlite-template.html').read_text()
gradio_lite_snippet_template = Path('templates/gradio-lite/gradio-lite-snippet-template.html').read_text()
stlite_snippet_template = Path('templates/stlite/stlite-snippet-template.html').read_text()


def starting_app_code(demo_type: DemoType) -> str:
    if demo_type == DemoType.GRADIO:
        return Path('templates/gradio-lite/gradio_lite_starting_code.py').read_text().replace('`', r'\`')
    elif demo_type == DemoType.STREAMLIT:
        return Path('templates/stlite/stlite_starting_code.py').read_text().replace('`', r'\`')
    raise NotImplementedError(f'{demo_type} is not a supported demo type')


def load_js(demo_type: DemoType) -> str:
    if demo_type == DemoType.GRADIO:
        return f"""() => {{
            if (window.gradioLiteLoaded) {{
                return
            }}
            
            // Get the query string from the URL
            const queryString = window.location.search;
            // Use a function to parse the query string into an object
            function parseQueryString(queryString) {{
                const params = {{}};
                const queryStringWithoutQuestionMark = queryString.substring(1); // Remove the leading question mark
                const keyValuePairs = queryStringWithoutQuestionMark.split('&');
            
                keyValuePairs.forEach(keyValue => {{
                    const [key, value] = keyValue.split('=');
                    if (value) {{
                        params[key] = decodeURIComponent(value.replace(/\+/g, ' '));
                    }}
                }});
            
                return params;
            }}
            // Parse the query string into an object
            const queryParams = parseQueryString(queryString);
            // Access individual parameters
            const codeValue = queryParams.code;
            const requirementsValue = queryParams.requirements;
                        
            const htmlString = '<iframe id="gradio-iframe" width="100%" height="512px" src="about:blank"></iframe>';
            const parser = new DOMParser();
            const doc = parser.parseFromString(htmlString, 'text/html');
            const iframe = doc.getElementById('gradio-iframe'); 
            const div = document.getElementById('gradioDemoDiv');
            div.appendChild(iframe);

            let template = `{gradio_lite_html_template.replace('STARTING_CODE', starting_app_code(demo_type))}`;    
            if (codeValue) {{
                template = `{gradio_lite_html_template}`.replace('STARTING_CODE', codeValue.replaceAll(String.fromCharCode(92), String.fromCharCode(92) + String.fromCharCode(92)).replaceAll('`', String.fromCharCode(92) + '`'));
            }}
            template = template.replace('STARTING_REQUIREMENTS', requirementsValue || '');
            const frame = document.getElementById('gradio-iframe');
            frame.contentWindow.document.open('text/html', 'replace');
            frame.contentWindow.document.write(template);
            frame.contentWindow.document.close();
            window.gradioLiteLoaded = true;
        }}"""
    elif demo_type == DemoType.STREAMLIT:
        return f"""() => {{
            if (window.stliteLoaded) {{
                return
            }}
            const htmlString = '<iframe id="stlite-iframe" width="100%" height="512px" src="about:blank"></iframe>';
            const parser = new DOMParser();
            const doc = parser.parseFromString(htmlString, 'text/html');
            const iframe = doc.getElementById('stlite-iframe'); 
            const div = document.getElementById('stliteDemoDiv');
            div.appendChild(iframe);
            
            const template = `{stlite_html_template.replace('STARTING_CODE', starting_app_code(demo_type))}`;            
            const frame = document.getElementById('stlite-iframe');
            frame.contentWindow.document.open();
            frame.contentWindow.document.write(template);
            frame.contentWindow.document.close();
            window.stliteLoaded = true;
        }}"""
    raise NotImplementedError(f'{demo_type} is not a supported demo type')


def update_iframe_js(demo_type: DemoType) -> str:
    if demo_type == DemoType.GRADIO:
        return f"""async (code, requirements) => {{
                const formattedRequirements = requirements.split('\\n').filter(x => x && !x.startsWith('#')).map(x => x.trim());
                let errorResult = null;
                const attemptedRequirements = new Set();
                const installedRequirements = [];
                async function update() {{
                    // Remove existing stylesheet so it will be reloaded;
                    // see https://github.com/gradio-app/gradio/blob/200237d73c169f39514465efc163db756969d3ac/js/app/src/lite/css.ts#L41
                    const demoFrameWindow = document.getElementById('gradio-iframe').contentWindow;
                    const oldStyle = demoFrameWindow.document.querySelector("head style");
                    oldStyle.remove();
                    const appController = demoFrameWindow.window.appController;
                    const newCode = code + ` # Update tag ${{Math.random()}}`;
                    try {{
                        await appController.install(formattedRequirements);
                        await appController.run_code(newCode);
                    }}
                    catch (e) {{                    
                        // Replace old style if code error prevented new style from loading.
                        const newStyle = demoFrameWindow.document.querySelector("head style");
                        if (!newStyle) {{
                            demoFrameWindow.document.head.appendChild(oldStyle);
                        }}
                        
                        // If the error is caused by a missing module try once to install it and update again.
                        if (e.toString().includes('ModuleNotFoundError')) {{
                            try {{
                                const guessedModuleName = e.toString().split("'")[1].replaceAll('_', '-');
                                if (attemptedRequirements.has(guessedModuleName)) {{
                                    throw Error(`Could not install pyodide module ${{guessedModuleName}}`);
                                }}
                                console.log(`Attempting to install missing pyodide module "${{guessedModuleName}}"`);
                                attemptedRequirements.add(guessedModuleName);
                                await appController.install([guessedModuleName]);
                                installedRequirements.push(guessedModuleName);
                                return await update();
                            }}
                            catch (err) {{
                                console.log(err);
                            }}
                        }}
                        
                        // Hide app so the error traceback is visible.
                        // First div in main is the error traceback, second is the app.
                        const appBody = demoFrameWindow.document.querySelectorAll("div.main > div")[1];
                        appBody.style.visibility = "hidden";
                        errorResult = e.toString();
                        const allRequirements = formattedRequirements.concat(installedRequirements);
                        console.log(allRequirements, errorResult);
                        return [allRequirements, errorResult];
                    }}
                }};
                await update();
                
                const allRequirements = formattedRequirements.concat(installedRequirements);
                // Update URL query params to include the current demo code state
                const currentUrl = new URL(window.location.href);
                currentUrl.searchParams.set('requirements', allRequirements.join('\\n'));
                currentUrl.searchParams.set('code', code);
                // Replace the current URL with the updated one
                history.replaceState({{}}, '', currentUrl.href);
                
                return [allRequirements, errorResult];
            }}"""
    elif demo_type == DemoType.STREAMLIT:
        return f"""async (code, requirements) => {{       
            async function update() {{
                const appController = document.getElementById('stlite-iframe').contentWindow.window.appController;
                const newCode = code + ` # Update tag ${{Math.random()}}`;
                const entrypointFile = "streamlit_app.py";
                appController.writeFile(entrypointFile, newCode);
            }};
            await update();
        }}"""
    raise NotImplementedError(f'{demo_type} is not a supported demo type')


def copy_share_link_js(demo_type: DemoType) -> str:
    if demo_type == DemoType.GRADIO:
        return f"""async () => {{
            const shareLink = document.URL;
            await navigator.clipboard.writeText(shareLink);
        }}"""
    raise NotImplementedError(f'{demo_type} is not a supported demo type')


def copy_snippet_js(demo_type: DemoType) -> str:
    if demo_type == DemoType.GRADIO:
        return f"""async (code, requirements) => {{
            const escapedCode = code.replaceAll(String.fromCharCode(92), String.fromCharCode(92) + String.fromCharCode(92) + String.fromCharCode(92) + String.fromCharCode(92)).replaceAll('`', String.fromCharCode(92) + '`');
            const template = `{gradio_lite_snippet_template}`;
            // Step 1: Generate the HTML content
            const completedTemplate = template.replace('STARTING_CODE', escapedCode).replace('STARTING_REQUIREMENTS', requirements);
            const snippet = completedTemplate;
            await navigator.clipboard.writeText(snippet);
            return [code, requirements];
        }}"""
    elif demo_type == DemoType.STREAMLIT:
        return f"""async (code) => {{
            const escapedCode = code.replaceAll(String.fromCharCode(92), String.fromCharCode(92) + String.fromCharCode(92) + String.fromCharCode(92)).replaceAll('`', String.fromCharCode(92) + '`');
            const template = `{stlite_snippet_template}`;
            // Step 1: Generate the HTML content
            const completedTemplate = template.replace('STARTING_CODE', code);
            
            const snippet = completedTemplate;            
            await navigator.clipboard.writeText(snippet);
        }}"""
    raise NotImplementedError(f'{demo_type} is not a supported demo type')


def download_code_js(demo_type: DemoType) -> str:
    if demo_type == demo_type.GRADIO:
        return f"""(code, requirements) => {{
            const escapedCode = code.replaceAll(String.fromCharCode(92), String.fromCharCode(92) + String.fromCharCode(92)).replaceAll('`', String.fromCharCode(92) + '`');
            // Step 1: Generate the HTML content
            const completedTemplate = `{gradio_lite_html_template}`.replace('STARTING_CODE', escapedCode).replace('STARTING_REQUIREMENTS', requirements);

            // Step 2: Create a Blob from the HTML content
            const blob = new Blob([completedTemplate], {{ type: "text/html" }});

            // Step 3: Create a URL for the Blob
            const url = URL.createObjectURL(blob);

            // Step 4: Create a download link
            const downloadLink = document.createElement("a");
            downloadLink.href = url;
            downloadLink.download = "gradio-lite-app.html"; // Specify the filename for the download

            // Step 5: Trigger a click event on the download link
            downloadLink.click();

            // Clean up by revoking the URL
            URL.revokeObjectURL(url);
        }}"""
    elif demo_type == demo_type.STREAMLIT:
        return f"""(code) => {{
            const escapedCode = code.replaceAll(String.fromCharCode(92), String.fromCharCode(92) + String.fromCharCode(92)).replaceAll('`', String.fromCharCode(92) + '`');
            // Step 1: Generate the HTML content
            const completedTemplate = `{stlite_html_template}`.replace('STARTING_CODE', escapedCode);
            
            // Step 2: Create a Blob from the HTML content
            const blob = new Blob([completedTemplate], {{ type: "text/html" }});
            
            // Step 3: Create a URL for the Blob
            const url = URL.createObjectURL(blob);
            
            // Step 4: Create a download link
            const downloadLink = document.createElement("a");
            downloadLink.href = url;
            downloadLink.download = "stlite-app.html"; // Specify the filename for the download
            
            // Step 5: Trigger a click event on the download link
            downloadLink.click();
            
            // Clean up by revoking the URL
            URL.revokeObjectURL(url);
        }}"""
    raise NotImplementedError(f'{demo_type} is not a supported demo type')