from __future__ import annotations

from contextlib import contextmanager

from modules import img2img, processing, shared


def cn_restore_unet_hook(p, cn_latest_network):
    if cn_latest_network is not None:
        unet = p.sd_model.model.diffusion_model
        cn_latest_network.restore(unet)


class CNHijackRestore:
    def __init__(self):
        self.process = hasattr(processing, "__controlnet_original_process_images_inner")
        self.img2img = hasattr(img2img, "__controlnet_original_process_batch")

    def __enter__(self):
        if self.process:
            self.orig_process = processing.process_images_inner
            processing.process_images_inner = getattr(
                processing, "__controlnet_original_process_images_inner"
            )
        if self.img2img:
            self.orig_img2img = img2img.process_batch
            img2img.process_batch = getattr(
                img2img, "__controlnet_original_process_batch"
            )

    def __exit__(self, *args, **kwargs):
        if self.process:
            processing.process_images_inner = self.orig_process
        if self.img2img:
            img2img.process_batch = self.orig_img2img


@contextmanager
def cn_allow_script_control():
    orig = False
    if "control_net_allow_script_control" in shared.opts.data:
        try:
            orig = shared.opts.data["control_net_allow_script_control"]
            shared.opts.data["control_net_allow_script_control"] = True
            yield
        finally:
            shared.opts.data["control_net_allow_script_control"] = orig
    else:
        yield