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

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


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


# ==================== CSS ====================
# Main CSS
css_file_path = f"{root_path}/CSS/auto_cleaner.css"
with open(css_file_path , "r") as f:
    CSS_AC = f.read()
display(HTML(f"<style>{CSS_AC}</style>"))


# ================ AutoCleaner function ================
directories = {
    "Изображения": f"{webui_path}/output",
    "Модели": f"{webui_path}/models/Stable-diffusion/",
    "Vae": f"{webui_path}/models/VAE/",
    "LoRa": f"{webui_path}/models/Lora/",
    "ControlNet Модели": f"{webui_path}/models/ControlNet/"
}

""" functions """
def clean_directory(directory):
    deleted_files = 0
    image_dir = directories['Изображения']

    for root, dirs, files in os.walk(directory):
        for file in files:
            file_path = os.path.join(root, file)

            if file.endswith(".txt"):
                continue
            if file.endswith((".safetensors", ".pt")) or root == image_dir: # fix for image counter
                deleted_files += 1

            os.remove(file_path)
    return deleted_files

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:.2f} GB <span style="color: #555">|</span> Используется: {used:.2f} GB <span style="color: #555">|</span> Свободно: {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 = {
        "Изображения": "Изображений",
        "Модели": "Моделей",
        "Vae": "Vae",
        "LoRa": "LoRa",
        "ControlNet Модели": "ControlNet Моделей"
    }
    for key, value in deleted_files_dict.items():
        object_word = word_variants.get(key)
        messages.append(f"Удалено {value} {object_word}")
    return messages


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


# ================ Widgets ================
# UI Code
AutoCleaner_options = AutoCleaner_options = list(directories.keys())
instruction_label = widgets.HTML('''
<span class="instruction_AC">Используйте <span style="color: #B2B2B2;">ctrl</span> или <span style="color: #B2B2B2;">shift</span> для множественного выбора.</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='Выполнить Очистку').add_class("button_execute_AC").add_class("button_AC")
execute_button.on_click(on_execute_button_press)
clear_button = Button(description='Скрыть Виджет').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:.2f} GB <span style="color: #555">|</span> Используется: {used:.2f} GB <span style="color: #555">|</span> Свободно: {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)