import subprocess
import sys
import os
import urllib.request

def colored(text, color):
    """
    Функция для вывода цветного текста в консоль.
    """
    colors = {
        'red': '\033[91m',
        'green': '\033[92m',
        'yellow': '\033[93m',
        'blue': '\033[94m',
        'magenta': '\033[95m',
        'cyan': '\033[96m',
        'white': '\033[97m',
    }
    return colors.get(color, '') + text + '\033[0m'

def uninstall_module(module_name, python_path):
    """
    Функция для деинсталляции модуля.
    """
    print(colored(f"Удаление модуля {module_name}...", 'red'))
    pip_path = os.path.join(python_path, "Scripts", "pip.exe")
    subprocess.run([
        pip_path,
        "uninstall",
        "-y",
        module_name
    ])

def download_file(url, filename):
    """
    Функция для загрузки файла с помощью urllib.request.
    """
    print(colored(f"Загрузка файла {filename}...", 'magenta'))
    try:
        with urllib.request.urlopen(url) as response, open(filename, 'wb') as f:
            f.write(response.read())
        print(colored(f"Файл {filename} успешно загружен.", 'magenta'))
    except Exception as e:
        print(colored(f"Ошибка при загрузке файла: {e}", 'red'))

def install_module(filename, python_path):
    """
    Функция для установки модуля из файла .whl.
    """
    print(colored(f"Установка модуля из {filename}...", 'yellow'))
    pip_path = os.path.join(python_path, "Scripts", "pip.exe")
    subprocess.run([
        pip_path,
        "install",
        filename
    ])

def install_from_index(module_name, index_url, python_path):
    """
    Функция для установки модуля из указанного индекса.
    """
    print(colored(f"Установка модуля {module_name} с индекса {index_url}...", 'yellow'))
    pip_path = os.path.join(python_path, "Scripts", "pip.exe")
    subprocess.run([
        pip_path,
        "install",
        "--extra-index-url",
        index_url,
        module_name
    ])

if __name__ == "__main__":
    # Получение пути к текущей директории
    current_dir = os.path.dirname(os.path.abspath(__file__))

    # Создание полного пути к папке Python
    python_path = os.path.join(current_dir, "python_embeded")

    # Создание полного пути к папке Scripts
    python_scripts_path = os.path.join(python_path, "Scripts")
    
    # Определение путей к git
    git_path = os.path.join(current_dir, "git")
    git_bin_path = os.path.join(git_path, "bin")
    git_exe_path = os.path.join(git_bin_path, "git.exe")
    git_libexec_path = os.path.join(git_path, "libexec")

    # Обновление переменной PATH
    os.environ['PATH'] = f"{git_path};{git_bin_path};{git_libexec_path};{python_path};{python_scripts_path}"
    # os.environ['PATH'] = f"{git_path};{git_bin_path};{git_libexec_path};{os.environ['PATH']}"

    # Установка Comfy
    print(colored("Установка Comfy", 'yellow'))
    subprocess.run([
        git_exe_path,
        "clone",
        "https://github.com/comfyanonymous/ComfyUI"
    ])
    
    # Обновление numpy/onnxruntime/onnxruntime-gpu
    modules_to_uninstall = ["numpy", "onnxruntime", "onnxruntime-gpu"]
    for module in modules_to_uninstall:
        uninstall_module(module, python_path)

    # Обновление модуля pip
    print(colored("Обновление модуля pip...", 'yellow'))
    pip_path = os.path.join(python_path, "Scripts", "pip.exe")
    subprocess.run([
        pip_path,  # Добавлен путь к pip.exe
        "install",
        "--upgrade",
        "pip"
    ])
    
    # Установка numpy==1.26.4
    print(colored("Установка модуля numpy==1.26.4...", 'yellow'))
    pip_path = os.path.join(python_path, "Scripts", "pip.exe")
    subprocess.run([
        pip_path,  # Добавлен путь к pip.exe
        "install",
        "numpy==1.26.4"
    ])

    # Загрузка файла insightface-0.7.3-cp311-cp311-win_amd64.whl
    url = "https://github.com/Gourieff/Assets/raw/main/Insightface/insightface-0.7.3-cp311-cp311-win_amd64.whl"
    filename = "insightface-0.7.3-cp311-cp311-win_amd64.whl"
    download_file(url, filename)

    # Установка модуля из файла insightface-0.7.3-cp311-cp311-win_amd64.whl
    install_module(filename, python_path)

    # Установка torch/torchvision/torchaudio
    modules_to_install = ["torch", "torchvision", "torchaudio"]
    pytorch_index = "https://download.pytorch.org/whl/cu124"
    for module in modules_to_install:
        install_from_index(module, pytorch_index, python_path)
    
    # Установка onnxruntime-gpu из сетевого ресурса
    onnxruntime_index = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/"
    install_from_index("onnxruntime-gpu", onnxruntime_index, python_path)

    print(colored("Установка завершена.", 'green'))