;;; Defines `zexp`, or zilch-expressions.
;;; A zexp is a Scheme expression that may reference other zexps, or
;;; for example `<store-path>` objects.
(define-library (zilch zexpr)
  (import
    (scheme base) (scheme read) (scheme write)
    (zilch nix drv)
    (chicken base) (chicken format))

  (cond-expand (chicken (import (chicken read-syntax))))

  (export
    <zexp> make-zexp zexp? zexp-thunk
    <zexp-context> make-zexp-context zexp-context?
    zexp-context-srcs set-zexp-context-srcs!
    zexp-context-drvs set-zexp-context-drvs!

    <zexp-evaluation> zexp-evaluation?
    zexp-evaluation-value zexp-evaluation-drvs
    zexp-evaluation-srcs

    zexp-context-register-items
    zexp zexp-quote-inner zexp-unquote
    zexp-add-unquote-handler zexp-unwrap
    zexp-with-injected-context zexp-with-context)

  (begin
    ;; A zexp (concept inspired from Guix g-expressions) is represented as a
    ;; thunk that returns the quoted value, and writes the metadata (e.g. string context) necessary
    ;; into `++*zexp-context*++`.

    ;; `(make-zexp thunk printer)`
    ;; `thunk` `(zexp-thunk zexp)` is the thunk called when evaluating the zexp.
    ;; `printer` `(zexp-printer zexp)` is a thunk that is called with a port to print a representation of the zexp.
    (define-record-type <zexp>
      (make-zexp thunk printer)
      zexp?
      (thunk zexp-thunk)
      (printer zexp-printer))

    (define-record-printer (<zexp> zexp out)
      (fprintf out "#<zexp val: ")
      ((zexp-printer zexp) out)
      (fprintf out ">"))

    ;; The context used to evaluate a zexp, stored in `++*zexp-context*++` during the evaluation.
    ;;
    ;; Stores a list of sources in `zexp-content-srcs` (settable using `set-zexp-context-srcs!`)
    ;; and an alist of derivations with a list of their outputs in `zexp-content-drvs` (settable using `set-zexp-context-drvs!`)
    ;;
    ;; Prefer using zexp-context-register-items over directly interacting with this record.
    (define-record-type <zexp-context>
      (make-zexp-context srcs drvs)
      zexp-context?
      (srcs zexp-context-srcs set-zexp-context-srcs!)
      (drvs zexp-context-drvs set-zexp-context-drvs!))

    (define-record-printer (<zexp-evaluation> zeval out)
      (fprintf out "#<zexp-context drvs: ~s; srcs: ~s>"
        (zexp-context-drvs zeval)
        (zexp-context-srcs zeval)))

    ;; The output of evaluating a `zexp`.
    ;;
    ;; drvs is an alist of derivation path to a list of outputs used.
    ;; srcs is a list of source store paths used.
    (define-record-type <zexp-evaluation>
      (make-zexp-evaluation value drvs srcs)
      zexp-evaluation?
      (value zexp-evaluation-value)
      (drvs zexp-evaluation-drvs)
      (srcs zexp-evaluation-srcs))
   
    (define-record-printer (<zexp-evaluation> zeval out)
      (fprintf out "#<zexp-evaluation val: ~s; drvs: ~s; srcs: ~s>"
        (zexp-evaluation-value zeval)
        (zexp-evaluation-drvs zeval)
        (zexp-evaluation-srcs zeval)))
    
    ;; Adds any new items from a list of sources and an alist of derivations to the current `++*zexp-context*++`.
    ;; drvs is an alist of derivation object to output. name.
    ;; TODO(puck): 'spensive?
    (define (zexp-context-register-items drvs srcs)
      (define ctx (*zexp-context*))
      (define ctx-src (and ctx (zexp-context-srcs ctx)))
      (define ctx-drvs (and ctx (zexp-context-drvs ctx)))
      (when ctx
         (for-each (lambda (src)
                       (when (eq? (member src ctx-src) #f)
                             (set! ctx-src (cons src ctx-src))
                             (set-zexp-context-srcs! ctx ctx-src))) srcs)
         (for-each (lambda (drv)
                       (define pair (assoc (car drv) ctx-drvs derivation-equal?))
                       (if (eq? pair #f)
                           (begin
                            (set! ctx-drvs (cons drv ctx-drvs))
                            (set-zexp-context-drvs! ctx ctx-drvs))
                           (for-each (lambda (output)
                                         (unless (member output (cdr pair)) (set-cdr! pair (cons output (cdr pair))))) (cdr drv)))) drvs)))

    ;; The current zexp evaluation context. #f if not evaluating a zexp.
    (define *zexp-context* (make-parameter #f))

    ; The actual zexp "quote" equivalent.
    (define-syntax zexp
      (syntax-rules (unquote)
        ((zexp-quote stuff) (make-zexp (lambda () (zexp-quote-inner stuff)) (lambda (port) (write (quote stuff) port))))))

    ; If external objects want to be unquotable, they can override this procedure.
    (define zexp-unquote-handler (lambda (v) v))

    (define zexp-unquote-handlers '())

    ;; Add a procedure to be called when unquotingg an unknown value.
    ;; This procedure should return #f if the value passed in cannot be unquoted by this handler.
    (define (zexp-add-unquote-handler handler) (set! zexp-unquote-handlers (cons handler zexp-unquote-handlers)))

    (define (iter-unquote-handler val handlers)
        (if (eq? handlers '())
            (error "Cannot unquote this value.")
            (let ((result ((car handlers) val)))
                 (if (eq? result #f)
                     (iter-unquote-handler val (cdr handlers))
                     result))))

    ;; Used in the `zexp` macro to zexp-unquote values.
    (define (zexp-unquote val)
      (cond
        ((pair? val) (cons (zexp-unquote (car val)) (zexp-unquote (cdr val))))
        ((vector? val) (vector-map (lambda (val) (zexp-unquote val)) val))

        ; (zexp (zexp-unquote (zexp (foo bar)))) -> (zexp (foo bar))
        ; TODO: keep this?
        ((zexp? val) ((zexp-thunk val)))
        ((or (boolean? val) (char? val) (null? val) (symbol? val) (bytevector? val) (eof-object? val) (number? val) (string? val)) val)
        (else (iter-unquote-handler val zexp-unquote-handlers))))

    ;; Unwraps a <zexp>, returning a <zexp-evaluation> containing the proper quoted expressions, and its dependencies.
    (define (zexp-unwrap val)
      (parameterize ((*zexp-context* (make-zexp-context '() '())))
                    (let ((nval (zexp-unquote val)))
                         (make-zexp-evaluation nval (zexp-context-drvs (*zexp-context*)) (zexp-context-srcs (*zexp-context*))))))

    ;;; Returns a `<zexp>` that returns the same value as `<val>`, but adds the drvs/srcs as context.
    (define (zexp-with-injected-context val drvs srcs)
      (make-zexp (lambda () (zexp-context-register-items drvs srcs) ((zexp-thunk val))) (lambda (port) (write val port))))

    (define (zexp-with-context fn)
      (parameterize ((*zexp-context* (make-zexp-context '() '())))
                    (let ((result (fn))) (list result  (zexp-context-drvs (*zexp-context*)) (zexp-context-srcs (*zexp-context*))))))

    ; If trying to quote a pair, we return a cons with both arguments recursively quoted.
    ; When an zexp-unquote (e.g. #~) is encountered, it is replaced with a call to the zexp-unquote procedure.
    (define-syntax zexp-quote-inner
      (syntax-rules (unquote unquote-splicing zexp-quote-inner zexp-unquote zexp-unquote-splicing)
       ((zexp-quote-inner ((zexp-unquote-splicing to-splice) . right))
        (apply
           append
           (list (map zexp-unquote (zexp-unquote to-splice))
                 (zexp-quote-inner right))))
       ((zexp-quote-inner (zexp-unquote item))
        (zexp-unquote item))

       ; (zexp-quote-inner (foo bar baz)) -> (cons (zexp-quote-inner foo) (cons (zexp-quote-inner bar) (zexp-quote-inner baz)))
       ((zexp-quote-inner (unquote item)) item)
       ((zexp-quote-inner ((unquote-splicing item) . right)) (append item (zexp-quote-inner right)))
       ((zexp-quote-inner (left)) (cons (zexp-quote-inner left) '()))
       ((zexp-quote-inner (left . right)) (cons (zexp-quote-inner left) (zexp-quote-inner right)))
       ((zexp-quote-inner item) (quote item))))

    (cond-expand
      (chicken
        (set-sharp-read-syntax! #\~
          (lambda (port) (define contents (read port)) (list 'zexp contents)))
        (set-sharp-read-syntax! #\$
          (lambda (port)
            (list
              (if (char=? (peek-char port) #\@)
                  (begin (read-char port) 'zexp-unquote-splicing)
                  'zexp-unquote)
              (read port))))))))