# custom_nodes/download_lora_node/__init__.py

import os
import shutil
import tarfile
import tempfile
import requests

class DownloadLoRANode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "lora_url": ("STRING", {"default": "https://example.com/flux-lora.tar"}),
                "lora_name": ("STRING", {"default": "my_lora"}),
                "loras_dir": ("STRING", {"default": "./ComfyUI/models/loras"}),
            }
        }

    RETURN_TYPES = ("STRING",)  # путь к safetensors
    FUNCTION = "download_lora"
    CATEGORY = "Utils/LoRA"
    NODE_NAME = "Download LoRA (.tar URL)"

    def download_lora(self, lora_url, lora_name, loras_dir):
        os.makedirs(loras_dir, exist_ok=True)

        with tempfile.TemporaryDirectory() as tmpdir:
            tar_path = os.path.join(tmpdir, f"{lora_name}.tar")

            # Скачиваем архив
            with requests.get(lora_url, stream=True, timeout=300) as r:
                r.raise_for_status()
                with open(tar_path, "wb") as f:
                    for chunk in r.iter_content(1 << 20):
                        f.write(chunk)

            # Распаковываем
            with tarfile.open(tar_path, "r") as t:
                t.extractall(tmpdir)

            # ищем .safetensors внутри архива
            safetensor_file = None
            for root, _, files in os.walk(tmpdir):
                for fn in files:
                    if fn.endswith(".safetensors"):
                        safetensor_file = os.path.join(root, fn)
                        break
                if safetensor_file:
                    break

            if not safetensor_file:
                raise RuntimeError("❌ Не найден .safetensors в архиве")

            dst = os.path.join(loras_dir, f"{lora_name}.safetensors")
            shutil.move(safetensor_file, dst)

        return (dst,)


NODE_CLASS_MAPPINGS = {
    "Download LoRA (.tar URL)": DownloadLoRANode
}