import os
import re
import sys
import subprocess
import shutil

def huggingface():
    os.rename("/mnt/workspace/stable-diffusion-webui/LoraDownload.py","/mnt/workspace/stable-diffusion-webui/LoraDownload.py1")
    os.rename("/mnt/workspace/stable-diffusion-webui/CheckpointDownload.py","/mnt/workspace/stable-diffusion-webui/CheckpointDownload.py1")
    directory = '/'
    for dirpath, _, filenames in os.walk(directory):
        for filename in filenames:
            filepath = os.path.join(dirpath, filename)
            if filename.endswith(('.py', '.txt', '.md', '.sh', '.json', '.yaml')):
                try:
                    with open(filepath, 'r', encoding='utf-8', errors='replace') as file:
                        content = file.read()
                    new_content = re.sub(r'https://huggingface\.co', 'https://hf-mirror.com', content)
                    with open(filepath, 'w', encoding='utf-8', errors='replace') as file:
                        file.write(new_content)
                except Exception as e:
                    print(f"Failed to process file {filepath}: {e}")
    os.rename("/mnt/workspace/stable-diffusion-webui/LoraDownload.py1","/mnt/workspace/stable-diffusion-webui/LoraDownload.py")
    os.rename("/mnt/workspace/stable-diffusion-webui/CheckpointDownload.py1","/mnt/workspace/stable-diffusion-webui/CheckpointDownload.py")

def API():
    source = "/mnt/workspace/stable-diffusion-webui/frpc_linux_amd64_v0.2"
    destination = "/usr/local/lib/python3.10/dist-packages/gradio"
    if os.path.exists(source):
        subprocess.run(["cp", "-f", source, destination], check=True)
        subprocess.run(['sudo', 'chmod', '+x', '/usr/local/lib/python3.10/dist-packages/gradio/frpc_linux_amd64_v0.2',])
        print("API success")
    else:
        pass

def move():
    source_files = [
        "/mnt/workspace/stable-diffusion-webui/SD！启动！.ipynb",
        "/mnt/workspace/stable-diffusion-webui/Photo",
    ]
    destination_dir = "/mnt/workspace"
    for file_path in source_files:
        if os.path.exists(file_path):
            try:
                shutil.move(file_path, destination_dir)
                pass
            except Exception as e:
                pass
        else:
            pass
        
def sd_requirements():
    os.chdir("/mnt/workspace/stable-diffusion-webui")
    subprocess.run(["pip", "install", "-r", "requirements_versions.txt","--extra-index-url","https://download.pytorch.org/whl/cu117"])

def controlnet_requirement():
    os.chdir('/mnt/workspace/stable-diffusion-webui')
    os.system('pip install handrefinerportable-2024.2.12.0-py2.py3-none-any.whl')
    
def main():
    sd_requirements()
    controlnet_requirement()
    huggingface()
    API()
    move()

if __name__ == "__main__":
    main()