##~ 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"


# ==================== 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;
}

/* 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 ====================
# Constants
HR = widgets.HTML('<hr>')
HEADER_DL = 'DOWNLOAD RESULTS'
VERSION = 'v0.31'

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

def get_files_list(directory, extensions):
    return [file for file in os.listdir(directory) if file.endswith(extensions)]

def get_folders_list(directory):
    return [folder for folder in os.listdir(extensions_dir) if os.path.isdir(os.path.join(extensions_dir, folder))]

def get_controlnets_list(directory, filter_pattern):
    filter_name = re.compile(filter_pattern)
    return [
        filter_name.match(file).group(1) if filter_name.match(file) else file
        for file in os.listdir(directory)
        if not file.endswith(('.txt', '.yaml'))
    ]

## Widgets
header_widget = widgets.HTML(value=f'''
<div><span class="header_dl">{HEADER_DL}</span> <span style="color: grey; opacity: 0.25;">| {VERSION}</span></div>
''')

# Models
models_list = get_files_list(models_dir, '.safetensors')
models_widget = output_container_generator('Models', models_list)
# Vaes
vaes_list = get_files_list(vaes_dir, '.safetensors')
vaes_widget = output_container_generator('VAEs', vaes_list)
# Embeddings
embeddings_list = get_files_list(embeddings_dir, ('.safetensors', '.pt'))
embeddings_widget = output_container_generator('Embeddings', embeddings_list)
# LoRAs
loras_list = get_files_list(loras_dir, '.safetensors')
loras_widget = output_container_generator('LoRAs', loras_list)
# Extensions
extensions_list = get_folders_list(extensions_dir)
extensions_widget = output_container_generator('Extensions', extensions_list).add_class("extension") # for fix height
# ControlNet
controlnets_list = get_controlnets_list(control_dir, r'^[^_]*_[^_]*_[^_]*_(.*)_fp16\.safetensors')
controlnets_widget = output_container_generator('ControlNets', controlnets_list)

## 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)