"""Drives a running ComfyUI through the browser to export a workflow as a script.

Opens a workflow, clicks the "Save as Script" button, and writes the downloaded Python
file to --output. Used by the daily CI run, but it works against any local ComfyUI:

    python e2e_save_as_script.py --workflow test_workflow.json --output out.py
"""

import argparse
import os
import re
import sys

from playwright.sync_api import sync_playwright, TimeoutError as PlaywrightTimeoutError

# The frontend takes a while to boot on a cold CPU-only container
READY_TIMEOUT = 180_000
ACTION_TIMEOUT = 60_000
# Generating the script imports every node module, so give it room
DOWNLOAD_TIMEOUT = 300_000


def dismiss_startup_dialogs(page):
    """Closes the welcome/template dialogs some frontend versions open on first load."""
    for _ in range(3):
        closed = False
        for selector in ('[data-pc-name="dialog"] [aria-label="Close"]',
                         '[data-pc-name="dialog"] .p-dialog-close-button',
                         '.p-dialog-header-close'):
            button = page.locator(selector)
            if button.count() and button.first.is_visible():
                button.first.click()
                closed = True
                page.wait_for_timeout(500)
        if not closed:
            break
    page.keyboard.press("Escape")
    page.wait_for_timeout(500)


def first_visible(locator):
    """Returns the first visible match, or None."""
    for i in range(locator.count()):
        if locator.nth(i).is_visible():
            return locator.nth(i)
    return None


def click_menu_item(page, label):
    """Clicks a menu entry by its visible label."""
    exact = re.compile(rf"^\s*{re.escape(label)}\s*$", re.IGNORECASE)
    for candidate in (page.get_by_role("menuitem", name=exact),
                      page.locator(f'xpath=//*[normalize-space(text())="{label}"]')):
        item = first_visible(candidate)
        if item:
            item.click()
            return True
    return False


def open_workflow(page, workflow):
    """Loads `workflow` the way a user opens a workflow file.

    Current ComfyUI has no File > Open menu; the open-workflow command
    (Comfy.OpenWorkflow) is bound to Ctrl+O and clicks a hidden file input. Older
    frontends exposed the same command under File > Open, so that is tried second.
    """
    try:
        with page.expect_file_chooser(timeout=20_000) as chooser_info:
            page.keyboard.press("Control+o")
        chooser_info.value.set_files(workflow)
        return
    except PlaywrightTimeoutError:
        print("Ctrl+O did not open a file chooser; trying the File > Open menu",
              file=sys.stderr)

    with page.expect_file_chooser(timeout=ACTION_TIMEOUT) as chooser_info:
        if not click_menu_item(page, "File"):
            raise RuntimeError(
                "Could not open a workflow: Ctrl+O opened no file chooser and there is "
                "no 'File' menu. ComfyUI's open-workflow entry point may have changed."
            )
        page.wait_for_timeout(500)
        if not click_menu_item(page, "Open"):
            raise RuntimeError("Found a 'File' menu but no 'Open' item inside it.")
    chooser_info.value.set_files(workflow)


def wait_for_workflow(page, expect_node):
    """Waits until the loaded graph actually contains `expect_node`."""
    if not expect_node:
        page.wait_for_timeout(2000)
        return
    try:
        page.wait_for_function(
            """(type) => (window.app?.graph?._nodes ?? []).some(
                   (n) => n?.type === type || n?.comfyClass === type)""",
            arg=expect_node,
            timeout=ACTION_TIMEOUT,
        )
    except PlaywrightTimeoutError:
        # Not every frontend exposes window.app; the exported script is checked
        # separately, so this is a warning rather than a failure.
        print(f"warning: could not confirm a {expect_node} node via window.app",
              file=sys.stderr)
    page.wait_for_timeout(1000)


def find_save_as_script_button(page):
    """Returns the visible "Save as Script" button.

    Matching is on button text, not accessible name: the menubar button carries an
    aria-label ("Save the current workflow as a Python script") that would otherwise
    hide it from a name-based lookup. The extension also adds a button to the legacy
    .comfy-menu sidebar, which is normally hidden.
    """
    label = re.compile("Save as Script", re.IGNORECASE)
    for candidate in (page.locator("button").filter(has_text=label),
                      page.get_by_role("button", name=label),
                      page.locator('button[aria-label*="Python script" i]')):
        button = first_visible(candidate)
        if button:
            return button

    raise RuntimeError(
        "The 'Save as Script' button is not present. The extension's setup() may have "
        "failed to run - check the browser console output above."
    )


def export_script(page, name):
    """Clicks "Save as Script", answers the filename prompt, returns the download."""
    button = find_save_as_script_button(page)

    # The extension asks for the filename with window.prompt(); without a handler
    # Playwright auto-dismisses it and no download is ever produced.
    page.once("dialog", lambda dialog: dialog.accept(name))

    with page.expect_download(timeout=DOWNLOAD_TIMEOUT) as download_info:
        button.click()
    return download_info.value


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--url", default="http://127.0.0.1:8188", help="Where ComfyUI is serving")
    parser.add_argument("--workflow", required=True, help="Workflow JSON to open")
    parser.add_argument("--output", required=True, help="Where to write the exported script")
    parser.add_argument("--expect-node", default=None,
                        help="Node type that must appear in the graph once loaded")
    parser.add_argument("--artifacts", default=None, help="Directory for failure screenshots")
    parser.add_argument("--headed", action="store_true", help="Show the browser (for debugging)")
    args = parser.parse_args()

    workflow = os.path.abspath(args.workflow)
    output = os.path.abspath(args.output)
    os.makedirs(os.path.dirname(output) or ".", exist_ok=True)
    if args.artifacts:
        os.makedirs(args.artifacts, exist_ok=True)

    with sync_playwright() as playwright:
        browser = playwright.chromium.launch(
            headless=not args.headed,
            # Chromium crashes with the small /dev/shm containers get by default
            args=["--disable-dev-shm-usage", "--no-sandbox"],
        )
        context = browser.new_context(accept_downloads=True)
        page = context.new_page()
        page.set_default_timeout(ACTION_TIMEOUT)
        page.on("console", lambda msg: print(f"[browser:{msg.type}] {msg.text}"))
        page.on("pageerror", lambda err: print(f"[browser:pageerror] {err}", file=sys.stderr))

        try:
            page.goto(args.url, wait_until="domcontentloaded", timeout=READY_TIMEOUT)
            page.wait_for_selector("#graph-canvas", timeout=READY_TIMEOUT)
            page.wait_for_timeout(3000)
            dismiss_startup_dialogs(page)

            open_workflow(page, workflow)
            wait_for_workflow(page, args.expect_node)

            download = export_script(page, os.path.splitext(os.path.basename(output))[0])
            download.save_as(output)
            print(f"Exported script to {output}")
        except Exception:
            if args.artifacts:
                page.screenshot(path=os.path.join(args.artifacts, "failure.png"), full_page=True)
                with open(os.path.join(args.artifacts, "failure.html"), "w", encoding="utf-8") as f:
                    f.write(page.content())
                print(f"Wrote failure artifacts to {args.artifacts}", file=sys.stderr)
            raise
        finally:
            context.close()
            browser.close()

    if not os.path.getsize(output):
        raise SystemExit("The downloaded script is empty")


if __name__ == "__main__":
    main()