import os
import hashlib
import requests
import folder_paths
import comfy.utils
import comfy.sd
from tqdm import tqdm

class LoadLoraFromURL:
    """Load a LoRA model from a URL"""
    
    def __init__(self):
        self.cache_dir = os.path.join(folder_paths.get_input_directory(), "url_loras")
        os.makedirs(self.cache_dir, exist_ok=True)

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

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

    def download_if_needed(self, url):
        """Download the file if not in cache"""
        filename = hashlib.md5(url.encode()).hexdigest() + ".safetensors"
        local_path = os.path.join(self.cache_dir, filename)
        
        if os.path.exists(local_path):
            return local_path
            
        print(f"Downloading LoRA from {url}")
        response = requests.get(url, stream=True)
        response.raise_for_status()
        
        total_size = int(response.headers.get('content-length', 0))
        
        with open(local_path, 'wb') as f, tqdm(
            desc=filename,
            total=total_size,
            unit='iB',
            unit_scale=True
        ) as pbar:
            for data in response.iter_content(chunk_size=8192):
                size = f.write(data)
                pbar.update(size)
                
        return local_path

    def load_lora(self, url, model, strength):
        try:
            # Download or get cached file
            lora_path = self.download_if_needed(url)
            
            # Load the LoRA using ComfyUI's built-in functions
            lora = comfy.utils.load_torch_file(lora_path)
            model_lora, _ = comfy.sd.load_lora_for_models(
                model, None, lora, strength, 0
            )
            return (model_lora,)
            
        except Exception as e:
            print(f"Error loading LoRA: {str(e)}")
            return (model,)