using JSON3, Downloads
using BinaryBuilder
using Base: thisminor

# GA releases have historically been published to developer.download.nvidia.com, while
# EA/preview releases (and possibly everything, going forward) only appear on
# packages.nvidia.com's bin-archive service. The latter serves the very same JSON manifest
# format from a different root, the only difference being that the relative paths point
# into a content-addressed pool instead of sitting next to the manifest.
redist_roots(product::String) = [
    "https://developer.download.nvidia.com/compute/$product/redist",
    "https://packages.nvidia.com/bin-archive/release/$product/redist",
]

function get_sources(product::String, components::Vector{String};
                     version::Union{VersionNumber,String}, platform::Platform,
                     variant::Union{Nothing,String}=nothing)
    errors = []
    for root in redist_roots(product)
        url = "$root/redistrib_$(version).json"
        json = try
            sprint(io->Downloads.download(url, io))
        catch err
            push!(errors, url => err)
            continue
        end
        return parse_sources(json, product, components; version, platform, variant, root)
    end
    error("Could not download a redistributable manifest for $product $version:\n" *
          join(("  $url: $(sprint(showerror, err))" for (url, err) in errors), "\n"))
end
# XXX: split, for now, so that we can use this with manual JSON
function parse_sources(json::String, product::String, components::Vector{String};
                       version::Union{VersionNumber,String}, platform::Platform,
                       variant::Union{Nothing,String}=nothing,
                       root::String=first(redist_roots(product)))
    redist = JSON3.read(json)
    if !haskey(platform, "cuda")
        error("Please provide a platform that has the 'cuda' tag set, indicating which CUDA toolkit version this product is to be used with.")
    end
    cuda_version = platform["cuda"]
    architecture = if Sys.islinux(platform)
        libc(platform) == "glibc" || error("Only glibc is supported on Linux")
        if arch(platform) == "x86_64"
            "linux-x86_64"
        elseif arch(platform) == "powerpc64le"
            "linux-ppc64le"
        elseif arch(platform) == "aarch64"
            if VersionNumber(cuda_version) >= v"13"
                haskey(platform, "cuda_platform") && error("CUDA 13 uses unified ARM platforms")
                "linux-sbsa"
            else
                haskey(platform, "cuda_platform") || error("CUDA 12 and earlier require the 'cuda_platform' tag to be set")
                if platform["cuda_platform"] == "jetson"
                    "linux-aarch64"
                elseif platform["cuda_platform"] == "sbsa"
                    "linux-sbsa"
                else
                    error("Unknown cuda_platform $(platform["cuda_platform"])")
                end
            end
        else
            error("Unsupported Linux architecture $(arch(platform))")
        end
    elseif Sys.iswindows(platform)
        arch(platform) == "x86_64" || error("Only x86_64 is supported on Windows")
        "windows-x86_64"
    end

    sources = []
    for component in components
        if !haskey(redist, component)
            error("No such component $component for $product $version")
        end
        data = redist[component]

        if !haskey(data, architecture)
            error("No $architecture binaries for $component in $product $version")
        end
        data = data[architecture]

        if variant !== nothing
            if !haskey(data, variant)
                error("No $variant variant for $architecture $component in $product $version")
            end
            data = data[variant]
        end

        push!(sources, ArchiveSource(
            resolve_redist_url(root, data["relative_path"]), data["sha256"]
        ))
    end
    sources
end

# resolve a manifest-relative path against the directory the manifest lives in, folding
# away any ".." segments (bin-archive manifests use e.g. ../../../pool/<platform>/<uuid>/).
# libcurl would normalize these anyway, but the URL also ends up in build metadata.
function resolve_redist_url(root::String, relative_path::AbstractString)
    resolved = String[]
    for part in split("$root/$relative_path", '/')
        if part == ".."
            isempty(resolved) && error("Invalid relative path $relative_path")
            pop!(resolved)
        elseif part != "."
            push!(resolved, part)
        end
    end
    join(resolved, '/')
