from IPython.display import display, HTML, clear_output, Image
from IPython import get_ipython
from ipywidgets import widgets
from pathlib import Path
import subprocess
import shlex
import json

home = Path.home()
src = home / '.gutris1'
css = src / 'segsmaker.css'
startup = home / '.ipython/profile_default/startup'
nenen = startup / 'nenen88.py'
key_file = src / 'api-key.json'
img = src / 'loading.png'

R = '\033[0m'
T = f'▶{R}'
BLUE = f'\033[38;5;33m{T}'
CYAN = f'\033[36m{T}'
PURPLE = f'\033[38;5;135m{T}'
PINK = f'\033[38;5;201m{T}'
RED = f'\033[31m{T}'
GREEN = f'\033[38;5;35m{T}'

Path(src).mkdir(parents=True, exist_ok=True)

SyS = get_ipython().system

main_output = widgets.Output()
save_button = widgets.Button(description='Save')
civitai_key_box = widgets.Text(placeholder='Enter Your Civitai API Key Here', layout=widgets.Layout(width='350px'))
hf_token_box = widgets.Text(placeholder='Huggingface READ Token (optional)', layout=widgets.Layout(width='350px'))
input_widget = widgets.Box(
    [civitai_key_box, hf_token_box, save_button], 
    layout=widgets.Layout(
        width='500px',
        height='200px',
        display='flex',
        flex_flow='column',
        align_items='center',
        justify_content='space-around',
        padding='10px'
    )
)

save_button.add_class('save-button')
civitai_key_box.add_class('api-input')
hf_token_box.add_class('api-input')
input_widget.add_class('boxs')

def CondaInstall():
    try:
        display(Image(filename=str(img)))

        cmd_list = [
            (f'rm -rf {home}/.condarc', None),
            ('conda config --add channels conda-forge', None),
            ('conda install -qy mamba', f'{BLUE} Installing Anaconda'),
            ('mamba install -y conda', None),
            ('mamba install -y python=3.10.13', f'{CYAN} Installing Python 3.10'),
            ('mamba install -y glib gperftools openssh pv gputil curl', f'{PURPLE} Installing Conda Packages'),
            ('pip install psutil aria2 gdown', f'{PINK} Installing Python Packages'),
            ('conda clean -qy --all', None)
        ]

        for cmd, msg in cmd_list:
            if msg is not None: print(msg)
            subprocess.run(shlex.split(cmd), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

        SyS(f'rm -rf {home}/.cache/* {home}/.condarc')
        clear_output()
        print(f'{GREEN} Done')

        get_ipython().kernel.do_shutdown(True)

    except KeyboardInterrupt:
        clear_output()
        print('^ Canceled')

def LoadCSS(): 
    display(HTML(f'<style>{open(css).read()}</style>'))

def KeyInject(civitai_key, hf_token):
    SyS(f'curl -sLo {nenen} https://github.com/gutris1/segsmaker/raw/main/script/nenen88.py')

    p = Path(nenen)
    v = p.read_text()
    v = v.replace("TOKET = ''", f"TOKET = '{civitai_key}'")
    v = v.replace("TOBRUT = ''", f"TOBRUT = '{hf_token}'")
    p.write_text(v)

def KeyWidget(civitai_key='', hf_token=''):
    civitai_key_box.value = civitai_key
    hf_token_box.value = hf_token

    def KeyInputs(b):
        civitai_key = civitai_key_box.value.strip()
        hf_token = hf_token_box.value.strip()

        with main_output:
            if not civitai_key:
                print('Please enter your Civitai API Key')
                return

            if len(civitai_key) < 32:
                print('API key must be at least 32 characters long')
                return

            civitai_key_value = {'civitai-api-key': civitai_key}
            hf_token_value = {'huggingface-read-token': hf_token}

            secrets = {**civitai_key_value, **hf_token_value}
            key_file.write_text(json.dumps(secrets, indent=4))

        KeyInject(civitai_key, hf_token)

        input_widget.close()
        main_output.clear_output()

        p = subprocess.run(['conda', '--version'], capture_output=True, text=True, check=True)
        cv = p.stdout.strip().split()[1]
        mv = int(cv.split('.')[0])

        with main_output:
            CondaInstall()
    save_button.on_click(KeyInputs)

def KeyCheck():
    if key_file.exists():
        v = json.loads(key_file.read_text())

        civitai_key = v.get('civitai-api-key', '')
        hf_token = v.get('huggingface-read-token', '')

        if not civitai_key or not hf_token:
            display(input_widget, main_output)
            KeyWidget(civitai_key, hf_token)
        else:
            KeyInject(civitai_key, hf_token)
            display(main_output)
            CondaInstall()
    else:
        display(input_widget, main_output)
        KeyWidget()

def CondaMisc():
    for scr in [
        f'curl -sLo {css} https://github.com/gutris1/segsmaker/raw/main/script/SM/segsmaker.css',
        f'curl -sLo {startup}/00-startup.py https://github.com/gutris1/segsmaker/raw/main/script/SM/00-startup.py',
        f'curl -sLo {startup}/util.py https://github.com/gutris1/segsmaker/raw/main/script/SM/util.py',
        f'curl -sLo {img} https://github.com/gutris1/segsmaker/raw/main/script/loading.png',
        f'curl -sLo {startup}/cupang.py https://github.com/gutris1/segsmaker/raw/main/script/cupang.py'
    ]: SyS(scr)

CondaMisc()
LoadCSS()
KeyCheck()