import json
import os.path
import re
from typing import List
import folder_paths  # noqa
from folder_paths import models_dir  # noqa
import comfy.sd  # noqa
import comfy.utils  # noqa
from nodes import LoraLoader  # noqa
from .utils.lora_enc_map import FELoraEmpFinder, read_automap, read_emp
from .utils.any_hack import any
from .utils.model_info import get_metadata


class LoraRef:
    def __init__(self, name: str, model_s: float, clip_s: float):
        self.name = name
        self.model_s = model_s
        self.clip_s = clip_s


def find_model(text, model_name):
    pattern = r'<lora:' + re.escape(model_name) + r':(\d+(?:\.\d+)?):(\d+(?:\.\d+)?)>'
    match = re.search(pattern, text)

    if match:
        model_s = float(match.group(1))
        clip_s = float(match.group(2))
        return match.start(), model_s, clip_s, match.end()
    else:
        return None


class FEEncLoraAutoLoaderStack(LoraLoader, FELoraEmpFinder):
    def __init__(self):
        super().__init__()
        self.loaded_lora = None

    @classmethod
    def INPUT_TYPES(s):
        s.EMP_CACHE = read_emp("loras")
        return {
            "required": {
                "model_type": (["sd1.5", "sdxl", "sd3", "sd3.5", 'pony', 'il', 'noob', 'flux'],),
                "prompt": ("STRING", {"default": "", 'forceInput': True}),
                "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
                "strength_clip": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
            }
        }

    RETURN_TYPES = ("LORA_STACK", 'STRING', any)
    RETURN_NAMES = ("lora_stack", "prompt", "EXTRA")
    FUNCTION = "load_lora_stack"

    CATEGORY = "loaders"

    @staticmethod
    def get_replace(lora_setting, trigger_wd):
        replace = lora_setting.get("replace", None)
        if replace is None:
            return trigger_wd
        if isinstance(replace, dict):
            replace = replace.get(trigger_wd, trigger_wd)
        if isinstance(replace, list):
            return replace[0]
        return str(replace)

    def load_lora_stack(self, model_type, prompt, strength_model, strength_clip):
        auto_loras = read_automap("loras")
        lora_need_load = []  # type:List[LoraRef]
        pe = str(prompt)
        pe_lower = pe  # .lower()

        # lora_pure_names
        for lora_name, tw in [(_, _.split("]")[-1].rsplit(".", 1)[0]) for _ in self.EMP_CACHE.keys()]:
            lora_find_rs = find_model(pe_lower, tw)
            if lora_find_rs is None:
                continue
            pe = pe[:lora_find_rs[0]] + pe[lora_find_rs[3]:]
            pe_lower = pe
            model_s = lora_find_rs[1]
            clip_s = lora_find_rs[2]
            lora_need_load.append(LoraRef(lora_name, model_s, clip_s))

        # auto lora maps
        for lora_name, lora_setting in auto_loras.items():
            if lora_setting.get("model", "?") != model_type:
                continue
            tw = lora_setting.get("trigger_word", "")
            if tw == "" or tw not in pe_lower:
                continue
            # 存在，检查是否为<xx:f:f>
            lora_find_rs = find_model(pe_lower, tw)
            if lora_find_rs is None:
                continue
            pe = pe[:lora_find_rs[0]] + pe[lora_find_rs[3]:]
            pe_lower = pe
            model_s = lora_find_rs[1]
            clip_s = lora_find_rs[2]
            lora_need_load.append(LoraRef(lora_name, model_s, clip_s))

        result = list()
        lcnt = 0

        for lora_ref in lora_need_load:
            lora_name = self.find_current_lora(lora_ref.name)
            if not lora_name:
                print("Failed to load lora with empty name", lora_ref.name)
                continue
            result.extend([(lora_ref.name, lora_ref.model_s, lora_ref.clip_s)]),
        print("Auto Load", lcnt, "LLor ClipData")
        return (result, pe, {},)