#!python

import sys
import os
import platform
import json
import requests
import subprocess
from tempfile import NamedTemporaryFile
from tqdm import tqdm

def get_python_version():
    return f"{sys.version_info.major}.{sys.version_info.minor}"

def is_windows():
    return platform.system().lower() == "windows"

def get_architecture():
    machine = platform.machine().lower()
    if is_windows():
        if machine == "amd64" or machine == "x86_64":
            return "win_amd64"
        elif machine == "x86" or machine == "i386":
            return "win32"
        elif machine == "arm64" or machine == "aarch64":
            return "win_arm64"
    return machine

def get_wheel_url(python_version):
    wheels_json_url = "https://raw.githubusercontent.com/celray/python-gdal-installer/refs/heads/main/gdal-wheels.json"
    response = requests.get(wheels_json_url)
    wheels = response.json()
    
    arch = get_architecture()
    version_wheels = wheels.get(python_version, {})
    return version_wheels.get(arch)

def main():
    if is_windows():
        python_version = get_python_version()
        wheel_url = get_wheel_url(python_version)
        
        if not wheel_url:
            print(f"No GDAL wheel found for Python {python_version} on {get_architecture()}")
            sys.exit(1)
        
        print(f"Downloading GDAL wheel for Python {python_version} on {get_architecture()}")
        
        # Download with progress bar
        response = requests.get(wheel_url, stream=True)
        total_size = int(response.headers.get('content-length', 0))
        
        with NamedTemporaryFile(suffix='.whl', delete=False) as tmp_file:
            with tqdm(total=total_size, unit='iB', unit_scale=True) as pbar:
                for data in response.iter_content(chunk_size=1024):
                    size = tmp_file.write(data)
                    pbar.update(size)
            wheel_path = tmp_file.name

        # install dependencies
        subprocess.check_call([sys.executable, "-m", "pip", "install", wheel_path])
        subprocess.check_call([sys.executable, "-m", "pip", "install", "gdal"])

        # Clean up
        os.unlink(wheel_path)


    else:
        # On Unix systems, just use pip but install to user site (but some system require --break-system-packages which is not used here)
        subprocess.check_call([sys.executable, "-m", "pip", "install", "GDAL", "--user"])

if __name__ == "__main__":
    main()
