(import (scheme base) (scheme write) (zilch statusbar) (zilch nix daemon) (zilch magic) (zilch lib getopt) (scheme process-context) (chicken process-context) (srfi 146) (chicken port) (chicken foreign) (chicken condition))
(foreign-declare "#include <sched.h>")
(foreign-declare "#include \"stock_overrides.h\"")
(define stock-overrides (foreign-value "stock_overrides" nonnull-c-string))

(define get-cpu-count
  (foreign-lambda* int ()
    "cpu_set_t set; sched_getaffinity(0, sizeof(set), &set); C_return(CPU_COUNT(&set));"))

(define (print-help msg)
  (when msg
    (write-string (string-append msg "\n\n") (current-error-port)))
  (write-string "Usage: zilch-cli-rust [OPTION] [TARGET...]
Process the given crate directory (or the current directory, if unspecified)
and output derivations for each target given on the command line (or all
executables in the crate, if unspecified)

  -h, --help              Print this help message.
  -b, --build             Build the store paths, rather than show their
                          derivations.
  -j, --max-jobs COUNT    The maximum amount of builds to run. Defaults
                          to the amount of cores.
  -v, --verbose           Increase the verbosity configured in the Nix
                          daemon.
  -L, --print-build-logs  Print derivation logs as they come in.
  -m, --crate-dir DIR     The directory to use as root crate.
  -r, --replace DIR       Replace the crate specified by the Cargo.toml
                          with this source directory, rather than using
                          the upstream crate. Can be specified more
                          than once.
  -z, --overrides PATH    Read build script overrides from this file.
                          By default, a stock set of overrides is read.
                          This can be disabled by passing `-z \"\"`
                          (a blank string).

   --debug                Crash on the first error, rather than
                          continuing with the next package.
" (current-error-port))
  (exit (or (not msg) 1)))

(define-values (options args)
 (getopt
  '((help #f #\h)
    (build #f #\b)
    (max-jobs #t #\j)
    (verbose #f #\v)
    (print-build-logs #f #\L)
    (crate-dir #t #\m)
    (replace #t #\r)
    (overrides #t #\z)
    (debug #f))
  (list->vector (cdr (command-line)))
  print-help))

(when (assoc 'help options) (print-help #f))

; Set up the logger.
(define (set-print-logs val) #f)
(let ((prev-error-handler (current-exception-handler))) (current-exception-handler (lambda data (set-print-logs #t) (apply prev-error-handler data))))
(when (terminal-port? (current-error-port))
 (let-values (((new-out new-err statusbar-set-print-logs logger) (statusbar-logger (current-output-port) (current-error-port) (assoc 'print-build-logs options))))
   (current-output-port new-out)
   (current-error-port new-err)
   (set! set-print-logs statusbar-set-print-logs)
   (*logger* logger)))

;; Flags passed to the nix daemon:
(define max-jobs (if (assoc 'max-jobs options) (string->number (cdr (assoc 'max-jobs options))) (get-cpu-count)))
(define verbosity 3)
(for-each (lambda (val) (when (eq? (car val) 'verbose) (set! verbosity (+ 1 verbosity)))) options)
(write-string (string-append "Connected to Nix daemon, version " (daemon-link-daemon-version (*daemon*)) "\n") (current-error-port))

(daemon-wop-set-options (*daemon*) verbosity max-jobs #t)
(import (scheme base) (scheme file) (scheme write) (zilch magic) (zilch zexpr) (zilch nix drv) (zilch file) (zilch lang rust registry) (zilch lang rust cargo) (zilch vfs) (zilch lang rust resolver) json (srfi 128) (srfi 146) (zilch nixpkgs))
(import (chicken format))

(define (read-file fname)
  (call-with-input-file fname (lambda (p) (read-string 999999 p))))

(define crate-dir (let ((m (assoc 'crate-dir options))) (if m (cdr m) (current-directory))))

(define (remove-target-dir vfs)
  (vfs-dir-filter-all
   (lambda (dirname)
     (not (string=? dirname "target")))
   vfs))

(define root-vfs (remove-target-dir (vfs-from-directory crate-dir)))

(define-values (cargo-toml cargo-workspace) (parse-cargo-toml root-vfs (read-file (string-append crate-dir "/Cargo.toml")) #f))
(define projects '())

(define (add-crate-targets dirname vfs crate workspace)
  ; (printf "Adding crate target.. ~A\n" crate)
  (set! projects (cons (cons crate vfs) projects))
  (define (check-dep dep)
   (define is-path (cargo-dep-path? (cargo-dependency-origin dep)))
   (define path (and is-path (cargo-dep-path-path (cargo-dependency-origin dep))))
   (when is-path
    (let*
     ((root (if (and (pair? path) (eq? (car path) 'workspace)) crate-dir dirname))
      (new-path (string-append root "/" (if (pair? path) (cdr path) path))))
     (append-dir new-path (remove-target-dir (vfs-from-directory new-path)) workspace))))
  (for-each check-dep (cargo-crate-dependencies crate))
  (for-each check-dep (cargo-crate-build-dependencies crate)))

(when cargo-toml
  (add-crate-targets crate-dir root-vfs cargo-toml #f))

(define seen-crate-names '())

(define (append-dir dirname vfs workspace)
  (define-values (parsed-cargo new-workspace) (parse-cargo-toml vfs (call-with-input-file (string-append dirname "/Cargo.toml") (lambda (p) (read-string 999999 p))) workspace))
  (when (and (not workspace) new-workspace)
    (set-print-logs #t)
    (fprintf (current-error-port) "Replaced directory ~S contains a Cargo workspace. This is unsupported.\n" dirname)
    (exit 1))
  (unless (member (cargo-crate-name parsed-cargo) seen-crate-names)
    (set! seen-crate-names (cons (cargo-crate-name parsed-cargo) seen-crate-names))
    (add-crate-targets dirname vfs parsed-cargo workspace)))

(when cargo-workspace
  (for-each
    (lambda (workspace-member)
      (append-dir (string-append crate-dir "/" workspace-member) (vfs-subdir root-vfs workspace-member) cargo-workspace))
    (cargo-workspace-members cargo-workspace)))

(for-each
  (lambda (kv)
    (when (eq? (car kv) 'replace)
     (append-dir (cdr kv) (remove-target-dir (vfs-from-directory (cdr kv))) #f)))
  options)

(define lockfile (parse-lockfile (read-file (string-append crate-dir "/Cargo.lock"))))

(define output (process-many-with-lockfile projects lockfile))

(define do-build (assoc 'build options))

(define (build-if-wanted path)
  (store-path-materialize path)
  (if do-build
    (begin
      (store-path-build path)
      (store-path-realisation path))
    (if (string=? (store-path-output path) "out")
      (derivation-path (store-path-drv path))
      (string-append (derivation-path (store-path-drv path)) "!" (store-path-output path)))))

(define build-script-overrides (mapping (make-default-comparator)))
(define build-script-dependency-overrides (mapping (make-default-comparator)))
(define rustc-overrides (mapping (make-default-comparator)))

(define (process-overrides data)
 (for-each
  (lambda (pair)
   (define crate-name (car pair))
   (define overrides (vector->list (cdr pair)))
   (define self (if (assoc "buildScript" overrides) (vector->list (cdr (assoc "buildScript" overrides))) (list)))
   (define dependency (if (assoc "buildScriptDependency" overrides) (vector->list (cdr (assoc "buildScriptDependency" overrides))) (list)))
   (define rustc (if (assoc "rustc" overrides) (vector->list (cdr (assoc "rustc" overrides))) (list)))
      
   (define (parse-inner vals)
     (map (lambda (kv) (cons (car kv) (nix-eval (string-append "with (import <nixpkgs> {}); \"" (cdr kv) "\"")))) vals))

   (set! build-script-overrides (mapping-set! build-script-overrides crate-name (parse-inner self)))
   (set! build-script-dependency-overrides (mapping-set! build-script-dependency-overrides crate-name (parse-inner dependency)))
   (set! rustc-overrides (mapping-set! rustc-overrides crate-name (parse-inner rustc))))
  (vector->list data)))


; If we have no request to disable stock overrides, apply them.
(unless (member '(overrides "") options)
  (process-overrides (call-with-port (open-input-string stock-overrides) json-read)))

(for-each
 (lambda (kv)
  (when (eq? (car kv) 'overrides) (unless (string=? (cdr kv) "") (process-overrides (call-with-input-file (cdr kv) json-read)))))
 options)

(define (build-script-env-overrides crate-name is-dependency)
  (mapping-ref/default (if is-dependency build-script-dependency-overrides build-script-overrides) crate-name '()))
(define (rustc-env-overrides crate-name)
  (mapping-ref/default rustc-overrides crate-name '()))

(define (should-skip pkg)
  (define is-skippable #t)
  (if (null? args) #f
    (begin
     (for-each
      (lambda (match)
       (when
        (or
         (string=? (cargo-target-name (resolved-package-cargo-target pkg)) match)
         (string=? (cargo-crate-name (resolved-package-crate pkg)) match))
        (set! is-skippable #f)))
      args)
     is-skippable)))

(for-each
  (lambda (pkg)
    (if (equal? 'bin (cargo-target-crate-type (resolved-package-cargo-target pkg)))
      (if (should-skip pkg)
        (printf "~A\t~A\t~A\tskipped\n" (cargo-crate-name (resolved-package-crate pkg)) (cargo-target-name (resolved-package-cargo-target pkg)) (cargo-target-crate-type (resolved-package-cargo-target pkg)))
        (let ((built (build-package pkg build-script-env-overrides rustc-env-overrides))) (printf "~A\t~A\tbin\t~A\n" (cargo-crate-name (resolved-package-crate pkg)) (cargo-target-name (resolved-package-cargo-target pkg)) (build-if-wanted built))))
      (printf "~A\t~A\t~A\tnot a binary\n" (cargo-crate-name (resolved-package-crate pkg)) (cargo-target-name (resolved-package-cargo-target pkg)) (cargo-target-crate-type (resolved-package-cargo-target pkg)))))
  output)