import os
import requests
from tqdm import tqdm
import time

def download_file(url, directory, chunk_size=1024, max_retries=5, retry_delay=5):
    if not os.path.exists(directory):
        os.makedirs(directory)

    local_filename = os.path.join(directory, url.split('/')[-1].split('?')[0])

    headers = {}
    mode = 'wb'
    initial_pos = 0

    if os.path.exists(local_filename):
        initial_pos = os.path.getsize(local_filename)
        headers['Range'] = f'bytes={initial_pos}-'
        mode = 'ab'

    for attempt in range(max_retries):
        try:
            with requests.get(url, headers=headers, stream=True) as r:
                r.raise_for_status()
                total_size = int(r.headers.get('content-length', 0)) + initial_pos

                with open(local_filename, mode) as f:
                    with tqdm(
                        desc=f"Попытка {attempt+1}/{max_retries}: {local_filename}",
                        total=total_size,
                        unit='iB',
                        unit_scale=True,
                        unit_divisor=1024,
                        initial=initial_pos
                    ) as progress_bar:
                        for chunk in r.iter_content(chunk_size=chunk_size):
                            size = f.write(chunk)
                            progress_bar.update(size)
                            initial_pos += size

                return local_filename

        except (requests.exceptions.RequestException, IOError) as e:
            print(f"Ошибка при скачивании {url}: {str(e)}")
            print(f"Повторная попытка через {retry_delay} секунд...")
            time.sleep(retry_delay)

            if os.path.exists(local_filename):
                initial_pos = os.path.getsize(local_filename)
                headers['Range'] = f'bytes={initial_pos}-'
                mode = 'ab'

    raise Exception(f"Не удалось скачать файл после {max_retries} попыток")

# Список URL-адресов и соответствующих директорий
files_to_download = [
    ("https://huggingface.co/daswer123/test/resolve/main/models/ControlNet/diffusion_pytorch_model.safetensors?download=true", "models/ControlNet"),
    ("https://huggingface.co/daswer123/test/resolve/main/models/ControlNet/ip-adapter.bin?download=true", "models/ControlNet"),
    ("https://huggingface.co/daswer123/test/resolve/main/models/Stable-diffusion/moxieDiffusionXL_v16.safetensors?download=true", "models/Stable-diffusion"),
    ("https://huggingface.co/daswer123/test/resolve/main/models/VAE/sdxl_vae.safetensors?download=true", "models/VAE"),
    ("https://huggingface.co/daswer123/test/resolve/main/models/ESRGAN/4x_NMKD-Siax_2000k.pth?download=true", "models/ESRGAN")
]

# Получаем текущую рабочую директорию
current_dir = os.getcwd()

# Скачиваем каждый файл
for url, directory in files_to_download:
    full_directory = os.path.join(current_dir, directory)
    try:
        downloaded_file = download_file(url, full_directory)
        print(f"Успешно скачан файл: {downloaded_file}")
    except Exception as e:
        print(f"Не удалось скачать {url}: {str(e)}")

print("Все файлы обработаны!")