(define-module (env)
  #:use-module (gcrypt hash)
  #:use-module (guix base32)
  #:use-module (guix build syscalls)
  #: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!
            environment->overlay-mount-maybe
            unmount-environment-maybe
            environment-mounted?))

(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) (tag onto))
  ;; 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)))

    (define (grow tree pinned layer)
      ;; TODO abstract away
      (let  ((atom? (compose not pair?)))
        (match tree
          (() (throw 'missformed-tree))
          ((root) (if (equal? pinned root)
                      (list root layer)
                      (throw 'missformed-tree)))
          ((root _ ...)
           (cons root
                 (let loop ((tree (cdr tree))
                            (hash (environment-root env)))
                   (match tree
                     (() '())
                     (((? atom? head))
                      (let ((h (decode-hash-link-encode hash head)))
                        (if (equal? pinned h)
                            (list head layer)
                            (list head))))
                     (((? atom? head) tail ...)
                      (let ((h (decode-hash-link-encode hash head)))
                        (if (equal? pinned h)
                            (list (list head (list layer) tail))
                            (cons head (loop tail h)))))
                     ((((? atom? head)
                        (? pair? branch1)
                        (? pair? branch2)
                        (? pair? branch3) ...))
                      (let* ((h (decode-hash-link-encode hash head)))
                        (if (eq? pinned h)
                            (list (cons head
                                        (cons (list layer)
                                              (cons* branch1 branch2 branch3))))
                            (list (cons head (map (λ (br) (loop br h))
                                                  (cons* branch1 branch2 branch3)))))))
                     (else (throw 'missformed-tree)))))))))

    ;; 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)
      (let ((tags/tag (string-append "./tags/" (symbol->string tag))))
        (when (and (file-exists? tags/tag) (symbolic-link? tags/tag))
          (delete-file tags/tag))
        (with-output-to-file tags/tag
          (λ () (display new-pin))))
      (with-output-to-file "./tree"
        (λ () (write (grow (environment-tree env) pin 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 tag)))

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

    (let* ((dir (begin
                  (mkdir-p "/tmp/genv")
                  (mkdtemp "/tmp/genv/XXXXXX")))
           (id  (string-take-right dir 6))
           (merged (let ((f (or target
                                (string-append %env-directory
                                               "/" (symbol->string name)))))
                     (if (directory-exists? f) f (begin (mkdir-p f) f))))
           (upper (string-append dir "/upper"))
           (work  (string-append dir "/work"))
           ;; TODO abstract away
           (layers-links (map (λ (l) (string-take l 7))
                              layers-stack)))
      (mkdir upper)
      (mkdir work)
      (let ((dirln (string-append "/tmp/genv/" (symbol->string name))))
        (when (file-exists? dirln)
          (delete-file dirln))
        (symlink dir dirln)
        (let loop ((links layers-links)
                   (n 0))
          (unless (nil? links)
            (symlink (car links) (string-append dir "/layer" (number->string n)))
            (loop (cdr links) (1+ n)))))
      (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* (unmount-environment-maybe env #:optional target)
  (let* ((name (symbol->string (environment-name env)))
         (dir (or target (string-append %env-directory "/" name)))
         (tmpdirln (string-append "/tmp/genv/" name)))
    (if (file-exists? (string-append tmpdirln "/lock"))
        #f
        (begin
          (when (file-exists? dir)
            ;; TODO check if actually mounted
            (system* "umount" dir))
          (when (file-exists? tmpdirln)
            (delete-file tmpdirln))))))

(define (fuse-overlayfs-mounts)
  (filter (λ (m) (equal? "fuse.fuse-overlayfs"
                         (mount-type m) ))
          (mounts)))

(define (fuse-overlayfs-mount-points)
  (map mount-point (fuse-overlayfs-mounts)))

(define (environment-mounted? env)
  (if (member  (string-append %env-directory
                              "/" (symbol->string (environment-name env)))
               (fuse-overlayfs-mount-points))
      #t
      #f))

(define* (environment->overlay-mount-maybe env #:optional target)
  (let* ((name (symbol->string (environment-name env)))
         (dir (readlink (string-append "/tmp/genv/" name))))
    (if (file-exists? dir)
        (let ((upper (string-append dir "/upper"))
              (work (string-append dir "/work"))
              (merged (or target (string-append %env-directory "/" name)))
              (layers-links
               (map (λ (s) (readlink (string-append dir "/" s)))
                    (filter (λ (s) (string-prefix? "layer" s))
                            (directory-files dir)))))
          (make-overlay-mount layers-links upper work merged))
        #f)))

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