import tempfile
import hashlib
from urllib.request import urlopen
import os
import base64


def __try_download_sha(base_url, use_base64):
    url = base_url + ".sha256"
    print(f"Loading {url}")

    try:
        url_result = urlopen(url)
    except:
        print(f"No uploaded sha256 hash at {url}")
        return None

    data = url_result.read()

    sha_hex = data.decode("utf-8")
    if use_base64:
        return "sha256-" + base64.b64encode(bytes.fromhex(sha_hex)).decode("utf-8")
    return sha_hex


def download_integrities(
    module_file: str,
    has_headers: bool,
    classifiers: list[str],
    url_base: str,
    use_base64: bool,
):
    hash_output = "# Generated by shared/bazel/thirdparty/integrity_helper.py\n"

    hash_output += f'url_pattern = "{url_base}"\n\n'

    if has_headers:
        url = url_base % "headers"
        sha = __try_download_sha(url, use_base64)
        hash_output += f'headers_integrity = "{sha}"\n\n'

    hash_output += "_LIB_ARTIFACTS = {\n"

    static_glob_lookup = {
        "linux": "**/*.a",
        "osx": "**/*.a",
        "windows": "**/*.lib",
    }

    shared_glob_lookup = {
        "linux": "**/*.so*",
        "osx": "**/*.dylib",
        "windows": "**/*.dll",
    }

    for suffix in classifiers:
        url = url_base % suffix
        sha = __try_download_sha(url, use_base64)

        glob_lookup = static_glob_lookup if "static" in suffix else shared_glob_lookup

        if "linux" in suffix:
            simple_suffix = "linux"
        elif "osx" in suffix:
            simple_suffix = "osx"
        elif "windows" in suffix:
            simple_suffix = "windows"

        hash_output += f'    "{suffix}": ("{simple_suffix}", "{glob_lookup[simple_suffix]}", "{sha}"),\n'

    hash_output += "}\n"
    hash_output += "# End auto-gen\n"

    updated_contents = ""
    with open(module_file, "r") as f:
        # Grab all the stuff before the auto gen
        for line in f:
            if line.startswith(
                "# Generated by shared/bazel/thirdparty/integrity_helper.py"
            ):
                break
            updated_contents += line

        updated_contents += hash_output

        # Skip the autogen chunk
        for line in f:
            if line.startswith("# End auto-gen"):
                break

        updated_contents += "".join(f)

    with open(module_file, "w") as f:
        f.write(updated_contents)

    # print(updated_contents)


def update_ceres():
    # Keep in sync with shared/ceres.gradle
    year = "frc2026"
    version = "2.2-1"

    has_headers = True
    classifiers = [
        "linuxarm32static",
        "linuxarm32staticdebug",
        "linuxarm64static",
        "linuxarm64staticdebug",
        "linuxx86-64static",
        "linuxx86-64staticdebug",
        "osxuniversalstatic",
        "osxuniversalstaticdebug",
        "windowsarm64static",
        "windowsarm64staticdebug",
        "windowsx86-64static",
        "windowsx86-64staticdebug",
    ]
    url_base = f"https://frcmaven.wpi.edu/artifactory/development/edu/wpi/first/thirdparty/{year}/ceres/ceres-cpp/{version}/ceres-cpp-{version}-%s.zip"

    download_integrities(
        "shared/bazel/thirdparty/ceres/ceres.MODULE.bazel",
        has_headers,
        classifiers,
        url_base,
        True,
    )


def update_libssh():
    # Keep in sync with shared/libssh.gradle
    version = "2027-0.120-1"

    has_headers = True
    classifiers = [
        "linuxarm64static",
        "linuxarm64staticdebug",
        "linuxx86-64static",
        "linuxx86-64staticdebug",
        "osxuniversalstatic",
        "osxuniversalstaticdebug",
        "windowsx86-64static",
        "windowsx86-64staticdebug",
        "windowsarm64static",
        "windowsarm64staticdebug",
    ]
    url_base = f"https://frcmaven.wpi.edu/release/org/wpilib/thirdparty/libssh/{version}/libssh-{version}-%s.zip"

    download_integrities(
        "shared/bazel/thirdparty/libssh/libssh.MODULE.bazel",
        has_headers,
        classifiers,
        url_base,
        False,
    )


def update_mrclib():
    # Keep in sync with shared/libmrclib.gradle
    version = "2027.1.0-alpha-1-65-g21f308e"

    has_headers = True
    classifiers = [
        "linuxarm64",
        "linuxx86-64",
        "osxuniversal",
        "linuxsystemcore",
        "windowsarm64",
        "windowsx86-64",
    ]
    url_base = f"https://frcmaven.wpi.edu/artifactory/development-2027/org/wpilib/mrclib/mrclib-cpp/{version}/mrclib-cpp-{version}-%s.zip"

    download_integrities(
        "shared/bazel/thirdparty/mrclib/mrclib.MODULE.bazel",
        has_headers,
        classifiers,
        url_base,
        False,
    )


def main():
    update_ceres()
    update_libssh()
    update_mrclib()


if __name__ == "__main__":
    # python3 shared/bazel/thirdparty/integrity_helper.py
    main()