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

from directory_setup import models_dir, vaes_dir, control_dir, loras_dir, output_dir

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


# Setup Env
env = os.getenv('ENV_NAME')
root_path = os.getenv('ROOT_PATH')
webui_path = os.getenv('WEBUI_PATH')
free_plan = os.getenv('FREE_PLAN')


# ==================== 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 = {
    "Images": output_dir,
    "Models": models_dir,
    "Vae": vaes_dir,
    "LoRa": loras_dir,
    "ControlNet Models": control_dir
}

""" functions """
def clean_directory(directory):
    deleted_files = 0
    image_dir = directories['Images']

    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 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": "Images",
        "Models": "Models",
        "Vae": "Vae",
        "LoRa": "LoRa",
        "ControlNet Models": "ControlNet Models"
    }
    for key, value in deleted_files_dict.items():
        object_word = word_variants.get(key)
        messages.append(f"Deleted {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)