##~   WIDGET CODE | BY: ANXETY   ~##

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 JS ====================
##~ custom background images V1.5 ~##
import argparse
parser = argparse.ArgumentParser(description='This script processes an background image.')
parser.add_argument('-i', '--image', type=str, help='URL of the image to process', metavar='')
parser.add_argument('-o', '--opacity', type=float, help='Opacity level for the image, between 0 and 1', metavar='', default=0.3)
parser.add_argument('-b', '--blur', type=str, help='Blur level for the image', metavar='', default=0)
parser.add_argument('-y', type=int, help='Y coordinate for the image in px', metavar='', default=0)
parser.add_argument('-x', type=int, help='X coordinate for the image in px', metavar='', default=0)
parser.add_argument('-s', '--scale', type=int, help='Scale image in %%', metavar='', default=100)
parser.add_argument('-m', '--mode', action='store_true', help='Removes repetitive image tiles')
parser.add_argument('-t', '--transparent', action='store_true', help='Makes input/selection fields 35%% more transparent')
parser.add_argument('-bf', '--blur-fields', type=str, help='Background blur level for input/selection fields', metavar='', default=2)

args = parser.parse_args()

url_img = args.image
opacity_img = args.opacity
blur_img = args.blur
y_img = args.y
x_img = args.x
scale_img = args.scale
blur_fields = args.blur_fields

## ---
""" WTF KAGGLE - WHAT THE FUCK IS THE DIFFERENCE OF 35 PIXELS!?!?!? """
fix_heigh_img = "-810px" if env == "Kaggle" else "-775px"

""" transperent fields """
t_bg_alpha = "1" if not args.transparent else "0.65"

""" mode img - repeats """
mode_img = "repeat" if not args.mode else "no-repeat"

container_background = f'''
<style>
:root {{
  /* for background container*/
  --img_background: url({url_img});
  --img_opacity: {opacity_img};
  --img_blur: {blur_img}px;
  --image_y: {y_img}px;
  --image_x: {x_img}px;
  --img_scale: {scale_img}%;
  --img_mode: {mode_img};
  --img_height_dif: {fix_heigh_img};

  /* for fields */
  --bg-field-color: rgba(28, 28, 28, {t_bg_alpha}); /* -> #1c1c1c */
  --bg-field-color-hover: rgba(38, 38, 38, {t_bg_alpha}); /* -> #262626; */
  --bg-field-blur-level: {blur_fields}px;
}}
'''

display(HTML(container_background))

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

# Main JS
JS = '''
<!-- TOGGLE 'CustomDL' SCRIPT -->
<script>
function toggleContainer() {
    let downloadContainer = document.querySelector('.container_custom_downlad');
    let info = document.querySelector('.info');

    downloadContainer.classList.toggle('expanded');
    info.classList.toggle('showed');
}
</script>
'''
display(HTML(JS))


# ==================== WIDGETS V2 ====================
HR = widgets.HTML('<hr>')

class WidgetFactory:
    def __init__(self, style=None, layout=None):
        self.style = style if style else {'description_width': 'initial'}
        self.layout = layout if layout else widgets.Layout(max_width='1080px', width='100%')

    def create_html(self, content, class_name=None):
        html_widget = widgets.HTML(content)
        if class_name:
            html_widget.add_class(class_name)
        return html_widget

    def create_header(self, name):
        return widgets.HTML(f'<div class="header">{name}<div>')

    def create_dropdown(self, options, value, description):
        return widgets.Dropdown(options=options, value=value, description=description, style=self.style, layout=self.layout)

    def create_text(self, description, placeholder='', value=''):
        return widgets.Text(description=description, placeholder=placeholder, value=value, style=self.style, layout=self.layout)

    def create_checkbox(self, value, description):
        return widgets.Checkbox(value=value, description=description, style=self.style, layout=self.layout)

    def create_button(self, description, class_name=None):
        button = widgets.Button(description=description)
        if class_name:
            button.add_class(class_name)
        return button

    def create_hbox(self, children):
        return widgets.HBox(children)

    def create_vbox(self, children, class_names=None):
        vbox = widgets.VBox(children)
        if class_names:
            for class_name in class_names:
                vbox.add_class(class_name)
        return vbox

    def display(self, widget):
        display(widget)

# Instantiate the factory
factory = WidgetFactory()

