##~ AutoCleaner V3.6 CODE | BY: ANXETY   ~##

# --change log--
"""
V3.6 | 13.03.24
Fixed selection window
Dynamic update of memory display
"""


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


#  ================= 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"),
        'SAGEMAKER_INTERNAL_IMAGE_URI': ('SageMaker Studio Lab', "/home/studio-lab-user/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"
#  ----------------------------------------------


directories = {
    "Images": f"{webui_path}/outputs",
    "Models": f"{webui_path}/models/Stable-diffusion/",
    "Vae": f"{webui_path}/models/VAE/",
    "LoRa": f"{webui_path}/models/Lora/",
    "ControlNet Models": f"{webui_path}/models/ControlNet/"
}


# ==================== CSS ====================
CSS = """
<style>
/* General Styles */

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

.instruction_AC {
    font-family: cursive;
    font-size: 18px;
    color: grey;
    user-select: none;
    cursor: default;
}


/* Container style */

.container_AC {
    position: relative;
    background-color: #232323;
    width: 800px;
    height: auto;
    padding: 15px;
    border-radius: 15px;
    box-shadow: 0 0 50px rgba(0, 0, 0, 0.3);
    margin-bottom: 5px;
    overflow: visible;
}

.container_AC::before {
    position: absolute;
    top: 5px;
    right: 10px;
    content: "AutoCleanerV3.6";
    font-weight: bold;
    font-size: 24px;
    color: rgba(0, 0, 0, 0.2);
}

.container_AC::after {
    position: absolute;
    top: 30px;
    right: 10px;
    content: "ANXETY";
    font-weight: bold;
    font-size: 18px;
    color: rgba(0, 0, 0, 0.2);
}

.custom-select-multiple_AC select {
    padding: 10px;
    font-family: cursive;
    border: 1px solid #262626 !important;
    border-radius: 10px;
    color: white;
    background-color: #1c1c1c;
    box-shadow: inset 0 0 5px rgba(0, 0, 0, 0.5);
}

.output_AC {
    padding: 10px;
    height: auto;
    border: 1px solid #262626;
    border-radius: 10px;
    background-color: #1c1c1c;
    box-shadow: inset 0 0 5px rgba(0, 0, 0, 0.5);
    width: auto;
    box-sizing: border-box;
}

.output_message_AC {
    font-family: cursive;
    color: white !important;
    font-size: 14px;
    user-select: none;
    cursor: default
}


.storage_info_AC {
    padding: 5px 20px;
    height: auto;
    border: 1px solid #262626;
    border-radius: 10px;
    background-color: #1c1c1c;
    box-shadow: inset 0 0 5px rgba(0, 0, 0, 0.5);
    width: auto;
    font-family: cursive;
    color: #B2B2B2 !important;
    font-size: 14px;
    user-select: none;
    cursor: default
}


/* Button and storage info layout */
.lower_information_panel_AC {
    display: flex;
    align-items: center;
    justify-content: space-between;
}


/* Button style */

.button_AC {
    width: auto;
    font-family: cursive;
    color: white !important;
    font-size: 14px;
    font-weight: bold;
    height: 35px;
    border-radius: 15px;
    background-image: radial-gradient(circle at top left, purple 10%, violet 90%);
    background-size: 200% 200%;
    background-position: left bottom;
    transition: background 0.5s ease-in-out, transform 0.3s ease;
}

.button_AC:hover {
    cursor: pointer;
    background-size: 200% 200%;
    background-position: right bottom;
    transform: translateY(1px);
}

.button_execute_AC:hover {
    background-image: radial-gradient(circle at top left, purple 10%, #93ac47 90%);
}

.button_clear_AC:hover {
    background-image: radial-gradient(circle at top left, purple 10%, #fc3468 90%);
}

.button_execute_AC:active,
.button_clear_AC:active {
    filter: brightness(0.75);
}

.jupyter-widgets.lm-Widget:focus {
    outline: none;
}


/* Animation of elements */

/* Emergence */
.container_AC {
    animation-name: slideInTopBlur;
    animation-duration: 0.7s;
    animation-fill-mode: forwards;
}

@keyframes slideInTopBlur {
    0% {
        transform: translate3d(0, 50%, 0) scale(0.85) rotate3d(1, 0, 0, -85deg);
        filter: blur(5px) grayscale(1) brightness(0.5);
        opacity: 0;
    }
    100% {
        transform: translate3d(0, 0, 0) scale(1) rotate3d(1, 0, 0, 0deg);
        filter: blur(0) grayscale(0) brightness(1);
        opacity: 1;
    }
}

/* Leaving */
.container_AC.hide {
    animation-name: slideOutTopBlur;
    animation-duration: 0.5s;
    animation-fill-mode: forwards;
}

@keyframes slideOutTopBlur {
    0% {
        transform: translate3d(0, 0, 0) scale(1);
        filter: blur(0) grayscale(0) brightness(1);
        opacity: 1;
    }
    100% {
        transform: translate3d(0, -100%, 0);
        filter: blur(5px) grayscale(1) brightness(0);
        opacity: 0;
    }
}
</style>
"""

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


# ================ AutoCleaner function ================
def clean_directory(directory):
    deleted_files = 0
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith(".txt"):
                continue
            os.remove(os.path.join(root, file))
            if not file.endswith(".yaml"):
                deleted_files += 1
    return deleted_files


def get_word_variant(n, variants):
    unit = abs(n) % 10
    tens = abs(n) % 100
    if tens in range(11, 15):
        return variants[2]
    elif unit == 1:
        return variants[0]
    elif unit in range(2, 5):
        return variants[1]
    else:
        return variants[2]


def update_memory_info():
    disk_space = psutil.disk_usage(os.getcwd())
    total = disk_space.total / (1024 ** 3)
    used = disk_space.used / (1024 ** 3)
    free = disk_space.free / (1024 ** 3)

    storage_info.value = f'''
    <div class="storage_info_AC">Total storage: {total:.2f} GB <span style="color: #555">|</span> Used: {used:.2f} GB <span style="color: #555">|</span> Free: {free:.2f} GB</div>
    '''


def on_execute_button_press(button):
    selected_cleaners = auto_cleaner_widget.value
    deleted_files_dict = {}

    for option in selected_cleaners:
        if option in directories:
            deleted_files_dict[option] = clean_directory(directories[option])

    output.clear_output()

    with output:
        for message in generate_messages(deleted_files_dict):
            message_widget = HTML(f'<p class="output_message_AC">{message}</p>')
            display(message_widget)

    update_memory_info()


def on_clear_button_press(button):
    container.add_class("hide")
    time.sleep(0.5)
    widgets.Widget.close_all()


def generate_messages(deleted_files_dict):
    messages = []
    word_variants = {
        "Images": ["Image", "Images", "Images"],
        "Models": ["Model", "Models", "Models"],
        "Vae": ["VAE", "VAE", "VAE"],
        "LoRa": ["LoRa", "LoRa", "LoRa"],
        "ControlNet Models": ["ControlNet Model", "ControlNet Models", "ControlNet Models"]
    }
    deleted_word_variants = ["Deleted", "Deleted", "Deleted"]
    for key, value in deleted_files_dict.items():
        word_variant = word_variants.get(key, ["", "", ""])
        deleted_word = get_word_variant(value, deleted_word_variants)
        object_word = get_word_variant(value, word_variant)
        messages.append(f"{deleted_word} {value} {object_word}")
    return messages
# ================ AutoCleaner function ================


# --- storage memory ---
import psutil
disk_space = psutil.disk_usage(os.getcwd())
total = disk_space.total / (1024 ** 3)
used = disk_space.used / (1024 ** 3)
free = disk_space.free / (1024 ** 3)


# UI Code
AutoCleaner_options = AutoCleaner_options = list(directories.keys())
instruction_label = widgets.HTML('''
<span class="instruction_AC">Use <span style="color: #B2B2B2;">ctrl</span> or <span style="color: #B2B2B2;">shift</span> for multiple selections.</span>
''')
auto_cleaner_widget = widgets.SelectMultiple(options=AutoCleaner_options, layout=widgets.Layout(width='auto')).add_class("custom-select-multiple_AC")
output = widgets.Output().add_class("output_AC")
# ---
execute_button = Button(description='Execute Cleaning').add_class("button_execute_AC").add_class("button_AC")
execute_button.on_click(on_execute_button_press)
clear_button = Button(description='Hide Widget').add_class("button_clear_AC").add_class("button_AC")
clear_button.on_click(on_clear_button_press)
# ---
storage_info = widgets.HTML(f'''
<div class="storage_info_AC">Total storage: {total:.2f} GB <span style="color: #555">|</span> Used: {used:.2f} GB <span style="color: #555">|</span> Free: {free:.2f} GB</div>
''')
# ---
buttons = widgets.HBox([execute_button, clear_button])
lower_information_panel = widgets.HBox([buttons, storage_info]).add_class("lower_information_panel_AC")

container = VBox([instruction_label, widgets.HTML('<hr>'), auto_cleaner_widget, output, widgets.HTML('<hr>'), lower_information_panel]).add_class("container_AC")

display(container)