##~   LAUNCH CODE | BY: ANXETY   ~##

import os
import re
import time
import json
import requests
import cloudpickle as pickle
from datetime import timedelta
from IPython.display import clear_output


#  ================= DETECT ENV =================
def detect_environment():
    free_plan = (os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / (1024 ** 3) <= 20)
    environments = {
        'COLAB_GPU': ('Google Colab', "/root" if free_plan else "/content"),
        'KAGGLE_URL_BASE': ('Kaggle', "/kaggle/working/content")
    }

    for env_var, (environment, path) in environments.items():
        if env_var in os.environ:
            return environment, path, free_plan

env, root_path, free_plan = detect_environment()
webui_path = f"{root_path}/sdw"
#  ----------------------------------------------

def load_settings():
    SETTINGS_FILE = f'{root_path}/settings.json'
    if os.path.exists(SETTINGS_FILE):
        with open(SETTINGS_FILE, 'r') as f:
            settings = json.load(f)
        return settings

settings = load_settings()
ngrok_token = settings['ngrok_token']
zrok_token = settings['zrok_token']
commandline_arguments = settings['commandline_arguments']
change_webui = settings['change_webui']


# ======================== TUNNEL V2 ========================
print('Please Wait...')

def get_public_ip(version='ipv4'):
    try:
        url = f'https://api64.ipify.org?format=json&{version}=true'
        response = requests.get(url)
        data = response.json()
        public_ip = data['ip']
        return public_ip
    except Exception as e:
        print(f"Error getting public {version} address:", e)

# Check if public IP is already saved, if not then get it
try:
    with open(f"{root_path}/public_ip.txt", "r") as file:
        public_ipv4 = file.read().strip()
except FileNotFoundError:
    public_ipv4 = get_public_ip(version='ipv4')
    with open(f"{root_path}/public_ip.txt", "w") as file:
        file.write(public_ipv4)

tunnel_class = pickle.load(open(f"{root_path}/new_tunnel", "rb"), encoding="utf-8")
tunnel_port = 1734
tunnel = tunnel_class(tunnel_port)
tunnel.add_tunnel(command="cl tunnel --url localhost:{port}", name="cl", pattern=re.compile(r"[\w-]+\.trycloudflare\.com"))
tunnel.add_tunnel(command="lt --port {port}", name="lt", pattern=re.compile(r"[\w-]+\.loca\.lt"), note="Password : " + "\033[32m" + public_ipv4 + "\033[0m" + " rerun cell if 404 error.")

''' add zrok tunnel '''
if zrok_token:
    get_ipython().system('zrok enable {zrok_token} &> /dev/null')
    tunnel.add_tunnel(command="zrok share public http://localhost:{port}/ --headless", name="zrok", pattern=re.compile(r"[\w-]+\.share\.zrok\.io"))

clear_output()


# =============== Automatic Fixing Path V3 ===============
paths_to_check = [
    ("tagger_hf_cache_dir", f"{webui_path}/models/interrogators/"),
    ("additional_networks_extra_lora_path", f"{webui_path}/models/Lora/"),
    ("ad_extra_models_dir", f"{webui_path}/models/adetailer/"),
    ("sd_checkpoint_hash", ""),
    ("sd_model_checkpoint", ""),
    ("sd_vae", "None")
]

config_path = f'{webui_path}/ui-config.json'

with open(config_path, 'r') as file:
    config_data = json.load(file)

for key, value in paths_to_check:
    if key in config_data and config_data[key] != value:
        sed_command = f"sed -i 's|\"{key}\": \".*\"|\"{key}\": \"{value}\"|' {config_path}"
        os.system(sed_command)

# Additional check for Kaggle
if env == 'Kaggle':
   get_ipython().system('sed -i \'s/"civitai_interface\\/NSFW content\\/value":.*/"civitai_interface\\/NSFW content\\/value": false/g\' {webui_path}/ui-config.json')
#  -------------------------------------------------------


with tunnel:
    get_ipython().run_line_magic('cd', '{webui_path}')

    commandline_arguments += f' --port={tunnel_port}'
    if ngrok_token:
        commandline_arguments += f' --ngrok {ngrok_token}'
    if env != "Google Colab":
        commandline_arguments += f' --encrypt-pass={tunnel_port} --api'

    # -- FORGE --
    if change_webui == 'Forge':
        commandline_arguments += ' --cuda-stream --pin-shared-memory'

    get_ipython().system('COMMANDLINE_ARGS="{commandline_arguments}" python launch.py')


# after runnig
start_colab = float(open(f'{webui_path}/static/colabTimer.txt', 'r').read())
time_since_start = str(timedelta(seconds=time.time()-start_colab)).split('.')[0]
print(f"\n⌚️ \033[0mYou have been conducting this session for - \033[33m{time_since_start}\033[0m\n\n")

''' del zrok tunnel '''
if zrok_token:
    get_ipython().system('zrok disable &> /dev/null')