# --- MODEL ---
model_header = factory.create_header('Model Selection')
model_options = ['none',
                 '1.Anime (by XpucT) + INP',
                 '2.BluMix [Anime] [V7] + INP',
                 '3.Cetus-Mix [Anime] [V4] + INP',
                 '4.Counterfeit [Anime] [V3] + INP',
                 '5.CuteColor [Anime] [V3]',
                 '6.Dark-Sushi-Mix [Anime]',
                 '7.Deliberate [Realism] [V6] + INP',
                 '8.Meina-Mix [Anime] [V11] + INP',
                 '9.Mix-Pro [Anime] [V4] + INP']
model_widget = factory.create_dropdown(model_options, '4.Counterfeit [Anime] [V3] + INP', 'Model:')
model_num_widget = factory.create_text('Model Number:', 'Enter the model numbers to be downloaded using comma/space.')
inpainting_model_widget = factory.create_checkbox(False, 'Inpainting Models')

# Display Model
all_model_box = factory.create_vbox([model_header, model_widget, model_num_widget, inpainting_model_widget], class_names=["container", "image_1"])
factory.display(all_model_box)

# --- VAE ---
vae_header = factory.create_header('VAE Selection')
vae_options = ['none',
               '1.Anime.vae',
               '2.Anything.vae',
               '3.Blessed2.vae',
               '4.ClearVae.vae',
               '5.WD.vae']
vae_widget = factory.create_dropdown(vae_options, '3.Blessed2.vae', 'Vae:')
vae_num_widget = factory.create_text('Vae Number:', 'Enter the vae numbers to be downloaded using comma/space.')

# Display Vae
all_vae_box = factory.create_vbox([vae_header, vae_widget, vae_num_widget], class_names=["container", "image_2"])
factory.display(all_vae_box)

# --- ADDITIONAL ---
additional_header = factory.create_header('Additional')
latest_webui_widget = factory.create_checkbox(True, 'Update WebUI')
latest_exstensions_widget = factory.create_checkbox(True, 'Update Extensions')
change_webui_widget = factory.create_dropdown(['A1111', 'Forge'], 'A1111', 'Change WebUI:')
detailed_download_widget = factory.create_dropdown(['off', 'on'], 'off', 'Detailed Download:')
choose_changes_widget = factory.create_hbox([latest_webui_widget, latest_exstensions_widget, change_webui_widget, detailed_download_widget])

controlnet_options = ['none', 'ALL', '1.canny',
                      '2.openpose', '3.depth',
                      '4.normal_map', '5.mlsd',
                      '6.lineart', '7.soft_edge',
                      '8.scribble', '9.segmentation',
                      '10.shuffle', '11.tile',
                      '12.inpaint', '13.instruct_p2p']
controlnet_widget = factory.create_dropdown(controlnet_options, 'none', 'ControlNet:')
controlnet_num_widget = factory.create_text('ControlNet Number:', 'Enter the ControlNet model numbers to be downloaded using comma/space.')
commit_hash_widget = factory.create_text('Commit Hash:')
huggingface_token_widget = factory.create_text('HuggingFace Token:')

ngrok_token_widget = factory.create_text('Ngrok Token:')
ngrock_button = factory.create_html('<a href="https://dashboard.ngrok.com/get-started/your-authtoken" target="_blank">Получить Ngrok Токен</a>', class_name="button_ngrok")
ngrok_widget = factory.create_hbox([ngrok_token_widget, ngrock_button])

zrok_token_widget = factory.create_text('Zrok Token:')
zrok_button = factory.create_html('<a href="https://colab.research.google.com/drive/1d2sjWDJi_GYBUavrHSuQyHTDuLy36WpU" target="_blank">Зарегать Zrok Токен</a>', class_name="button_ngrok")
zrok_widget = factory.create_hbox([zrok_token_widget, zrok_button])

commandline_arguments_options = "--listen --enable-insecure-extension-access --theme dark --no-half-vae --disable-console-progressbars --xformers"
commandline_arguments_widget = factory.create_text('Arguments:', value=commandline_arguments_options)

# Display Additional
additional_widget_list = [additional_header,
                          choose_changes_widget,
                          HR,
                          controlnet_widget,
                          controlnet_num_widget,
                          commit_hash_widget,
                          huggingface_token_widget,
                          ngrok_widget,
                          zrok_widget,
                          HR,
                          commandline_arguments_widget]

