import os
import requests
import zipfile
from tqdm import tqdm

def download_file(url, filename):
    response = requests.get(url, stream=True)
    total_size = int(response.headers.get('content-length', 0))
    block_size = 1024
    t = tqdm(total=total_size, unit='iB', unit_scale=True)
    with open(filename, 'wb') as f:
        for data in response.iter_content(block_size):
            t.update(len(data))
            f.write(data)
    t.close()
    if total_size != 0 and t.n != total_size:
        print("ERROR, something went wrong")

def main():
    url = "https://github.com/akanametov/cyclegan/releases/download/1.0/horse2zebra.zip"
    dest_path = "data/horse2zebra.zip"
    os.makedirs("data", exist_ok=True)
    
    print(f"Downloading {url}...")
    try:
        download_file(url, dest_path)
        print("Extracting...")
        with zipfile.ZipFile(dest_path, 'r') as zip_ref:
            zip_ref.extractall("data")
        os.remove(dest_path)
        print("Done!")
    except Exception as e:
        print(f"Failed: {e}")

if __name__ == "__main__":
    main()