"""ComfyUI custom node: load a LoRA fetched from a URL at request time.

Lets serverless workflows reference LoRAs by URL (HF, civitai, anywhere)
instead of requiring the file to be pre-baked into the image. Downloads are
cached on the worker's disk keyed by URL hash. Supports HF_TOKEN and
CIVITAI_TOKEN env vars for authenticated downloads.
"""

import hashlib
import logging
import os
import urllib.request

import comfy.lora
import comfy.sd
import comfy.utils

CACHE_DIR = "/lora-cache"
os.makedirs(CACHE_DIR, exist_ok=True)


def _fetch(url: str) -> str:
    path = os.path.join(CACHE_DIR, hashlib.sha1(url.encode()).hexdigest() + ".safetensors")
    if os.path.exists(path):
        return path
    headers = {}
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token and "huggingface.co" in url:
        headers["Authorization"] = f"Bearer {hf_token}"
    civitai_token = os.environ.get("CIVITAI_TOKEN")
    if civitai_token and "civitai" in url:
        url += ("&" if "?" in url else "?") + "token=" + civitai_token
    req = urllib.request.Request(url, headers=headers)
    tmp = path + ".part"
    with urllib.request.urlopen(req, timeout=600) as r, open(tmp, "wb") as f:
        while chunk := r.read(1 << 20):
            f.write(chunk)
    os.rename(tmp, path)
    return path


class LoraFromURLModelOnly:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "url": ("STRING", {"default": ""}),
                "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
            }
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load"
    CATEGORY = "loaders"

    def load(self, model, url, strength):
        url = (url or "").strip()
        name = url.rsplit("/", 1)[-1] or "lora"
        if not url or strength == 0:
            logging.info(f"[lora_from_url] {name}: skipped (empty url / strength 0)")
            return (model,)
        lora_sd = comfy.utils.load_torch_file(_fetch(url), safe_load=True)
        total = len(lora_sd)
        # How many lora keys actually map onto this model's weights — the
        # definitive "did it bind" signal. load_lora consumes matched keys and
        # returns the leftovers it could not place.
        key_map = comfy.lora.model_lora_keys_unet(model.model, {})
        loaded = comfy.lora.load_lora(lora_sd, key_map)
        matched = len(loaded)
        model_lora, _ = comfy.sd.load_lora_for_models(model, None, lora_sd, strength, 0)
        logging.info(
            f"[lora_from_url] {name}: matched {matched} target modules "
            f"from {total} tensors at strength {strength}"
            + ("  ⚠ ZERO MATCHES — lora did NOT bind" if matched == 0 else "")
        )
        return (model_lora,)


NODE_CLASS_MAPPINGS = {"LoraFromURLModelOnly": LoraFromURLModelOnly}
NODE_DISPLAY_NAME_MAPPINGS = {"LoraFromURLModelOnly": "Load LoRA from URL (model only)"}