import os
import ipywidgets as widgets
from IPython.display import display
from subprocess import call

# Define the sd15 models and their URLs
sd15_models = {
    "ControlNet++ Canny": [
        "https://huggingface.co/huchenlei/ControlNet_plus_plus_collection_fp16/resolve/main/controlnet%2B%2B_canny_sd15_fp16.safetensors"
    ],
    "ControlNet++ Depth": [
        "https://huggingface.co/huchenlei/ControlNet_plus_plus_collection_fp16/resolve/main/controlnet%2B%2B_depth_sd15_fp16.safetensors"
    ],
    "ControlNet-v1-1 Inpaint": [
        "https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_inpaint_fp16.safetensors",
        "https://huggingface.co/ckpt/ControlNet-v1-1/raw/main/control_v11p_sd15_inpaint_fp16.yaml"
    ],
    "ControlNet++ Lineart": [
        "https://huggingface.co/huchenlei/ControlNet_plus_plus_collection_fp16/resolve/main/controlnet%2B%2B_lineart_sd15_fp16.safetensors"
    ],
    "ControlNet-v1-1 NormalBae": [
        "https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_normalbae_fp16.safetensors",
        "https://huggingface.co/ckpt/ControlNet-v1-1/raw/main/control_v11p_sd15_normalbae_fp16.yaml"
    ],
    "ControlNet-v1-1 Openpose": [
        "https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_openpose_fp16.safetensors",
        "https://huggingface.co/ckpt/ControlNet-v1-1/raw/main/control_v11p_sd15_openpose_fp16.yaml"
    ],
    "ControlNet-v1-1 Tile": [
        "https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11f1e_sd15_tile_fp16.safetensors",
        "https://huggingface.co/ckpt/ControlNet-v1-1/raw/main/control_v11f1e_sd15_tile_fp16.yaml"
    ],
    "ControlNet-v1-1 ip2p": [
        "https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11e_sd15_ip2p_fp16.safetensors",
        "https://huggingface.co/ckpt/ControlNet-v1-1/raw/main/control_v11e_sd15_ip2p_fp16.yaml"
    ],
    "ControlNet-v1-1 shuffle": [
        "https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11e_sd15_shuffle_fp16.safetensors",
        "https://huggingface.co/ckpt/ControlNet-v1-1/raw/main/control_v11e_sd15_shuffle_fp16.yaml"
    ],
    "ControlNet-v1-1 mlsd": [
        "https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_mlsd_fp16.safetensors",
        "https://huggingface.co/ckpt/ControlNet-v1-1/raw/main/control_v11p_sd15_mlsd_fp16.yaml"
    ],
    "ControlNet-v1-1 scribble": [
        "https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15_scribble_fp16.safetensors",
        "https://huggingface.co/ckpt/ControlNet-v1-1/raw/main/control_v11p_sd15_scribble_fp16.yaml"
    ],
    "ControlNet++ segment": [
        "https://huggingface.co/huchenlei/ControlNet_plus_plus_collection_fp16/resolve/main/controlnet%2B%2B_seg_sd15_fp16.safetensors"
    ],
    "ControlNet++ softedge": [
        "https://huggingface.co/huchenlei/ControlNet_plus_plus_collection_fp16/resolve/main/controlnet%2B%2B_hed_softedge_sd15_fp16.safetensors"
    ],
    "ControlNet-v1-1 lineart anime": [
        "https://huggingface.co/ckpt/ControlNet-v1-1/resolve/main/control_v11p_sd15s2_lineart_anime_fp16.safetensors",
        "https://huggingface.co/ckpt/ControlNet-v1-1/raw/main/control_v11p_sd15s2_lineart_anime_fp16.yaml"
    ],
    "ControlNet-v1-1 depth anything": [
        ("https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints_controlnet/diffusion_pytorch_model.safetensors", "control_sd15_depth_anything.safetensors")
    ],
    # IP-Adapter and FaceID models
    "IP-Adapter Full Face": [
        "https://huggingface.co/h94/IP-Adapter/resolve/main/models/ip-adapter-full-face_sd15.safetensors"
    ],
    "IP-Adapter Plus Face": [
        "https://huggingface.co/h94/IP-Adapter/resolve/main/models/ip-adapter-plus-face_sd15.safetensors"
    ],
    "IP-Adapter Plus": [
        "https://huggingface.co/h94/IP-Adapter/resolve/main/models/ip-adapter-plus_sd15.safetensors"
    ],
    "IP-Adapter": [
        "https://huggingface.co/h94/IP-Adapter/resolve/main/models/ip-adapter_sd15.safetensors"
    ],
    "ControlNet HandRefiner Inpaint Depth": [
        "https://huggingface.co/hr16/ControlNet-HandRefiner-pruned/resolve/main/control_sd15_inpaint_depth_hand_fp16.safetensors"
    ],
    "IP-Adapter FaceID Plus": [
        "https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main/ip-adapter-faceid-plus_sd15.bin"
    ],
    "IP-Adapter FaceID Plus V2": [
        "https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main/ip-adapter-faceid-plusv2_sd15.bin"
    ],
    "IP-Adapter FaceID Portrait": [
        "https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main/ip-adapter-faceid-portrait_sd15.bin"
    ],
}

