import subprocess, sys, os, time, errno, shlex
from IPython.display import clear_output, Image, display
from IPython import get_ipython
from pathlib import Path

xxx = Path('/kaggle/working')
script = xxx / 'venv.py'
img = xxx / "loading.png"
vnv = Path('/kaggle/venv')
url = 'https://huggingface.co/pantat88/back_up/resolve/main/venv-torch241-cu121.tar.lz4'
fn = Path(url).name

os.chdir(xxx)
sys.path.append(str(xxx))

get_ipython().system('curl -sLO https://raw.githubusercontent.com/gutris1/segsmaker/main/script/loading.png')
get_ipython().system('curl -sLO https://raw.githubusercontent.com/gutris1/segsmaker/main/kaggle/script/pantat88.py')
display(Image(filename=str(img)))
time.sleep(1)

from pantat88 import say, download
say('【{red} Installing VENV{d} 】{red}')

req_list = [
    "curl -LO https://github.com/DEX-1101/sd-webui-notebook/raw/main/res/new_tunnel",
    "curl -Lo /usr/bin/cl https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64",
    "apt-get update",
    "apt -y install pv",
    "pip install -q aria2 cloudpickle",
    "chmod +x /usr/bin/cl"
]

for items in req_list:
    subprocess.run(shlex.split(items), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

def she_bang():
    vnv_bin = vnv / 'bin'
    old_shebang = b'#!/home/studio-lab-user/tmp/venv/bin/python3\n'
    new_shebang = f"#!{vnv}/bin/python3\n"
    
    for script in vnv_bin.glob('*'):
        if script.is_file():
            try:
                with open(script, 'r+b') as file:
                    lines = file.readlines()
                    if lines and lines[0] == old_shebang:
                        lines[0] = new_shebang.encode('utf-8')
                        file.seek(0)
                        file.writelines(lines)
                        file.truncate()
                        print(f"Updated shebang in {script.name} to {new_shebang.strip()}")

            except OSError as e:
                if e.errno == 26:
                    print(f"Skipped {script.name}")
                else:
                    print(f"Failed to update {script.name}: {e}")

def venv_install():
    os.chdir('/kaggle')
    download(url)
    
    extract_venv = f'pv {fn} | lz4 -d | tar xf -'
    get_ipython().system(extract_venv)
    Path(fn).unlink()

    get_ipython().system(f'rm -rf {vnv / "bin" / "pip*"}')
    get_ipython().system(f'rm -rf {vnv / "bin" / "python*"}')
    get_ipython().system(f'python3 -m venv {vnv}')
    get_ipython().system(f'{vnv}/bin/python3 -m pip install -q --upgrade --force-reinstall pip')

venv_install()
she_bang()

clear_output(wait=True)
script.unlink()
os.chdir(xxx)