;; 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 132)
    (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-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 5000) (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))

   ;; Ensures the `<store-path>` exists, then opens an input port to allow reading from it.
   (define (store-path-open path)
     (increment-counter 2)
     (if (store-path? path)
      (let ((out-path (store-path-realisation path)))
       (unless (and out-path (file-exists? out-path)) (store-path-materialize path) (store-path-build path))
       (unless out-path (set! out-path (store-path-realisation path)))
       (open-input-file out-path))
      (let* ((ctx (zexp-unwrap (zexp (zexp-unquote path)))) (val (zexp-evaluation-value ctx)))

       ; TODO(puck): big hack to make file->store work
       (for-each
        (lambda (drv)
         (for-each
          (lambda (output)
           (when (and (string=? (string-append (derivation-output-path (cdr output)) "/-") val)
                      (not (or (not (derivation-output-hash (cdr output)))
                               (bytevector? (derivation-output-hash (cdr output))))))
            (daemon-wop-build-paths (*daemon*) (vector (string-append (derivation-path (car drv)) "!" (car output))))
            (set! val (string-append (cdr (assoc (car output) (daemon-wop-query-derivation-output-map (*daemon*) (derivation-path (car drv))))) "/-"))))
          (derivation-outputs (car drv))))
        (zexp-evaluation-drvs ctx))

       (unless (file-exists? val)
        (for-each
         (lambda (path)
          (for-each
           (lambda (output)
            (daemon-wop-build-paths (*daemon*) (vector (string-append (derivation-path (car path)) "!" output))))
           (cdr path)))
         (zexp-evaluation-drvs ctx)))

       (open-input-file val))))

   (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)))))