# We use objectid to find objects on the js side
js_uuid(object) = string(objectid(object))

function Bonito.print_js_code(io::IO, plot::AbstractPlot, context::Bonito.JSSourceContext)
    uuids = js_uuid.(Makie.collect_atomic_plots(plot))
    # This is a bit more complicated then it has to be, since evaljs / on_document_load
    # isn't guaranteed to run after plot initialization in an App... So, if we don't find any plots,
    # we have to check again after inserting new plots
    return Bonito.print_js_code(
        io, js"""(new Promise(resolve => {
            $(WGL).then(WGL=> {
                const find = ()=> {
                    const plots = WGL.find_plots($(uuids))
                    if (plots.length > 0) {
                        resolve(plots)
                    } else {
                        WGL.on_next_insert(find)
                    }
                };
                find()
            })
        }))""", context
    )
end

function Bonito.print_js_code(io::IO, scene::Scene, context::Bonito.JSSourceContext)
    code = js"""$(WGL).then(WGL=> {
        function try_find_scene(_retries) {
            let retries = _retries || 0;
            const max_retries = 100;
            const retry_delay = 100;
            const scene = WGL.find_scene($(js_uuid(scene)));
            if (scene) {
                return Promise.resolve(scene);
            } else if (retries < max_retries) {
                return new Promise(resolve => {
                    setTimeout(() => {
                        try_find_scene(retries + 1).then(resolve);
                    }, retry_delay);
                });
            } else {
                return Promise.reject(new Error("Scene not found after retries"));
            }
        }
        return try_find_scene();
    })"""
    return Bonito.print_js_code(io, code, context)
end


function get_order!(session::Session)
    order = Bonito.get_metadata(session, :wglmakie_scene_order, 1)
    Bonito.set_metadata!(session, :wglmakie_scene_order, order + 1)
    return order
end

function three_display(screen::Screen, session::Session, scene::Scene)
    config = screen.config
    order = get_order!(session)
    window_open = scene.events.window_open
    initial_size = size(scene)
    canvas_width = lift(x -> [round.(Int, widths(x))...], scene, viewport(scene))
    is_offline = Bonito.root_session(session).connection isa Bonito.NoConnection
    # Observable to receive the actual canvas size from JS after resize_to calculation
    real_size = Observable{Any}(nothing)
    # Create observable for scene serialization that updates asynchronously
    scene_serialized = Observable{Union{Nothing, Dict{Symbol, Any}}}(nothing)
    done_init = Observable{Any}(nothing)
    if is_offline
        # For offline connections, we have to serialize immediately
        # Since we cant do any round trip communication
        scene_serialized[] = serialize_scene(scene)
    else
        # Query the real canvas size (resize_to) from JS first, resize the scene to
        # it, THEN serialize — so the browser renders at the final size directly
        # instead of rendering at `initial_size` and re-laying-out via observable
        # updates afterwards (which caused a visible resize + a race on capture).
        on(session, real_size) do size_arr
            Makie.async_tracked() do should_close
                try
                    size_tuple = (round.(Int, (size_arr))...,)
                    if size_tuple != initial_size
                        resize!(scene, size_tuple...)
                    end
                    scene_serialized[] = serialize_scene(scene)
                catch e
                    @warn "Error resizing/serializing scene" exception = (e, catch_backtrace())
                    done_init[] = e
                end
            end
            return
        end
    end
    width, height = initial_size
    # Create canvas
    canvas = DOM.m(
        "canvas";
        tabindex = "0",
        # Set with/height to have a good inital size - might not match the final size with scaling etc, but this
        # will be adjusted in JS - this helps with less re-layoting
        width = "$(width)px",
        height = "$(height)px",
        style = "display: block",
        # Pass JupyterLab specific attributes to prevent it from capturing keyboard shortcuts
        # and to suppress the JupyterLab context menu in Makie plots, see:
        # https://jupyterlab.readthedocs.io/en/4.2.x/extension/notebook.html#keyboard-interaction-model
        # https://jupyterlab.readthedocs.io/en/4.2.x/extension/extension_points.html#context-menu
        dataLmSuppressShortcuts = true, dataJpSuppressContextMenu = nothing,
    )

    # Get spinner from config (will be constructed from theming)
    # Wrapper contains canvas and spinner as siblings, plus widgets will be added later
    # position: relative is needed for:
    # 1. absolute positioning of spinner on top of canvas
    # 2. absolute positioning of widgets (HTML widgets, etc.)
    wrapper = DOM.div(canvas, config.spinner; style = "width: 100%; height: 100%; position: relative;")
    comm = Observable(Dict{String, Any}())

    # Keep texture atlas in parent session, so we don't need to send it over and over again
    evaljs(
        session, js"""
        $(WGL).then(WGL => {
            WGL.execute_in_order($order, ()=> {
                WGL.setup_scene_init(
                    $wrapper,
                    $canvas,
                    $width,
                    $height,
                    $(config.resize_to),
                    $(config.px_per_unit),
                    $(config.scalefactor),
                    $(real_size),
                    $canvas_width,
                    $(scene_serialized),
                    $comm,
                    $(config.framerate),
                    $(done_init)
                )
            })
        })
        """
    )
    # push!(Bonito.children(wrapper), jss)
    on(session, done_init) do val
        window_open[] = true
    end
    connect_scene_events!(screen, scene, comm)
    return wrapper, done_init
end
