;; Defines procedures to interact with the Nix store by way of zexpressions.
;; This library defines the `<store-path>` record type, which can be used in zexps.
;; A `<store-path>` unquotes in `zexp`s as its store path.
(define-library (zilch magic)
  (import
    (scheme base) (scheme file)
    (zilch lib hash) (zilch nix daemon) (zilch nix drv) (zilch nix path)
    (zilch zexpr)
    (srfi 128) (srfi 132) (srfi 146) (srfi 152)
    (chicken base) (chicken format) socket)

  (export
    *daemon* *use-ca*
    <store-path>
    make-store-path store-path?
    store-path-drv store-path-output

    store-path-path store-path-build store-path-materialize store-path-realisation
    store-path-for-text store-path-for-fod store-path-for-drv
    store-path-for-impure-drv store-path-for-ca-drv store-path-for-ca-drv*
    store-path-realised store-path-open
   
    zilch-magic-counters)

  (begin
   ;; The daemon connection used by `(zilch magic)`.
   (define *daemon*
     (make-parameter
       (parameterize
        ((socket-send-buffer-size 4096) (socket-send-size 4096) (socket-receive-timeout 60000) (socket-send-timeout 5000))
        (let ((unix-socket (socket af/unix sock/stream)))
             (socket-connect unix-socket (unix-address "/nix/var/nix/daemon-socket/socket"))
             (let-values (((in-port out-port) (socket-i/o-ports unix-socket)))
                         (make-daemon-link in-port out-port))))))
   (daemon-wop-handshake (*daemon*))

   ;; If set to `#f`, `store-path-for-ca-drv*` will not generate
   ;; content-addressed derivations.
   (define *use-ca* (make-parameter #t))
   
   ;; A vector of counters, counting the amount of derivations made, built, and IFD'd.
   (define zilch-magic-counters (vector 0 0 0))
   
   (define (increment-counter index)
     (vector-set! zilch-magic-counters index (+ 1 (vector-ref zilch-magic-counters index))))
 
   ;; Represents a reference to an output path of a derivation, or a source file.
   ;; if `output` is `""`, `drv` is the store path to a source file.
   (define-record-type <store-path>
     (make-store-path drv output written)
     store-path?
     (drv store-path-drv)
     (output store-path-output)
     (written store-path-written set-store-path-written!))
   
   (define-record-printer (<store-path> rt out)
     (if (eqv? (store-path-output rt) "")
       (fprintf out "#<store path ~A>" (store-path-path rt))
       (fprintf out "#<store path ~A (~A!~A)>" (store-path-path rt) (derivation-path (store-path-drv rt)) (store-path-output rt))))

   ;; Returns the store path for the output associated with this `<store-path>`.
   (define (store-path-path path)
     (derivation-output-path (cdr (assoc (store-path-output path) (derivation-outputs (store-path-drv path))))))
   
   ;; Makes sure the derivation referenced by this store path exists in the daemon.
   (define (store-path-materialize path)
     (unless (store-path-written path)
       (write-drv-to-daemon (store-path-drv path))
       (set-store-path-written! path #t)))
   
   ;; Returns the output path of this store path; fetching it from the daemon if
   ;; the derivation is content-addressed.
   (define (store-path-realisation path)
     (define drv (store-path-drv path))
     (define output (store-path-output path))
     (define drv-output (cdr (assoc output (derivation-outputs drv))))
     (if (or (not (derivation-output-hash drv-output)) (bytevector? (derivation-output-hash drv-output)))
       (derivation-output-path drv-output)
       (begin
         (store-path-materialize path)
         (let ((outputs (daemon-wop-query-derivation-output-map (*daemon*) (derivation-path drv))))
           (cdr (assoc output outputs))))))

   ;; Requests that the daemon build this store path.
   (define (store-path-build path)
     (increment-counter 1)
     (daemon-wop-build-paths (*daemon*) (vector (string-append (derivation-path (store-path-drv path)) "!" (store-path-output path)))))

   ;; Writes the `<derivation>` to the Nix store, via the currently specified `*daemon*`.
   (define (write-drv-to-daemon drv)
       (define path (derivation-path drv))
       (unless (file-exists? path)
        (let ((out (open-output-string)))
            (derivation-serialize drv out)
            (daemon-wop-add-text-to-store (*daemon*) (string-append (derivation-name drv) ".drv") (get-output-string out) (derivation-path-references drv))))
       (make-store-path path "" #t))

   ;; Returns a store path representing the text.
   (define (store-path-for-text name text)
     (increment-counter 0)
     (define goal-path (make-text-path "sha256" (sha256 text) name '()))
     (unless (file-exists? goal-path) (daemon-wop-add-text-to-store (*daemon*) name text '()))
     (make-store-path goal-path "" #t))

   ;; Returns a `<store-path>` for a fixed output derivation.
   (define (store-path-for-fod name platform builder env hash-algo hash-value hash-recursive)
     (increment-counter 0)
     (define collected-env (zexp-unwrap env))
     (define collected-builder (zexp-unwrap builder))

     (define input-drvs (merge-drvs (zexp-evaluation-drvs collected-env) (zexp-evaluation-drvs collected-builder)))
     (define input-srcs (merge-srcs (zexp-evaluation-srcs collected-env) (zexp-evaluation-srcs collected-builder)))
     (define drv (make-fixed-output-derivation name platform input-drvs input-srcs (zexp-evaluation-value collected-builder) (zexp-evaluation-value collected-env) hash-algo hash-value hash-recursive))
     (make-store-path drv "out" #f))

   ;; Returns an alist of output -> `<store-path>` for an input-addressed derivation.
   (define (store-path-for-drv name platform builder env outputs)
     (increment-counter 0)
     (define collected-env (zexp-unwrap env))
     (define collected-builder (zexp-unwrap builder))

     (define input-drvs (merge-drvs (zexp-evaluation-drvs collected-env) (zexp-evaluation-drvs collected-builder)))
     (define input-srcs (merge-srcs (zexp-evaluation-srcs collected-env) (zexp-evaluation-srcs collected-builder)))
     (define drv (make-input-addressed-derivation name platform input-drvs input-srcs (zexp-evaluation-value collected-builder) (zexp-evaluation-value collected-env) outputs))
     (map (lambda (l) (cons (car l) (make-store-path drv (car l) #f))) (derivation-outputs drv)))

   ;; Returns an alist of output -> `<store-path>` for an impure derivation.
   (define (store-path-for-impure-drv name platform builder env outputs)
     (increment-counter 0)
     (define collected-env (zexp-unwrap env))
     (define collected-builder (zexp-unwrap builder))

     (define input-drvs (merge-drvs (zexp-evaluation-drvs collected-env) (zexp-evaluation-drvs collected-builder)))
     (define input-srcs (merge-srcs (zexp-evaluation-srcs collected-env) (zexp-evaluation-srcs collected-builder)))
     (define drv (make-impure-derivation name platform input-drvs input-srcs (zexp-evaluation-value collected-builder) (zexp-evaluation-value collected-env) outputs))
     (map (lambda (l) (cons (car l) (make-store-path drv (car l) #f))) (derivation-outputs drv)))

   ;; Returns an alist of output -> `<store-path>` for a content-addressed derivation.
   (define (store-path-for-ca-drv name platform builder env outputs)
     (increment-counter 0)
     (define collected-env (zexp-unwrap env))
     (define collected-builder (zexp-unwrap builder))

     (define input-drvs (merge-drvs (zexp-evaluation-drvs collected-env) (zexp-evaluation-drvs collected-builder)))
     (define input-srcs (merge-srcs (zexp-evaluation-srcs collected-env) (zexp-evaluation-srcs collected-builder)))
     (define drv (make-ca-derivation name platform input-drvs input-srcs (zexp-evaluation-value collected-builder) (zexp-evaluation-value collected-env) outputs))
     (map (lambda (l) (cons (car l) (make-store-path drv (car l) #f))) (derivation-outputs drv)))
   
   ;; Calls either `store-path-for-ca-drv` or `store-path-for-drv` depending on `*use-ca*`.
   (define (store-path-for-ca-drv* name platform builder env outputs)
     (if (*use-ca*) (store-path-for-ca-drv name platform builder env outputs)
                    (store-path-for-drv name platform builder env outputs)))

   (define (merge-drvs left right)
     ; Create a new pair for the head of each drvs list
     (define drvs (map (lambda (l) (cons (car l) (cdr l))) left))

     (for-each
       (lambda (item)
         (define left (assoc (car item) drvs derivation-equal?))
         (if left
           (for-each
             (lambda (output)
               (unless (member output (cdr left))
                 (set-cdr! left (cons output (cdr left)))))
             (cdr item))
           (set! drvs (cons item drvs))))
       right)
     (list-sort (lambda (l r) (string<? (derivation-path (car l)) (derivation-path (car r)))) (map (lambda (a) (cons (car a) (list-sort string<? (cdr a)))) drvs)))

   (define (merge-srcs left right)
     (for-each (lambda (item) (when (eq? (member item left) #f) (set! left (cons item left)))) right)
     (list-sort string<? left))

   (define (resolve-upstream-output-placeholders path drv-context)
     (define known-placeholders (mapping (make-default-comparator)))
     ; Returns #t if placeholder was replaced
     (define (replace-placeholder placeholder replacement start-index)
       (define index (string-contains path placeholder start-index))
       (if index
         (begin
           (set! path (string-replace path replacement index (+ index (string-length placeholder))))
           (replace-placeholder placeholder replacement (+ index (string-length replacement))))
         (> start-index 0)))

     (define context-to-build '())
     (define drv-output-map #f)
     (define placeholders-to-build '())

     (define (check-output-needs-building output-name output drv)
       (when (and (derivation-output-placeholder? output) (string-contains path (derivation-output-path output)))
         (unless drv-output-map
           (set! drv-output-map (daemon-wop-query-derivation-output-map (*daemon*) (derivation-path drv))))
         (let ((known-path (cdr (assoc output-name drv-output-map))))
           (unless known-path
             (set! placeholders-to-build (cons (string-append (derivation-path drv) "!" output-name) placeholders-to-build))))))

     (define (process-output output-name output drv)
       (define needs-building #f)
       (when (and (derivation-output-placeholder? output) (string-contains path (derivation-output-path output)))
         (unless drv-output-map
           (set! drv-output-map (daemon-wop-query-derivation-output-map (*daemon*) (derivation-path drv))))
         (let* ((known-path (cdr (assoc output-name drv-output-map)))
                (is-replaced (replace-placeholder (derivation-output-path output) known-path 0)))
           (when (and is-replaced (not (file-exists? known-path)))
             (set! needs-building #t))))
       (when (or needs-building (not (or (derivation-output-placeholder? output) (file-exists? (derivation-output-path output)))))
         (set! context-to-build (cons (string-append (derivation-path drv) "!" output-name) context-to-build))))

     (for-each
       (lambda (drv-outputs)
         (set! drv-output-map #f)
         (for-each
           (lambda (output-name)
             (check-output-needs-building output-name (cdr (assoc output-name (derivation-outputs (car drv-outputs)))) (car drv-outputs)))
           (cdr drv-outputs)))
       drv-context)
     (unless (null? placeholders-to-build)
       (daemon-wop-build-paths (*daemon*) (list->vector placeholders-to-build)))
     (for-each
       (lambda (drv-outputs)
         (set! drv-output-map #f)
         (for-each
           (lambda (output-name)
             (process-output output-name (cdr (assoc output-name (derivation-outputs (car drv-outputs)))) (car drv-outputs)))
           (cdr drv-outputs)))
       drv-context)
     (values path context-to-build))

   (define (zexp-ctx-has-placeholder drv-context)
     (if (null? drv-context)
       #f
       (let ((drv (caar drv-context))
             (outputs (cdar drv-context))
             (has-placeholder #f))
         (for-each (lambda (output) (set! has-placeholder (or has-placeholder (derivation-output-placeholder? (cdr (assoc output (derivation-outputs drv))))))) outputs)
         (or has-placeholder (zexp-ctx-has-placeholder (cdr drv-context))))))

   (define (store-path-realised path)
     (define ctx (zexp-unwrap (zexp (zexp-unquote path))))
     (define val (zexp-evaluation-value ctx))
     (define to-build '())
     (if (string? val)
       (let-values (((resolved resolved-to-build) (resolve-upstream-output-placeholders val (zexp-evaluation-drvs ctx))))
         (set! val resolved)
         (set! to-build resolved-to-build))
       (when (zexp-ctx-has-placeholder (zexp-evaluation-drvs ctx))
         (error "store-path-realised: expression has dependencies on placeholder context, but isn't a string" (list path val))))
     (when (and (string? val) (not (file-exists? val)) (not (null? to-build)))
       (daemon-wop-build-paths (*daemon*) (list->vector to-build)))
     val)

   ;; Ensures the `<store-path>` exists, then opens an input port to allow reading from it.
   (define (store-path-open path)
     (increment-counter 2)
     (define output-path (store-path-realised path))
     (open-input-file output-path))

   (zexp-add-unquote-handler
     (lambda (val)
       (if (store-path? val)
           (begin
             (if (string=? (store-path-output val) "")
                 (begin (zexp-context-register-items '() (list (store-path-drv val))) (store-path-drv val))
                 (begin (store-path-materialize val) (zexp-context-register-items `((,(store-path-drv val) ,(store-path-output val))) '()) (store-path-path val))))
           #f)))))