;;; SPDX-License-Identifier: GPL-3.0-or-later
;;; Copyright © 2025 Nicolas Graves <ngraves@ngraves.fr>

;;; This file is not part of GNU Guix but requires this naming scheme
;;; so that the %cuda-updater is properly read when using
;;; `guix refresh -L$(pwd) cuda-cccl' in nonguix root.

(define-module (guix import cuda)
  #:use-module (gcrypt hash)
  #:use-module (gnu packages)
  #:use-module (guix base16)
  #:use-module (guix base32)
  #:use-module (guix http-client)
  #:use-module (guix import json)
  #:use-module (guix import utils)
  #:use-module (guix memoization)
  #:use-module (guix packages)
  #:use-module (guix records)
  #:use-module (guix upstream)
  #:use-module (guix utils)
  #:use-module (ice-9 match)
  #:use-module (ice-9 regex)
  #:use-module (ice-9 textual-ports)
  #:use-module (json)
  #:use-module (nonguix build-system cuda)
  #:use-module (srfi srfi-1)
  #:export (%cuda-updater))

(define %cuda-repository-url
  "https://developer.download.nvidia.com/compute/cuda/redist/")

(define (cuda-system->guix-system system)
  (match system
    ("linux-x86_64" "x86_64-linux")
    ("linux-aarch64" "aarch64-linux")
    ("linux-ppc64le" "powerpc64le-linux")
    (_ #f)))

(define (valid-version? version-string)
  (false-if-exception (version-prefix version-string 3)))

(define-record-type* <cuda-package>
  cuda-package make-cuda-package
  cuda-package? this-cuda-package
  (name cuda-package-name)
  (version cuda-package-version valid-version?)
  (hash-info cuda-package-hash-info cuda-hash-info?))

(define-record-type* <cuda-hash-info>
  cuda-hash-info make-cuda-hash-info
  cuda-hash-info? this-cuda-hash-info
  (system cuda-hash-info-system)
  (sha256 cuda-hash-info-sha256))

(define (cuda-toolkit-latest-version)
  (let* ((url (string-append %cuda-repository-url "index.html"))
         (port (http-fetch url #:text? #t))  ; FIXME no error management
         (html (get-string-all port))
         (regex "redistrib_[0-9.]*.json")
         (file-string
          (fold-matches regex html ""
                        (lambda (matching void)
                          (match:substring matching))))
         (version-string
          (string-drop-right
           (string-drop file-string (string-length "redistrib_"))
           (string-length ".json"))))
    (close-port port)
    version-string))

(define (cuda-json-pkg-alist->cuda-package cuda-pkg-alist)
  (make-cuda-package
   (snake-case (first cuda-pkg-alist))
   (assoc-ref cuda-pkg-alist "version")
   (filter
    identity
    (map (lambda (system)
           (let ((inner-alist (assoc-ref cuda-pkg-alist system)))
             (if inner-alist
                 (make-cuda-hash-info
                  (cuda-system->guix-system system)
                  (bytevector->nix-base32-string
                   (base16-string->bytevector
                    (assoc-ref inner-alist "sha256"))))
                 #f)))
         (list "linux-x86_64" "linux-aarch64" "linux-ppc64le")))))

(define cuda-db-fetch
  (memoize
   (lambda (toolkit-version)
     (map
      cuda-json-pkg-alist->cuda-package
      (filter list?
              (json-fetch
               (string-append %cuda-repository-url
                              "redistrib_" toolkit-version ".json")))))))

(define (cuda-fetch name tk-version)
  (let ((candidates (filter
                     (lambda (pkg) (equal? (cuda-package-name pkg) name))
                     (cuda-db-fetch tk-version))))
    (and (not (null? candidates)) (car candidates))))

(define* (latest-release package #:key (version #f))
  "Return an <upstream-source> for the latest-release of PACKAGE."
  (let* ((name (package-name package))
         (version (or version (cuda-toolkit-latest-version)))
         (package (cuda-fetch name version))
         (version (and=> package cuda-package-version)))
    (and version
         (upstream-source
          (package name)
          (version version)
          (urls (list (cuda-module-url name version)))))))

(define (make-cuda-sexp cuda-package)
  `(define-public ,(string->symbol (cuda-package-name cuda-package))
     (package
       (name ,(cuda-package-name cuda-package))
       (version ,(cuda-package-version cuda-package))
       (source
        (origin
          (method url-fetch)
          (uri (cuda-module-url name version))
          (sha256
           (base32
            (match (or (%current-target-system) (%current-system))
              ,@(map
                 (lambda (info)
                   (list (cuda-hash-info-system info)
                         (cuda-hash-info-sha256 info)))
                 (cuda-package-hash-info cuda-package)))))))
       (build-system cuda-build-system)
       (synopsis #f)
       (description #f)
       (home-page "https://developer.nvidia.com/cuda-toolkit")
       (license (cuda-license name)))))

(define (guix-name->cuda-name package)
  (string-join (string-split package  #\-) "_"))

(define (cuda-package? package)
  "Return true if PACKAGE is a CUDA Toolkit package."
  (eq? (package-build-system package) cuda-build-system))

(define %cuda-updater
  (upstream-updater
   (name 'cuda)
   (description "Updater for Cuda packages")
   (pred cuda-package?)
   (import latest-release)))

;; The issue with guix native importer is that it will only update
;; the x64_86 hash, but we do have different sources based on
;; (%current-target-system).

;; To update all hashes of a package, use:
;; (make-cuda-sexp (cuda-fetch "cuda-profiler-api" "12.1.1"))

;; To update all hashes of all packages, use:
;; (use-modules (ice-9 pretty-print))
;; (for-each
;;  (lambda (name)
;;    (pretty-print (make-cuda-sexp (cuda-fetch name "12.1.1"))))
;;  '("cuda-cccl"
;;    "cuda-cudart"
;;    "cuda-cuobjdump"
;;    "cuda-cuxxfilt"
;;    "cuda-cupti"
;;    "cuda-gdb"
;;    "cuda-nvcc"
;;    "cuda-nvml-dev"
;;    "cuda-nvdisasm"
;;    "cuda-nvprune"
;;    "cuda-nvrtc"
;;    "cuda-nvtx"
;;    "cuda-opencl"
;;    "cuda-sanitizer-api"
;;    "libcublas"
;;    "libcufft"
;;    "libcurand"
;;    "libcusolver"
;;    "libcusparse"
;;    ;; "libnvfatbin"
;;    "libnvjitlink"
;;    "libnvjpeg"
;;    "libnpp"))

;; cuda.scm ends here.