import gradio as gr

from modules import scripts, sd_samplers_common


store_latent = sd_samplers_common.store_latent
steps = []


class StepExplorerScript(scripts.Script):
    def title(self):
        return "Step Explorer"

    def show(self, is_img2img):
        return scripts.AlwaysVisible

    def setup(self, p, *args):
        return super().setup(p, *args)

    def ui(self, is_img2img):
        with gr.Row():
            with gr.Column():
                is_enabled = gr.Checkbox(
                    label="Enable Step Exploration",
                    value=False,
                    elem_id=self.elem_id("is_enabled"),
                )
        gr.HTML("<br />")

        return [is_enabled]

    def process(self, p, is_enabled):
        global steps, store_latent
        if not is_enabled:
            return

        steps = []

        def capture_latent(decoded):
            steps.append(sd_samplers_common.sample_to_image(decoded))
            store_latent(decoded)

        setattr(sd_samplers_common, "store_latent", capture_latent)

    def postprocess(self, p, processed, is_enabled):
        global steps, store_latent
        if not is_enabled:
            return

        processed.images = steps
        processed.index_of_first_image = 0
        prompt, info = processed.infotexts[0].split("\n")
        processed.infotexts = [
            f"{prompt}\nStep: {i+1}, {info}" for i in range(len(steps))
        ]

        setattr(sd_samplers_common, "store_latent", store_latent)