##~ Display Download Results Widgets | by: ANXETY ~##

import re
import os
import json
import time
import ipywidgets as widgets
from ipywidgets import widgets, Layout, Label, Button, VBox, HBox
from IPython.display import display, HTML, Javascript, clear_output


#  ================= DETECT ENV =================
def detect_environment():
    free_plan = (os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / (1024. ** 3) <= 20)
    environments = {
        'COLAB_GPU': ('Google Colab', "/root" if free_plan else "/content"),
        'KAGGLE_URL_BASE': ('Kaggle', "/kaggle/working/content")
    }

    for env_var, (environment, path) in environments.items():
        if env_var in os.environ:
            return environment, path, free_plan

env, root_path, free_plan = detect_environment()
webui_path = f"{root_path}/sdw"
#  ----------------------------------------------


# CONFIG DIR
models_dir = f"{webui_path}/models/Stable-diffusion"
vaes_dir = f"{webui_path}/models/VAE"
embeddings_dir = f"{webui_path}/embeddings"
loras_dir = f"{webui_path}/models/Lora"
extensions_dir = f"{webui_path}/extensions"
control_dir = f"{webui_path}/models/ControlNet"
adetailer_dir = f"{webui_path}/models/adetailer"


# ==================== CSS ====================
CSS = """
<style>
/* General Styles */
.header_dl,
.header_outputs_dl {
    font-family: cursive;
    font-size: 20px;
    font-weight: bold;
    text-align: center;
}
.header_dl {
    color: #AC8FA5;
    margin-bottom: 15px;
}
.header_outputs_dl {
    color: #0083C0;
}

hr {
    border-color: grey;
    background-color: grey;
    opacity: 0.25;
    width: 1000px;
}

/* Element text style */
.widget-html {
    font-family: cursive;
    font-size: 14px;
    color: white !important;
    user-select: none;
}


/* Container style */
.container_dl {
    position: relative;
    flex-direction: column;
    align-items: center;
    background-color: #232323;
    width: 1200px;
    height: auto;
    margin: 40px 10px 10px 10px;
    padding: 10px 15px;
    border-radius: 15px;
    box-shadow: 0 0 50px rgba(0, 0, 0, 0.3), inset 0 0 10px rgba(0, 0, 0, 0.3);
    transition: all 0.5s ease-in-out;
    overflow: visible;
}
.container_dl::after {
    position: absolute;
    top: 5px;
    right: 10px;
    content: "ANXETY";
    font-weight: bold;
    font-size: 24px;
    color: rgba(0, 0, 0, 0.2);
}

.result_output_dl {
    flex-wrap: wrap;
    box-shadow: inset 0 0 10px rgba(0, 0, 0, 0.3);
    border-radius: 8px;
    background-color: #1F1F1F;
    width: 95%;
    padding: 10px 15px;
    overflow: visible;
}
.outputs_dl {
    flex-grow: 1;
    flex-wrap: wrap;
    background-color: #181818;
    padding: 10px 15px;
    border-radius: 15px;
    box-shadow: 0 0 8px rgba(0, 0, 0, 0.3), inset 0 0 10px rgba(0, 0, 0, 0.3);
    color: #C1C1C1;
    margin: 8px;
    transition: all 0.15s ease-in-out;
}

/* setting the height of "outputs" containers */
.outputs_dl { max-height: 400px; }

/* good use of space */
.extension {
  display: grid;
  grid-template-columns: repeat(3, 1fr);
}


/* Animation of elements */

.container_dl,
.outputs_dl {
    animation: showedResult 1s;
}
.items_dl {
    animation: showedText 1s;
}

@keyframes showedResult {
    0% {
        transform: translate3d(0, 15%, 0);
        opacity: 0;
    }
    100% {
        transform: translate3d(0, 0, 0);
        opacity: 1;
    }
}

@keyframes showedText {
    0% {
        transform: translate3d(-30%, 0, 0);
        opacity: 0;
    }
    100% {
        transform: translate3d(0, 0, 0);
        opacity: 1;
    }
}
</style>
"""

display(HTML(CSS))
# ==================== CSS ====================


# ==================== WIDGETS ====================
# --- global widgets ---
HR = widgets.HTML('<hr>')

header_widget = widgets.HTML(value='''
<div class="header_dl">DOWNLOAD RESULTS <span style="color: grey; opacity: 0.25;">| v0.2.8</span></div>
''')

# Adding content to "outputs_widgets_list"
def output_container_generator(header, items):
    header_widget = widgets.HTML(value=f'<div class="header_outputs_dl">{header} ➤</div>')
    content = [widgets.HTML(value=f'<div class="items_dl">{item}</div>') for item in items]
    container_widget = widgets.VBox([header_widget] + content).add_class("outputs_dl")
    return container_widget

# Models
models_list = [file for file in os.listdir(models_dir) if file.endswith('.safetensors')]
models_widget = output_container_generator('Models', models_list).add_class("model")
# Vaes
vaes_list = [file for file in os.listdir(vaes_dir) if file.endswith('.safetensors')]
vaes_widget = output_container_generator('VAEs', vaes_list).add_class("vae")
# Embeddings
embeddings_list = [file for file in os.listdir(embeddings_dir) if file.endswith(('.safetensors', '.pt')) and not os.path.isdir(os.path.join(embeddings_dir, file))]
embeddings_widget = output_container_generator('Embeddings', embeddings_list).add_class("embedding")
# LoRAs
# loras_list = [file for file in os.listdir(loras_dir) if not file.endswith('.txt')]
loras_list = [file for file in os.listdir(loras_dir) if file.endswith('.safetensors')]
loras_widget = output_container_generator('LoRAs', loras_list).add_class("lora")
# Extensions
extensions_list = [folder for folder in os.listdir(extensions_dir) if os.path.isdir(os.path.join(extensions_dir, folder))]
extensions_widget = output_container_generator('Extensions', extensions_list).add_class("extension")
# ControlNet
# controlnets_list = [file for file in os.listdir(control_dir) if not file.endswith(('.txt', '.yaml'))]
filter_name = re.compile(r'^[^_]*_[^_]*_[^_]*_(.*)_fp16\.safetensors$')
controlnets_list = [
    filter_name.match(file).group(1) if filter_name.match(file) else file
    for file in os.listdir(control_dir)
    if not file.endswith(('.txt', '.yaml'))
]
controlnets_widget = output_container_generator('ControlNets', controlnets_list).add_class("controlnet")


## Sorting and Output
widgets_dict = {
    models_widget: models_list,
    vaes_widget: vaes_list,
    embeddings_widget: embeddings_list,
    loras_widget: loras_list,
    extensions_widget: extensions_list,
    controlnets_widget: controlnets_list
}
outputs_widgets_list = [widget for widget, widget_list in widgets_dict.items() if widget_list]
result_output_widget = widgets.HBox(outputs_widgets_list).add_class("result_output_dl")

container_widget = widgets.VBox([header_widget, HR, result_output_widget, HR]).add_class("container_dl")

display(container_widget)