(define-module (env)
  #:use-module (gcrypt hash)
  #:use-module (guix base32)
  #:use-module (guix build utils)
  #:use-module (guix hash)
  #:use-module (ice-9 ftw)
  #:use-module (ice-9 match)
  #:use-module (ice-9 textual-ports)
  #:use-module (rnrs bytevectors)
  #:use-module ((scheme base) #:select (bytevector-append))
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:export (add-layer
            call-with-env
            collect-garbage!
            <overlay-mount>
            make-overlay-mount
            overlay-mount?
            overlay-mount-lower
            overlay-mount-upper
            overlay-mount-work
            overlay-mount-merged
            <environment>
            environment-current
            environment-name
            environment-root
            environment-tag-layers
            environment-tag-root?
            environment-tag-value
            environment-tags
            environment-tree
            environment?
            init-environment
            layers-in-use
            list-environments
            lookup-environment
            mount-environment
            remove-env-no-checks!
            roll-back!))

(define directory-files (@@ (ice-9 ftw) directory-files))

(define (overlayfs-command lower upper work merged)
  ;; NOTE lower dirs are listed top to bottom from left to right!
  (let ((lower* (match lower
                  (() (throw 'overlayfs-no-lower-dir))
                  ((l ...) (string-join l  ":"))
                  (l l))))
    (list "fuse-overlayfs"
          "-o"
          (string-append "lowerdir=" lower*
                         ",upperdir=" upper
                         ",workdir=" work)
          merged)))

;; TODO handle missing env variables
;; TODO handle missing directory
(define-syntax-rule (define-env-path name s s* ...)
  ;; Canonicalize path?
  (define name (string-append s s* ...)))

(define-env-path %cache-directory (or (getenv "XDG_CACHE_HOME")
                                      (string-append (getenv "HOME")
                                                     "/.cache")))
(define-env-path %env-cache     %cache-directory "/env")
(define-env-path %env-layers    %env-cache       "/layers")
(define-env-path %env-links     %env-layers      "/l")
(define-env-path %env-by-name   %env-cache       "/by-name")
(define-env-path %env-directory (getenv "HOME")  "/env")

(define (metadata-file name . path)
  (apply
   string-append (cons* %env-by-name
                        "/" (symbol->string name)
                        path)))

(define (layer-file hash)
  (string-append  %env-layers "/" hash))

(define (symlink-layer! hash)
  ;; TODO: check for collisions
  (symlink (string-append "../" hash)
           (string-append %env-links "/" (string-take hash 7))))

(define (ensure-environment-metadata-directory! name)
  (for-each mkdir-p (list %env-layers
                          %env-links
                          (metadata-file name)
                          %env-directory
                          (metadata-file name "/tags"))))

;; TODO Should we detect symlink loops here?
(define (file->string file)
  (call-with-input-file file get-string-all))

(define (file->sexp file)
  (with-input-from-file file read))

(define-record-type <environment>
  (%make-environment name current root tags tree)
  environment?
  (name environment-name)
  (current environment-current)
  (root environment-root)
  (tags environment-tags)
  (tree environment-tree))

(define-record-type <overlay-mount>
  (make-overlay-mount lower upper work merged)
  overlay-mount?
  (lower overlay-mount-lower)
  (upper overlay-mount-upper)
  (work overlay-mount-work)
  (merged overlay-mount-merged))

(define (lookup-environment name)
  ;; NOTE will fail if .../by-name/name is a symlink
  (if (directory-exists? (metadata-file name))
      (%make-environment
       name
       (file->string (metadata-file name "/tags/current"))
       (file->string (metadata-file name "/tags/root"))
       (map string->symbol (directory-files (metadata-file name "/tags")))
       (file->sexp (metadata-file name "/tree")))
      #f))

(define (environment-tag-value env tag)
  (file->string (metadata-file (environment-name env)
                               "/tags/" (symbol->string tag))))

(define (hash-link v1 v2)
  (bytevector-hash (bytevector-append v1 v2)
                   (hash-algorithm sha256)))

(define (decode-hash-link-encode s1 s2)
  (bytevector->base32-string
   (apply hash-link
          (map base32-string->bytevector (list s1 s2)))))

(define (decode-hash-link-encode/chain hashes)
  (match hashes
    (() (throw 'missformed-argument))
    ((h) h)
    ((h h* ...) (fold (λ (e acc) (decode-hash-link-encode acc e))
                      h h*))))

(define (file-hash/overlay file)

  (define trim
    (let ((n (1+ (string-length file))))
      (λ (f) (string-drop f n))))

  (define (canonicalize-deleted l)
    (if (null? l)
        l
        (let ((bvs (map string->utf8 (sort (map trim l) string<?))))
          (apply bytevector-append
                 (cdr (reverse
                       (fold (λ (a b) (cons* a #vu8(0) b))
                             '()
                             bvs)))))))

  (let* ((res '())
         (select? (λ (f stat)
                    (if (eq? (stat:type stat) 'char-special)
                        (begin (set! res (cons f res))
                               #f)
                        #t)))
         (dirhash (file-hash* file #:select? select?)))
    (if (null? res)
        dirhash
        (hash-link dirhash
                   (canonicalize-deleted res)))))

(define (environment-tag-root? env tag)
  (equal? (environment-tag-value env tag) (environment-root env)))

(define (environment-tag-layers env tag)
  (let ((root-hash (environment-root env))
        (pinned    (environment-tag-value env tag))
        (atom? (compose not pair?)))
    (call/cc
     (λ (cont)
       (when (environment-tag-root? env tag)
         (cont (list (environment-root env))))
       (let loop ((tree  (cdr (environment-tree env)))
                  (hash  root-hash)
                  (route (list root-hash)))
         (match tree
           (() route)
           (((? atom? head) tail ...)
            (let ((h (decode-hash-link-encode hash head))
                  (r (cons head route)))
              (if (equal? pinned h)
                  (cont (reverse r))
                  (loop tail h r))))
           ((((? atom? head)
              (? pair? branch1)
              (? pair? branch2)
              (? pair? branch3) ...)
             tail ...)
            (let* ((h (decode-hash-link-encode hash head))
                   (r (cons head route)))
              (if (eq? pinned h)
                  (cont (reverse r))
                  (begin
                    (loop tail h r)
                    (for-each (λ (br) (loop br h r))
                              (cons* branch1 branch2 branch3))))))
           (else (throw 'missformed-tree))))))))

(define (copy-file/overlay source destination)
  (if (eq? (stat:type (stat source)) 'char-special)
      (mknod destination 'char-special #o000 0)
      (copy-file source destination)))

(define (copy-recursively/overlay source dest)
  (copy-recursively source dest #:copy-file copy-file/overlay))

(define (init-environment name dir)
  ;; should we include vcs files?
  (if (lookup-environment name)
      (throw 'environment-esists!)
      (begin
        (ensure-environment-metadata-directory! name)
        (let* ((dirhash (bytevector->base32-string (file-hash* dir)))
               (loc (layer-file dirhash)))
          (copy-recursively dir loc)
          (symlink-layer! dirhash)
          (with-directory-excursion (metadata-file name)
            (with-output-to-file "./tags/root" (λ () (display dirhash)))
            (symlink "./root" "./tags/current")
            (symlink "./tags/current" "current")
            (with-output-to-file "tree" (λ () (write (list dirhash))))))
        (lookup-environment name))))

(define* (add-layer env layer #:optional (onto 'current))
  ;; TODO currently ignores stree structure of tree file!!
  ;; TODO deduplicate with init-env
  (unless (file-exists? layer)
    (throw 'add-layer-directory-doesnt-exist!))
  (let* ((name (environment-name env))
         (dirhash (bytevector->base32-string (file-hash/overlay layer)))
         (metadata-dir (metadata-file (environment-name env)))
         (pin (environment-tag-value env onto))
         (new-pin (decode-hash-link-encode pin dirhash))
         (stack (environment-tag-layers env onto))
         (loc (layer-file dirhash)))
    ;; TODO check if appendable
    (unless (file-exists? loc)
      (copy-recursively/overlay layer loc)
      ;; TODO: check for collisions
      (symlink-layer! dirhash))
    (with-directory-excursion (metadata-file name)
      (when (symbolic-link? "./tags/current")
        (delete-file "./tags/current"))
      (with-output-to-file  "./tags/current"
        (λ () (display new-pin)))
      (with-output-to-file "./tree"
        (λ () (write (append stack (list dirhash))))))
    (lookup-environment name)))

(define* (roll-back! env #:optional (tag 'current))
  (let* ((stack (environment-tag-layers env tag))
         (new-stack (match stack
                      (() (throw 'empty-stack))
                      ((a) (throw 'first-generation-cant-roll-back))
                      ((a a* ...) (drop-right (cons a a*) 1)))))
    (with-directory-excursion (metadata-file (environment-name env))
      (with-output-to-file "./tree"
        (λ () (write new-stack)))
      (with-output-to-file  "./tags/current"
        (λ () (display (decode-hash-link-encode/chain new-stack)))))
    (values (lookup-environment (environment-name env))
            (last stack))))

(define (remove-env-no-checks! name)
  "Remove environment brutally. No checks whether any layers are shared
by other environments made."
  (let* ((metadata-dir (string-append %env-by-name "/" name))
         (tree (with-input-from-file (string-append metadata-dir "/tree")
                 read))
         (layers (let loop ((tree tree)
                            (layers '()))
                   (match tree
                     (() layers)
                     (((a ...) tail ...)
                      (append (loop a '())
                              (loop tail layers)))
                     ((a tail ...)
                      (loop tail
                            (cons a layers)))))))
    (for-each
     (λ (l)
       (format (current-error-port) "Deleting layer: ~a\n" l)
       (delete-file-recursively (string-append %env-layers "/" l))
       (format (current-error-port) "Cleaning up links for: ~a\n" l)
       (delete-file (string-append %env-links "/" (string-take l 7))))
     layers)
    (format (current-error-port) "Deleting metadata: ~a\n" name)
    (delete-file-recursively metadata-dir)))

(define (list-environments)
  (let loop ((names (map string->symbol (directory-files %env-by-name)))
             (envs '()))
    (match names
      (() envs)
      ((n n* ...)
       (let ((e (lookup-environment n)))
         (if e
             (loop n* (cons e envs))
             (begin
               (format (current-error-port)
                       "Warning: Clobbered by-name directory: ~a\n" n)
               (loop n* envs))))))))

(define (layers-in-use)
  (let loop ((envs (list-environments))
             (used '()))
    (match envs
      (() used)
      ((e e* ...)
       (loop e*
             (fold
              (λ (t acc)
                (lset-union equal?
                            acc
                            (environment-tag-layers e t)))
              used
              (environment-tags e)))))))

(define* (collect-garbage! #:optional dry-run)
  (with-directory-excursion %env-layers
    (for-each
     (λ (l)
       (if dry-run
           (format (current-error-port) "Would delete layer ~a\n" l)
           (begin
             (format (current-error-port) "Deleting layer ~a\n" l)
             (delete-file-recursively l)
             (delete-file (string-append "l/" (string-take l 7))))))
     (lset-xor equal?
               (delete "l" (directory-files "."))
               (layers-in-use)))))

(define* (mount-environment env #:key (tag 'current) target)
  (let* ((name (environment-name env))
         (metadata-dir (metadata-file name))
         (tree (environment-tree env))
         ;; TODO sanity check
         (current (environment-current env))
         (layers-stack (environment-tag-layers env 'current)))

    (when (or (not layers-stack) (null? layers-stack))
      (throw 'no-layers-stack-found))

    (let* ((mkdtemp*
            (λ (pfix)
              ;; /tmp/env name is too common
              (mkdir-p "/tmp/env")
              (mkdtemp (string-append
                        "/tmp/env/" (symbol->string name)
                        "." pfix ".XXXXXX"))))
           (merged (let ((f (or target
                                (string-append %env-directory
                                               "/" (symbol->string name)))))
                     (if (directory-exists? f) f (begin (mkdir-p f) f))))
           (upper  (mkdtemp* "upper"))
           (work   (mkdtemp* "work"))
           ;; TODO abstract away
           (layers-links (map (λ (l) (string-take l 7))
                              layers-stack)))

      (with-directory-excursion %env-links
        ;; TODO report missing links? or global sanity check.
        ;; TODO handle errors
        (match (apply system*
                      (overlayfs-command (reverse layers-links)
                                         upper
                                         work
                                         merged))
          (0 (values (make-overlay-mount layers-links upper work merged) 0))
          (err (values #f err)))))))

(define* (call-with-env env proc #:optional save? clean?)
  (match (mount-environment env)
    (($ <overlay-mount> lower upper work merged)
     (let ((res (proc merged)))
       (sync)
       (when save?
         (when (not (equal? 0 (system* "umount" merged)))
           (throw 'umount-error))
         (add-layer env upper))
       (when clean?
         (when file-exists? upper
               (delete-file-recursively upper))
         (delete-file-recursively merged)
         (delete-file-recursively work))
       res))
    (#f (throw 'mount-failure))))