# Define the xl models and their URLs
xl_models = {
    "SAI XL Canny": [
        "https://huggingface.co/lllyasviel/sd_control_collection/resolve/main/sai_xl_canny_256lora.safetensors"
    ],
    "SAI XL Depth": [
        "https://huggingface.co/lllyasviel/sd_control_collection/resolve/main/sai_xl_depth_256lora.safetensors"
    ],
    " XL Tile": [
        "https://huggingface.co/bdsqlsz/qinglong_controlnet-lllite/resolve/main/bdsqlsz_controlllite_xl_tile_realistic.safetensors"
    ],
    "MistoLine": [
        ("https://huggingface.co/TheMistoAI/MistoLine/resolve/main/mistoLine_rank256.safetensors", "mistoLine_softedge_sdxl_fp16.safetensors")
    ],
    "IP-Adapter Plus SDXL": [
        "https://huggingface.co/h94/IP-Adapter/resolve/main/sdxl_models/ip-adapter-plus_sdxl_vit-h.safetensors"
    ],
    "IP-Adapter Plus Face SDXL": [
        "https://huggingface.co/h94/IP-Adapter/resolve/main/sdxl_models/ip-adapter-plus-face_sdxl_vit-h.safetensors"
    ],
    "IP-Adapter FaceID Plus V2 SDXL": [
        "https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main/ip-adapter-faceid-plusv2_sdxl.bin"
    ],
     "Instant ID SDXL": [
        ("https://huggingface.co/InstantX/InstantID/resolve/main/ip-adapter.bin", "ip-adapter_instant_id_sdxl.bin"),
        ("https://huggingface.co/InstantX/InstantID/resolve/main/ControlNetModel/diffusion_pytorch_model.safetensors", "control_instant_id_sdxl.safetensors")
    ],
}

# Create checkboxes for each model
sd15_checkboxes = [widgets.Checkbox(value=False, description=key) for key in sd15_models.keys()]
xl_checkboxes = [widgets.Checkbox(value=False, description=key) for key in xl_models.keys()]

# "Select All" checkboxes
select_all_sd15 = widgets.Checkbox(description="Select All SD15", value=False)
select_all_xl = widgets.Checkbox(description="Select All XL", value=False)

# Function to update the selection state of all checkboxes based on "Select All" state
def update_selection(checkboxes, select_all_checkbox):
    for cb in checkboxes:
        cb.value = select_all_checkbox.value

# Attach the update_selection function to the "Select All" checkboxes
def attach_select_all(select_all_checkbox, checkboxes):
    def on_select_all_change(change):
        update_selection(checkboxes, select_all_checkbox)
    select_all_checkbox.observe(on_select_all_change, names='value')

attach_select_all(select_all_sd15, sd15_checkboxes)
attach_select_all(select_all_xl, xl_checkboxes)

# Function to display checkboxes for selected model type
checkbox_container = widgets.VBox(children=[])
def display_checkboxes(model_type):
    if model_type == 'sd15':
        checkbox_container.children = [select_all_sd15] + sd15_checkboxes
    elif model_type == 'xl':
        checkbox_container.children = [select_all_xl] + xl_checkboxes

# Create buttons for selecting model type
sd15_button = widgets.Button(description="SD15 Models")
xl_button = widgets.Button(description="XL Models")

# Button click events
sd15_button.on_click(lambda b: display_checkboxes('sd15'))
xl_button.on_click(lambda b: display_checkboxes('xl'))

# Display buttons
buttons_container = widgets.HBox(children=[sd15_button, xl_button])
display(buttons_container)

# Initially empty container for checkboxes
display(checkbox_container)

# Function to download selected models
def download_models(b):
    target_directory = "~/stable-diffusion-webui/extensions/sd-webui-controlnet/models"
    for checkbox in checkbox_container.children[1:]:  # Skip the "Select All" checkbox
        if checkbox.value:  # If the checkbox is checked
            model_name = checkbox.description
            model_files = sd15_models.get(model_name) or xl_models.get(model_name)
            for file_info in model_files:
                if isinstance(file_info, tuple):
                    url, filename = file_info
                    filepath = os.path.join(os.path.expanduser(target_directory), filename)
                else:
                    url = file_info
                    filename = url.split('/')[-1]
                    filepath = os.path.join(os.path.expanduser(target_directory), filename)
                
                # Check if the file already exists
                if not os.path.isfile(filepath):
                    # If the file does not exist, download it
                    if isinstance(file_info, tuple):
                        !wget {url} -O {filepath}
                    else:
                        !wget {url} -P {target_directory}
                else:
                    print(f"File {filename} already exists. Skipping download.")

# Create a button to start the download
download_button = widgets.Button(description="Download Selected Models")
download_button.on_click(download_models)
display(download_button)