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

from directory_setup import *

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


# 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/dl_display_result.css"
with open(css_file_path , "r") as f:
    CSS = f.read()
display(HTML(f"<style>{CSS}</style>"))


# ==================== WIDGETS ====================
# Constants
HR = widgets.HTML('<hr>')
HEADER_DL = 'DOWNLOAD RESULTS'
VERSION = 'v0.33'

## 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', '.ckpt'))
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)