import re
from typing import TypeVar
from sd_mecha.extensions.merge_methods import merge_method, StateDict, Parameter, Return
from sd_mecha.extensions import model_configs


sdxl_sgm_config = model_configs.resolve("sdxl-sgm")


re_inp = re.compile(r"\.input_blocks\.(\d+)\.")  # 12
re_mid = re.compile(r"\.middle_block\.(\d+)\.")  # 1
re_out = re.compile(r"\.output_blocks\.(\d+)\.")  # 12


T = TypeVar("T")


# srcs, in order of priority when disagreements occur:
# - https://github.com/hako-mikan/sd-webui-supermerger/blob/f14b3e5d0be9c510d199cca502c4148160f901bb/scripts/mergers/mergers.py#L1376
# - https://github.com/s1dlx/meh/blob/04af2c8d63744fb6c02d35d328a2c84380cca444/sd_meh/merge.py#L360
# - https://github.com/vladmandic/automatic/blob/e22d0789bddd3894364b0d59a4c9b3e456e89079/modules/merging/merge_utils.py#L64
@merge_method(is_conversion=True, reuse_outputs=False)
class convert_sdxl_blocks_to_sgm:
    @staticmethod
    def map_keys(b):
        for output_key in sdxl_sgm_config.keys():
            input_key = "BASE"
            if output_key.startswith("model.diffusion_model."):
                input_key = "OUT08"
                if ".time_embed" in output_key or ".label_emb" in output_key:
                    input_key = "BASE"  # before input blocks
                elif ".out." in output_key:
                    input_key = "OUT08"  # after output blocks
                elif m := re_inp.search(output_key):
                    input_key = f"IN{int(m.groups(1)[0]):02}"
                elif re_mid.search(output_key):
                    input_key = "M00"
                elif m := re_out.search(output_key):
                    input_key = f"OUT{int(m.groups(1)[0]):02}"
            elif output_key.startswith("first_stage_model."):
                input_key = "VAE"

            b[output_key] = b.keys[input_key]

    def __call__(
        self,
        blocks: Parameter(StateDict[T], model_config="sdxl-supermerger_blocks"),
        **kwargs,
    ) -> Return(T, model_config="sdxl-sgm"):
        relation = kwargs["key_relation"]
        block_key = relation.inputs["blocks"][0]
        return blocks[block_key]
