import os
import json
import folder_paths
import comfy.sd
import comfy.utils
from safetensors import safe_open
import hashlib
import urllib.request

CACHE_FILE = os.path.join(os.path.dirname(__file__), "dolphin_tags_cache.json")

def load_dolphin_cache():
    if os.path.exists(CACHE_FILE):
        try:
            with open(CACHE_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except: return {}
    return {}

def save_dolphin_cache(cache_data):
    try:
        with open(CACHE_FILE, "w", encoding="utf-8") as f:
            json.dump(cache_data, f, indent=4)
    except: pass

class DolphinArchitectV15:
    @classmethod
    def INPUT_TYPES(s):
        loras = ["None"] + folder_paths.get_filename_list("loras")
        
        res = {
            "required": {
                "global_switch": ("BOOLEAN", {"default": True, "label_on": "GLOBAL ON", "label_off": "GLOBAL OFF (SKIP)"}),
                "model_base": ("MODEL",), "clip_base": ("CLIP",),
                "model_high": ("MODEL",), "clip_high": ("CLIP",),
                "model_low": ("MODEL",), "clip_low": ("CLIP",),
            },
            "optional": {}
        }

        # [다이어트 완료] 오직 이름과 웨이트만 남겨서 가장 얇은 UI를 만듭니다.
        res["optional"]["header_master"] = ("STRING", {"default": "─── [ MASTER: B+H ] ───", "multiline": False})
        for i in range(1, 5): 
            res["optional"][f"master_l_{i}"] = (loras, {"default": "None"})
            res["optional"][f"master_b_w_{i}"] = ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01})
            res["optional"][f"master_h_w_{i}"] = ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01})
            
        res["optional"]["header_low"] = ("STRING", {"default": "─── [ LOW ONLY ] ───", "multiline": False})
        for i in range(1, 5): 
            res["optional"][f"low_l_{i}"] = (loras, {"default": "None"})
            res["optional"][f"low_w_{i}"] = ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01})

        res["optional"]["manual_triggers"] = ("STRING", {"default": "", "multiline": True})
        return res

    RETURN_TYPES = ("MODEL", "CLIP", "MODEL", "CLIP", "MODEL", "CLIP", "STRING")
    RETURN_NAMES = ("M_BASE", "C_BASE", "M_HIGH", "C_HIGH", "M_LOW", "C_LOW", "triggers")
    FUNCTION = "exec"
    CATEGORY = "Dolphin"

    def get_trig(self, lora_name):
        if lora_name == "None": return []
        path = folder_paths.get_full_path("loras", lora_name)
        if not path or not os.path.exists(path): return []
        
        triggers = []
        try:
            with safe_open(path, framework="pt", device="cpu") as f:
                m = f.metadata()
                if m and "modelspec.trigger_words" in m:
                    tw = json.loads(m["modelspec.trigger_words"])
                    if isinstance(tw, list): triggers.extend(tw)
                    elif isinstance(tw, str): triggers.append(tw)
        except: pass
        if triggers: return triggers 

        cache = load_dolphin_cache()
        if lora_name in cache and cache[lora_name]: return cache[lora_name]

        try:
            sha256_hash = hashlib.sha256()
            with open(path, "rb") as f:
                for byte_block in iter(lambda: f.read(4096 * 1024), b""):
                    sha256_hash.update(byte_block)
            
            url = f"https://civitai.com/api/v1/model-versions/by-hash/{sha256_hash.hexdigest()}"
            req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
            with urllib.request.urlopen(req, timeout=5) as response:
                data = json.loads(response.read().decode('utf-8'))
                api_triggers = data.get("trainedWords", [])
                if api_triggers:
                    cache[lora_name] = api_triggers
                    save_dolphin_cache(cache)
                    return api_triggers
        except: pass
        return triggers

    def exec(self, global_switch, model_base, clip_base, model_high, clip_high, model_low, clip_low, manual_triggers="", **kwargs):
        mb, cb = model_base, clip_base
        mh, ch = model_high, clip_high
        ml, cl = model_low, clip_low
        
        if not global_switch:
            trigs = [t.strip() for t in manual_triggers.split(',') if t.strip()]
            ut = ", ".join(list(dict.fromkeys(trigs)))
            return (mb, cb, mh, ch, ml, cl, ut)

        applied_names = []

        for i in range(1, 5): 
            n = kwargs.get(f"master_l_{i}", "None")
            w_b = kwargs.get(f"master_b_w_{i}", 1.0)
            w_h = kwargs.get(f"master_h_w_{i}", 1.0)
            
            # [최적화] None이 아니면서, 웨이트가 0.0이 아닐 때만 적용 (0.0이면 스위치 OFF 효과)
            if n != "None" and (w_b != 0.0 or w_h != 0.0):
                applied_names.append(n)
                lora_data = comfy.utils.load_torch_file(folder_paths.get_full_path("loras", n))
                mb, cb = comfy.sd.load_lora_for_models(mb, cb, lora_data, w_b, w_b)
                mh, ch = comfy.sd.load_lora_for_models(mh, ch, lora_data, w_h, w_h)

        for i in range(1, 5): 
            n = kwargs.get(f"low_l_{i}", "None")
            w = kwargs.get(f"low_w_{i}", 1.0)
            
            if n != "None" and w != 0.0:
                applied_names.append(n)
                lora_data = comfy.utils.load_torch_file(folder_paths.get_full_path("loras", n))
                ml, cl = comfy.sd.load_lora_for_models(ml, cl, lora_data, w, w)

        trigs = []
        for n in set(applied_names):
            trigs.extend(self.get_trig(n))
            
        if manual_triggers.strip():
            manual_list = [t.strip() for t in manual_triggers.split(',') if t.strip()]
            trigs.extend(manual_list)
        
        unique_trigs = ", ".join(list(dict.fromkeys(trigs)))
        
        return (mb, cb, mh, ch, ml, cl, unique_trigs)