if free_plan and env == "Google Colab": # remove ngrok from colab
    additional_widget_list.remove(ngrok_widget)

all_additional_box = factory.create_vbox(additional_widget_list, class_names=["container", "image_3"])
factory.display(all_additional_box)

# --- CUSTOM DOWNLOAD ---
custom_download_header_popup = factory.create_html('''
<style>
/* Term Colors */
.sample_label {color: #dbafff;}
.braces {color: #ffff00;}
.extension {color: #eb934b;}
.file_name {color: #ffffd8;}
</style>

<div class="header" style="cursor: pointer;" onclick="toggleContainer()">Custom Download</div>
<!-- PopUp Window -->
<div class="info">INFO</div>
<div class="popup">
    Separate multiple URLs with a comma/space. For a <span class="file_name">custom name</span> file/extension, specify it with <span class="braces">[]</span>
    after the URL without spaces.
    <span style="color: #ff9999">For files, be sure to specify</span> - <span class="extension">Filename Extension.</span>
    <div class="sample">
        <span class="sample_label">Example for File:</span>
        https://civitai.com/api/download/models/229782<span class="braces">[</span><span class="file_name">Detailer</span><span class="extension">.safetensors</span><span class="braces">]</span>
        <br>
        <span class="sample_label">Example for Extension:</span>
        https://github.com/hako-mikan/sd-webui-regional-prompter<span class="braces">[</span><span class="file_name">Regional-Prompter</span><span class="braces">]</span>
    </div>
</div>
''')

Model_url_widget = factory.create_text('Model:')
Vae_url_widget = factory.create_text('Vae:')
LoRA_url_widget = factory.create_text('LoRa:')
Embedding_url_widget = factory.create_text('Embedding:')
Extensions_url_widget = factory.create_text('Extensions:')
custom_file_urls_widget = factory.create_text('File (txt):')

# Display CustomDl
all_custom_box = factory.create_vbox([
    custom_download_header_popup, Model_url_widget, Vae_url_widget, LoRA_url_widget, Embedding_url_widget, Extensions_url_widget, custom_file_urls_widget
], class_names=["container", "image_4", "container_custom_downlad"])
factory.display(all_custom_box)

# --- Save Button ---
save_button = factory.create_button('Save', class_name="button_save")
factory.display(save_button)


# ==================== OTHER FUNC ====================
# Setup UI
def setup_webui(change_webui_widget):
    if 'SDW_UI' in os.environ:
        UI = os.environ['SDW_UI']
    else:
        os.environ['SDW_UI'] = change_webui_widget
        os.environ['SDW_OLD_UI'] = change_webui_widget

    UI = os.getenv('SDW_UI')

    if UI != change_webui_widget:
        os.environ['SDW_UI'] = change_webui_widget


# ============ Load / Save - Settings V2 ============
SETTINGS_FILE = f'{root_path}/settings.json'

SETTINGS_KEYS = [
    'model', 'model_num', 'inpainting_model',
    'vae', 'vae_num', 'latest_webui', 'latest_exstensions',
    'change_webui', 'detailed_download', 'controlnet',
    'controlnet_num', 'commit_hash', 'huggingface_token',
    'ngrok_token', 'zrok_token', 'commandline_arguments',
    'Model_url', 'Vae_url', 'LoRA_url', 'Embedding_url',
    'Extensions_url', 'custom_file_urls'
]

def save_settings():
    settings = {key: globals()[f"{key}_widget"].value for key in SETTINGS_KEYS}
    with open(SETTINGS_FILE, 'w') as f:
        json.dump(settings, f, indent=2)

def load_settings():
    if os.path.exists(SETTINGS_FILE):
        with open(SETTINGS_FILE, 'r') as f:
            settings = json.load(f)
            for key in SETTINGS_KEYS:
                globals()[f"{key}_widget"].value = settings.get(key, "")

def hide_widgets():
    widgets_list = [all_model_box, all_vae_box, all_additional_box, all_custom_box, save_button]
    for widget in widgets_list:
        widget.add_class("hide")
    time.sleep(0.5)
    widgets.Widget.close_all()

def save_data(button):
    save_settings()
    setup_webui(change_webui_widget.value)
    hide_widgets()

load_settings()
save_button.on_click(save_data)