from IPython.display import display, HTML, clear_output, Image
from ipywidgets import widgets, Layout
from IPython import get_ipython
from pathlib import Path
import subprocess, time, os, shlex
from nenen88 import pull, say, download, clone, tempe

version = "v1.9.4"
repo = f"git clone -q -b {version} https://github.com/gutris1/asd"

home = Path.home()
webui = home / "asd"
img = home / ".conda/loading.png"

os.chdir(home)

if webui.exists():
    git_dir = webui / '.git'
    if git_dir.exists():
        os.chdir(webui)
        commit_hash = os.popen('git rev-parse HEAD').read().strip()

        if commit_hash != version:
            get_ipython().system(f"git pull origin {version}")
            get_ipython().system("git fetch --tags")

else:
    css = home / ".conda/setup.css"
    devnull = {"stdout": subprocess.DEVNULL, "stderr": subprocess.DEVNULL}

    get_ipython().system(f"curl -sLo {css} https://github.com/gutris1/segsmaker/raw/main/ui/sd/asd/setup.css")

    loading = widgets.Output()
    button1 = widgets.Button(description='SD 1.5')
    install_button = widgets.Button(description='Install')

    button_button = widgets.HBox([button1], layout=Layout(justify_content='space-between'))
    sd_setup = widgets.Output()
    panel = widgets.VBox([button_button, install_button],
                         layout=Layout(
                             width='400px',
                             height='150px',
                             display='flex',
                             flex_flow='column',
                             align_items='center',
                             justify_content='space-between',
                             padding='20px'))
    button1.add_class("b1")
    install_button.add_class("out")
    panel.add_class("boxxx")
    selected = [None]

    def load_css(css):
        with css.open("r") as file:
            ccs = file.read()

        display(HTML(f"<style>{ccs}</style>"))

    def req_list(home, webui):
        return [
            f"rm -rf {home}/tmp/* {home}/tmp {webui}/models/Stable-diffusion/tmp_ckpt {webui}/models/Lora/tmp_lora {webui}/models/ControlNet",
            f"mkdir -p {webui}/models/Lora",
            f"mkdir -p {webui}/models/ESRGAN",
            f"ln -vs /tmp {home}/tmp",
            f"ln -vs /tmp/ckpt {webui}/models/Stable-diffusion/tmp_ckpt",
            f"ln -vs /tmp/lora {webui}/models/Lora/tmp_lora",
            f"ln -vs /tmp/controlnet {webui}/models/ControlNet"]

    def sd_clone(home, webui, devnull):
        time.sleep(1)
        pull(f"https://github.com/gutris1/segsmaker sd {webui}")

        os.chdir(webui)
        req = req_list(home, webui)

        for lines in req:
            subprocess.run(shlex.split(lines), **devnull)
            
        scripts = [
            f"https://github.com/gutris1/segsmaker/raw/main/script/zrok_reg.py {webui}/asd",
            f"https://github.com/gutris1/segsmaker/raw/main/script/zrok.py {webui}",
            f"https://github.com/gutris1/segsmaker/raw/main/script/pinggy.py {webui}",
            f"https://github.com/gutris1/segsmaker/raw/main/script/ngrok.py {webui}",
            f"https://github.com/gutris1/segsmaker/raw/main/script/venv.py {webui}"]
            
        upscalers = [
            f"https://huggingface.co/pantat88/ui/resolve/main/4x-UltraSharp.pth {webui}/models/ESRGAN",
            f"https://huggingface.co/pantat88/ui/resolve/main/4x-AnimeSharp.pth {webui}/models/ESRGAN",
            f"https://huggingface.co/pantat88/ui/resolve/main/4x_NMKD-Superscale-SP_178000_G.pth {webui}/models/ESRGAN",
            f"https://huggingface.co/pantat88/ui/resolve/main/4x_RealisticRescaler_100000_G.pth {webui}/models/ESRGAN",
            f"https://huggingface.co/pantat88/ui/resolve/main/8x_RealESRGAN.pth {webui}/models/ESRGAN",
            f"https://huggingface.co/pantat88/ui/resolve/main/4x_foolhardy_Remacri.pth {webui}/models/ESRGAN"]
        
        line = scripts + upscalers
        for item in line:
            download(item)

    def sd_1_5(home, webui, devnull):
        sd_clone(home, webui, devnull)

        extras = [
            f"https://github.com/gutris1/segsmaker/raw/main/script/controlnet/cn-1_5.css {webui}/asd",
            f"https://github.com/gutris1/segsmaker/raw/main/script/controlnet/cn-1_5.py {webui}/asd",
            f"https://huggingface.co/Carlos2312/asd/resolve/main/modelos.py {webui}/asd",
            f"https://huggingface.co/Carlos2312/asd/resolve/main/ext-1_5.txt {webui}/asd",
            f"https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors {webui}/models/VAE"]

        for items in extras:
            download(items)

        say("<br><b>【{red} Installing Extensions{d} 】{red}</b>")
        os.chdir(webui / "extensions")
        clone(str(webui / "asd/ext-1_5.txt"))
        os.rename(str(webui / "asd/cn-1_5.py"), str(webui / "asd/controlnet.py"))
        
        get_ipython().system(f"rm -rf {webui}/embeddings")
        get_ipython().system(f"git clone https://huggingface.co/Carlos2312/emb {webui}/embeddings")

        get_ipython().system(f"rm -rf {webui}/models/Lora")
        get_ipython().system(f"git clone https://huggingface.co/Carlos2312/lrs {webui}/models/Lora")

        get_ipython().system(f"rm -rf {webui}/models/adetailer")
        get_ipython().system(f"git clone https://huggingface.co/Carlos2312/md {webui}/models/adetailer")

        get_ipython().system(f"rm -rf {webui}/extensions/sd-dynamic-prompts/wildcards")
        get_ipython().system(f"git clone https://huggingface.co/Carlos2312/Wilcards {webui}/extensions/sd-dynamic-prompts/wildcards")
        
    def sd_install(selection):
        with loading:
            clear_output()
            display(Image(filename=str(img)))

        with sd_setup:
            sd_setup.clear_output()
            say("<b>【{red} Installing Stable Diffusion{d} 】{red}</b>")
            get_ipython().system(f"{repo}")

            if selection == 'SD 1.5':
                sd_1_5(home, webui, devnull)
                
            with loading:
                loading.clear_output()
                
            say("<br><b>【{red} Done{d} 】{red}</b>")

    def button_panel(button):
        selected[0] = button.description

    def installing(button):
        selection = selected[0]

        if selection:
            widgets.Widget.close(panel)
            sd_setup.clear_output()
            sd_install(selection)

        else:
            with sd_setup:
                print("at least make a choice")
                print("少なくとも選択してよ。")
            return

    button1.on_click(button_panel)
    install_button.on_click(installing)

    load_css(css)
    display(panel, sd_setup, loading)
    tempe()