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

import os
import re
import time
import json
import requests
import subprocess
import cloudpickle as pickle
from datetime import timedelta
from IPython.display import clear_output
import psutil 
def is_process_running(script_name):
    """检查脚本是否正在运行"""
    for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
        try:
            if proc.info['name'] == 'python' and script_name in ' '.join(proc.info['cmdline']):
                return True
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            pass
    return False


# 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')


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

settings = load_settings()
ngrok_token = settings.get('ngrok_token', "")
zrok_token = settings.get('zrok_token', "")
commandline_arguments = settings.get('commandline_arguments', "")
change_webui = settings.get('change_webui', "")


# ========================== OTHER ==========================
def is_package_installed(package_name):
    try:
        subprocess.run(["npm", "ls", "-g", package_name], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        return True
    except subprocess.CalledProcessError:
        return False

lt_tunnel = is_package_installed('localtunnel')


# ======================== 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)
        return response.json().get('ip', 'N/A')
    except Exception as e:
        print(f"Error getting public {version} address:", e)

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

tunnel_class = pickle.load(open(f"{root_path}/new_tunnel", "rb"), encoding="utf-8")
tunnel_port = 1834
tunnel = tunnel_class(tunnel_port)
tunnel.add_tunnel(command="cl tunnel --url localhost:{port}", name="cl", pattern=re.compile(r"[\w-]+\.trycloudflare\.com"))

if lt_tunnel:
    tunnel.add_tunnel(command="lt --port {port}", name="lt", pattern=re.compile(r"[\w-]+\.loca\.lt"), note="Password : " + "\033[32m" + public_ipv4 + "\033[0m")

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/",
    "ad_extra_models_dir": f"{webui_path}/models/adetailer/",
    "sd_checkpoint_hash": "",
    "sd_model_checkpoint": "",
    "sd_vae": "None"
}

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

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

    for key, value in paths_to_check.items():
        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)


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'

    if change_webui == 'A1111':
        commandline_arguments += ' --xformers'
    else:
        commandline_arguments += ' --disable-xformers --opt-sdp-attention --cuda-stream --pin-shared-memory'
    process_name = "python upload.py"  
    try:
        subprocess.check_output(["pgrep", "-f", process_name])
        print("Upload 脚本已经在运行中")
    except subprocess.CalledProcessError:
        # 脚本没有运行，启动它
        print("启动 Upload 脚本")
        # subprocess.Popen(["nohup", "python", "./upload.py", ">", "/dev/null", "2>&1", "&"])
        get_ipython().system('COMMANDLINE_ARGS="{commandline_arguments}" python upload.py &> /dev/null &')
    # get_ipython().system('COMMANDLINE_ARGS="{commandline_arguments}" python upload.py &> /dev/null &')
    get_ipython().system('COMMANDLINE_ARGS="{commandline_arguments}" python launch.py')

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")

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