end

function build_sdk(name::String, version::VersionNumber, platforms::Vector{Platform};
                   static::Bool=false)
    script = "static=\"$(static)\"\n" * raw"""
install_license cuda_cudart*/LICENSE

mkdir ${prefix}/cuda
if [[ ${target} == *-linux-gnu ]]; then
    for project in *-archive; do
        cp -a ${project}/* ${prefix}/cuda
    done

    # keep or remove static libraries
    if [[ ${static} == "false" ]]; then
        find ${prefix}/cuda -type f -name "*_static*.a" -exec rm -f {} +
    else
        find ${prefix}/cuda \( -type f -o -type l \) ! -name "*_static*.a" -exec rm -f {} +
    fi
    find ${prefix}/cuda -type d -empty -delete
elif [[ ${target} == x86_64-w64-mingw32 ]]; then
    for project in *-archive; do
        cp -a ${project}/* ${prefix}/cuda
    done

    # keep or remove static libraries
    if [[ ${static} == "false" ]]; then
        find ${prefix}/cuda -type f -name "*_static*.lib" -exec rm -f {} +
    else
        find ${prefix}/cuda \( -type f -o -type l \) ! -name "*_static*.lib" -exec rm -f {} +
    fi
    find ${prefix}/cuda -type d -empty -delete

    # fixup
    find ${prefix}/cuda -type f \( -name "*.exe" -or -name "*.dll" \) -exec chmod +x {} +
fi"""

    # determine exactly which tarballs we should build
    builds = []
    components = [
        "cuda_cudart",
        "cuda_cuobjdump",
        "cuda_cupti",
        "cuda_nvcc",
        "cuda_nvdisasm",
        "cuda_nvml_dev",
        "cuda_nvprune",
        "cuda_nvrtc",

        "cuda_sanitizer_api",

        "libcublas",
        "libcufft",
        "libcurand",
        "libcusolver",
        "libcusparse",
        "libnpp"
    ]
    if version >= v"11.8"
        push!(components, "cuda_profiler_api")
    end
    if version >= v"12"
        push!(components, "libnvjitlink")
    end
    if version >= v"12.2"
        # available earlier, but not for aarch64
        push!(components, "libnvjpeg")
    end
    if version >= v"13"
        push!(components, "cuda_crt")
        push!(components, "libnvvm")
    end
    if version >= v"13.3"
        push!(components, "cccl")
    else
        push!(components, "cuda_cccl")
    end
    for platform in platforms
        should_build_platform(triplet(platform)) || continue

        # add the CUDA tag to the platform used by get_sources to determine compatibility.
        # note that this is different from `version`, since `get_sources` is also used by
        # recipes to select products (e.g. CUDNN) whose version differs from the CTK one.
        # we don't actually use the augmented platform to generate a JLL since the CUDA SDK
        # JLLs also need to be loadable on systems without CUDA (i.e. the musl builders).
        augmented_platform = deepcopy(platform)
        augmented_platform["cuda"] = CUDA.platform(version)

        push!(builds,
                (; script, platforms=[platform], products=Product[],
                   sources=get_sources("cuda", components; version, platform=augmented_platform)
        ))
    end

    # don't allow `build_tarballs` to override platform selection based on ARGS.
    # we handle that ourselves by calling `should_build_platform`
    non_platform_ARGS = filter(arg -> startswith(arg, "--"), ARGS)

    # `--register` should only be passed to the latest `build_tarballs` invocation
    non_reg_ARGS = filter(arg -> arg != "--register", non_platform_ARGS)

    for (i,build) in enumerate(builds)
        build_tarballs(i == lastindex(builds) ? non_platform_ARGS : non_reg_ARGS,
                    name, version, build.sources, build.script,
                    build.platforms, build.products, [];
                    skip_audit=true, julia_compat="1.6", compression_format="xz")
    end
end
