# -*- coding: utf-8 -*-
import datetime
import os

# print(f"pip3 uninstall gradio")
# results = os.popen(f"pip3 uninstall gradio -y")
# print(results.readlines())
results = os.popen(f"ls ~/.cache/huggingface/hub")
print("step0:", results.readlines())
results = os.popen(f"pip3 install -r requirements.txt -i https://pypi.org/simple")
print("step1:", results.readlines())
results = os.popen(f"pip3 uninstall scepter -y")
print("step2:", results.readlines())
scepter_whl = os.environ.get("SCEPTER_WHL","scepter")
results = os.popen(f"pip3 install {scepter_whl} -i https://pypi.org/simple")
print("step3:", scepter_whl, results.readlines())
results = os.popen(f"pip3 list")
print("step4:", results.readlines())
#results = os.popen(f"sudo chmod 777 /data")
#print("step5:", results.readlines())
print("finish install")

import random
import gradio as gr

import scepter
from scepter.modules.utils.config import Config
from scepter.modules.utils.file_system import FS
from scepter.modules.utils.logger import get_logger, init_logger

def link_model():
    import os
    ms_root_dir = "cache/cache_data/"
    hf_root_dir = "/home/user/.cache/huggingface/hub/"
    # hf_root_dir = "/root/.cache/huggingface/hub/"
    model_map = {
        "95c8c1dcca842facbdcb9606c314f7ae/AI-ModelScope/stable-diffusion-v1-5": "models--runwayml--stable-diffusion-v1-5/snapshots/1d0c4ebf6ff58a5caecab40fa1406526bca4b5b9",
        "601abc72ccca722467955bb9c36da4b3/AI-ModelScope/stable-diffusion-2-1": "models--stabilityai--stable-diffusion-2-1/snapshots/5cae40e6a2745ae2b01ad92ae5043f95f23644d6",
        "76ee531623cf180d9c98a6fbf5fa0c45.0/AI-ModelScope/stable-diffusion-xl-base-1.0": "models--stabilityai--stable-diffusion-xl-base-1.0/snapshots/462165984030d82259a11f4367a4eed129e94a7b",
        "37fb512464d701519e3be1d14c97f8c0/AI-ModelScope/clip-vit-large-patch14": "models--openai--clip-vit-large-patch14/snapshots/32bd64288804d66eefd0ccbe215aa642df71cc41",
        "c2412c56551dfb7baa71c84978ee22eb/iic/scepter": "models--scepter-studio--scepter/snapshots/d78fd86694550622e1d256865fe444282fcccb19",
        "d5bec966515f3501f44f4844c8f8e5cb/iic/scepter_scedit": "models--scepter-studio--scepter_scedit/snapshots/c8905ae6dda1ebafd2db20261ebe98cf0d9a2cd1",
        "17ad8bd379610c0c40fae1dac5fae822/iic/LARGEN": "models--scepter-studio--LARGEN/snapshots/0d8f361b9391094d3b60b9ab75c52d1d0855654d",
        "fe8fd4e872d00cf0a3fa3a4838ee16e1/iic/stylebooth": "models--scepter-studio--stylebooth/snapshots/c993df593ed451c137118b5426a1dcab29e38d5c",
        "7168a5d7e0eb13ef4facee9a8d2da34c/cubeai/blip-image-captioning-base": "models--Salesforce--blip-image-captioning-base/snapshots/89b09ea1789f7addf2f6d6f0dfc4ce10ab58ef84"
    }
    for ms_name, hf_name in model_map.items():
        ms_path =os.path.join(ms_root_dir, ms_name)
        hf_path = os.path.join(hf_root_dir, hf_name)
        cmd = f"mkdir -p {ms_path} && ln -s {hf_path}/* {ms_path}"
        os.system(cmd)
        print(cmd)

link_model()

def prepare(config):
    if 'FILE_SYSTEM' in config:
        for fs_info in config['FILE_SYSTEM']:
            FS.init_fs_client(fs_info)

    if 'LOG_FILE' in config:
        logger = get_logger()
        tid = '{0:%Y%m%d%H%M%S%f}'.format(datetime.datetime.now()) + ''.join(
            [str(random.randint(1, 10)) for i in range(3)])
        init_logger(logger, log_file=config['LOG_FILE'].format(tid))

class TabManager():
    def __init__(self):
        pass

config_file = os.path.join(os.path.dirname(scepter.dirname), "scepter/methods/studio/scepter_ui.yaml")
config = Config(load=True, cfg_file=config_file)

# for v in config.FILE_SYSTEM:
    # v.TEMP_DIR = os.path.join("./", v.TEMP_DIR)
    # v.TEMP_DIR = os.path.join("/data/.huggingface", v.TEMP_DIR)
    # v.TEMP_DIR = os.path.join("~/.cache/huggingface/hub", v.TEMP_DIR)
    # os.makedirs(v.TEMP_DIR, exist_ok=True)

prepare(config)

language = "en"
debug = False


tab_manager = TabManager()
interfaces = []
for info in config['INTERFACE']:
    name = info.get('NAME_EN',
                    '') if language == 'en' else info['NAME']
    ifid = info['IFID']
    if not FS.exists(info['CONFIG']):
        info['CONFIG'] = os.path.join(os.path.dirname(scepter.dirname), info['CONFIG'])
    if not FS.exists(info['CONFIG']):
        raise f"{info['CONFIG']} doesn't exist."
    interface = None
    if ifid == 'home':
        from scepter.studio.home.home import HomeUI
        interface = HomeUI(info['CONFIG'],
            is_debug=debug,
            language=language,
            root_work_dir=config.WORK_DIR) 
    if ifid == 'tuner_manager':
        from scepter.studio.tuner_manager.tuner_manager import TunerManagerUI
        interface = TunerManagerUI(info['CONFIG'],
                                    is_debug=debug,
                                    language=language,
                                    root_work_dir=config.WORK_DIR)
    if ifid == 'inference':
        from scepter.studio.inference.inference import InferenceUI
        interface = InferenceUI(
            info['CONFIG'],
            is_debug=debug,
            language=language,
            root_work_dir=config.WORK_DIR)
    if ifid == '':
        pass  # TODO: Add New Features
    if interface:
        interfaces.append((interface, name, ifid))
        setattr(tab_manager, ifid, interface)

with gr.Blocks() as demo:
    if 'BANNER' in config:
        gr.HTML(config.BANNER)
    else:
        gr.Markdown(
            f"<h2><center>{config.get('TITLE', 'SCEPTER Studio')}</center></h2>"
        )
    with gr.Tabs(elem_id='tabs') as tabs:
        setattr(tab_manager, 'tabs', tabs)
        for interface, label, ifid in interfaces:
            with gr.TabItem(label, id=ifid, elem_id=f'tab_{ifid}'):
                interface.create_ui()
        for interface, label, ifid in interfaces:
            interface.set_callbacks(tab_manager)

demo.queue(status_update_rate=1, api_open=False).launch(share=False, show_error=True, enable_queue=True)