(import (scheme)
        (chicken base)
        (chicken condition)
        (chicken errno)
        (chicken file posix)
        (chicken file)
        (chicken foreign)
        (chicken format)
        (chicken io)
        (chicken pathname)
        (chicken platform)
        (chicken port)
        (chicken process-context posix)
        (chicken process-context)
        (chicken process)
        (chicken sort)
        (chicken string)
        (chicken tcp))

(include "socketpair")
(include "egg")
(include "http")
(include "list")
(include "string")

(define verbose (make-parameter #f))

(define default-zcat (cond-expand (macosx "gzcat") ; zcat on macOS expects the filename to end with a .Z
                                  (else "zcat")))
(define default-jobs 4)
(define default-chicken-install (make-pathname
                                  (list (foreign-value "C_INSTALL_PREFIX" c-string) "bin")
                                  chicken-install-program))

(define chicken-major (car (string-split (chicken-version) ".")))
(define egg-tarballs-server "code.call-cc.org")
(define egg-tarballs-path   (string-append "/egg-tarballs/" chicken-major "/index.gz"))

(define cache-dir (or (get-environment-variable "CHICKEN_LAY_CACHE")
                      (make-pathname (system-cache-directory) "chicken-lay")))
(define eggs-file (make-pathname cache-dir "index.gz"))
(define etag-file (make-pathname cache-dir "etag"))
(define dag-file  (make-pathname cache-dir "dag"))

(define (write/flush data port)
  (write data port)
  (flush-output port))

;;
;; Cache utilities
;;
(define (cache-check)
  (all file-exists? (list eggs-file etag-file dag-file)))

(define (cache-clear)
  (handle-exceptions
    exn
    (cond-expand
      ((or chicken-5.0 chicken-5.1 chicken-5.2 chicken-5.3)
       (void))
      (else
        (unless (eq? (get-condition-property exn 'exn 'errno #f) errno/noent)
          (abort exn))))
    (delete-directory cache-dir #t)))

(define (cache-save etag data zcat)
  (create-directory (pathname-directory etag-file) #t)
  (with-output-to-file etag-file (cut write etag))

  (create-directory (pathname-directory eggs-file) #t)
  (with-output-to-file eggs-file (cut write-string data)) 

  (call-with-input-pipe
    (sprintf "~A ~A" zcat eggs-file)
    (lambda (in)
      (call-with-output-file
        dag-file
        (lambda (out)
          (dag-generate in out))))))

;;
;; DAG utilities
;;
(define (dag-resolve-aliases dag aliases)
  (define (latest egg)
    (let* ((name (egg-name egg))
           (alias (alist-ref name aliases eq?)))
      (cond (alias
              (if (version>=? (egg-version alias) (egg-version egg))
                alias
                (error "egg version not satisfied" egg alias)))
            ((egg-to-skip? egg) #f)
            (else (error "egg not found" egg)))))
  (let loop ((dag dag)
             (seen '())
             (out '()))
    (if (null? dag)
      out
      (let* ((rest (cdr dag))
             (rule (car dag))
             (head (car rule))
             (deps (cdr rule))
             (head-name (egg-name head))
             (head-last (latest head)))
        (if (memq head-name seen)
          (loop
            rest
            seen
            out)
          (loop
            rest
            (cons head-name seen)
            (cons
              (cons head-last (filter-map latest deps))
              out)))))))

(define (dag-generate in out)
  (define (make-alias name version)
    (cons name (normalize-egg (list name version))))
  (let ((version (read in)))
    (unless (string=? version "2")
      (error "unhandled eggs index version" version))
    (let loop ((form (read in))
               (last-name #f)
               (last-version #f)
               (aliases '())
               (dag '()))
      (if (eof-object? form)
        (begin
          (let ((aliases (cons (make-alias last-name last-version) aliases)))
            (write (cons aliases (dag-resolve-aliases dag aliases)) out)))
        (let-values (((name version size checksum deps test-deps) (apply values form)))
          (let ((new-name (and last-name (not (eq? last-name name))))
                (new-rule (cons (normalize-egg (list name version))
                                (map normalize-egg deps))))
            ; We assume that all versions of an egg are listed in a contiguous
            ; block of lines in the eggs index file: whenever we encounter a
            ; new egg name, we know that we have encountered all versions of
            ; the previous egg. It's easy enough to compare versions, that we
            ; don't need to assume that the versions are listed in order.
            ; This is so we can alias egg-name to egg-name:latest-version.
            (loop (read in)
                  name
                  (if (or (not last-name)
                          new-name
                          (version>=? version last-version))
                    version
                    last-version)
                  (if new-name
                    (cons (make-alias last-name last-version) aliases)
                    aliases)
                  (cons new-rule dag))))))))

(define (dag-prune dag-data eggs)
  (let ((aliases (car dag-data))
        (full-dag (cdr dag-data)))
   (let loop ((eggs (filter (compose not egg-to-skip?) eggs))
               (dag '())
               (seen '()))
      (if (null? eggs)
        dag
        (let ((egg (or (alist-ref (car eggs) aliases eq?) (car eggs))))
          (cond
            ((memq egg seen)
             (loop (cdr eggs) dag seen))
            (else
              (let* ((rule (assq egg full-dag))
                     (_ (unless rule (error "no rule found for egg" egg)))
                     (deps (filter (lambda (node) (memq (car node) rule)) full-dag)))
                (loop
                  (flatten deps (cdr eggs))
                  (append deps dag)
                  (append rule seen))))))))))

(define (dag-to-levels dag)
  (define (depends-on? elem lvl)
    (let ((adj-list (alist-ref elem dag eq?)))
      (and adj-list (any (lambda (x) (memq x adj-list)) lvl))))
  (let loop ((topo (reverse (topological-sort dag eq?))) (lvls '()))
    (cond ((null? topo) (reverse lvls))
          ((null? lvls) (loop (cdr topo) (list (list (car topo)))))
          (else
            (let ((curr (car topo))
                  (next (cdr topo))
                  (x (car lvls))
                  (y (cdr lvls)))
              (let lvl-loop ((seen-lvl '())
                             (cand-lvl '())
                             (curr-lvl x)
                             (next-lvl y))
                (if (depends-on? curr curr-lvl)
                  (loop next
                        (append seen-lvl
                                (cons (cons curr cand-lvl)
                                      (cons curr-lvl next-lvl))))
                  (if (null? next-lvl)
                    (loop next
                          (append seen-lvl
                                  (list cand-lvl)
                                  (list (cons curr curr-lvl))))

                    (lvl-loop
                      (append seen-lvl (if (null? cand-lvl) '() (list cand-lvl)))
                      curr-lvl
                      (car next-lvl)
                      (cdr next-lvl))))))))))

;;
;; Workers utility
;;
(define (worker-free? w) (car w))
(define (worker-free! w) (set-car! w #t))
(define (worker-busy! w) (set-car! w #f))
(define (worker-out w) (cadr w))
(define (worker-in w) (caddr w))

(define (worker-ack! w)
  (or (worker-free? w)
      (and (read (worker-in w))
           (worker-free! w)))
  w)

(define (worker-create chicken-install)
  (define redirect (if (verbose) "" " > /dev/null"))
  (let-values (((p c) (socketpair)))
    (define (child)
      (let ((out (open-output-file* c))
            (in  (open-input-file* c)))
        (let loop ()
          (let ((cmd (read in)))
            (cond 
              ((eq? 'build (car cmd))
               (let* ((egg (cadr cmd))
                      (line (string-append
                              chicken-install 
                              " -no-install-dependencies -no-install "
                              (symbol->string egg)
                              redirect)))
                 (when (verbose) (print line))
                 (system* line)
                 (write/flush #t out)
                 (loop)))
              ((eq? 'install (car cmd))
               (let* ((eggs (cdr cmd))
                      (line (string-append
                              chicken-install
                              " -cached "
                              (apply string-append
                                     (map
                                       (lambda (egg)
                                         (string-append " " (symbol->string egg)))
                                       eggs))
                              redirect)))
                 (when (verbose) (print line))
                 (system* line)
                 (write/flush #t out)
                 (loop)))
              ((eq? 'exit (car cmd)))
              (else (error "got unexpected command" cmd)))))))

    (let ((pid (process-fork child #t)))
      (if (zero? pid)
        (exit)
        (list #t ; worker is free
              (open-output-file* p)
              (open-input-file* p))))))

(define (worker-build w egg)
  (worker-busy! w)
  (write/flush `(build ,egg) (worker-out w)))

(define (worker-install w eggs)
  (worker-busy! w)
  (write/flush `(install ,@eggs) (worker-out w)))

(define (workers-next-free workers)
  (let loop ((w workers))
    (cond
      ((null? w)
       (call-with-values
         (cut file-select (map (lambda (w) (port->fileno (worker-in w))) workers) #f)
         (lambda (ready-fds _)
           (car
             (map
               worker-ack!
               (filter
                 (lambda (w)
                   (memq (port->fileno (worker-in w)) ready-fds))
                 workers))))))
      ((worker-free? (car w))
       (car w))
      ((char-ready? (worker-in (car w)))
       (worker-ack! (car w)))
      (else (loop (cdr w))))))

(define (workers-wait workers)
  (for-each worker-ack! workers))

(define (workers-finish workers)
  (for-each
    (lambda (w)
      (write/flush '(exit) (worker-out w)))
    workers))

;;
;; Main logic
;;
(define (update-eggs-index)
  (let* ((local-etag (and (cache-check) 
                          (with-input-from-file etag-file read)))
         (headers (if local-etag
                    `((If-None-Match . ,(sprintf "~A" local-etag)))
                    '())))
    (let-values (((status headers body) (http-get egg-tarballs-server
                                                  egg-tarballs-path
                                                  headers)))
      (case status
        ((304) #f)
        ((200) (cons (alist-ref 'etag headers) body))
        (else
          (error (sprintf "cannot handle HTTP response with status: ~A, headers: ~A, body len: ~A"
                          status headers (string-length body))))))))

(define (resolve levels jobs chicken-install)
  (let ((workers (tabulate jobs (cut worker-create chicken-install))))
    (print "plan " levels)
    (map
      (lambda (eggs)
        (print "level " eggs)
        (map
          (lambda (egg)
            (let ((w (workers-next-free workers)))
              (print "build " egg)
              (worker-build w egg)))
          eggs)
        (workers-wait workers)
        (worker-install (workers-next-free workers) eggs)
        (workers-wait workers))
      levels)
    (workers-finish workers)))

(define (run jobs zcat chicken-install eggs)
  ; chicken-install creates these directories; make sure they are created here
  ; so we can call chicken-install in parallel without issues
  (create-directory chicken-install-cache-directory #t)
  (create-directory chicken-install-cache-metadata-directory #t)

  (let ((updated (update-eggs-index)))
    (when updated (cache-save (car updated) (cdr updated) zcat)))
  (let ((dag (dag-prune (with-input-from-file dag-file read) eggs)))
    (unless (null? dag)
      (resolve (dag-to-levels dag) jobs chicken-install))))

(let loop ((args (command-line-arguments))
           (jobs default-jobs)
           (zcat default-zcat)
           (chicken-install default-chicken-install)
           (eggs '()))
  (if (null? args)
    (unless (null? eggs)
      (run jobs zcat chicken-install (reverse (map string->symbol eggs))))
    (case (string->symbol (car args))
      ((|-i|)
       (loop (cddr args) jobs zcat (cadr args) eggs))
      ((|-c|)
       (cache-clear)
       (loop (cdr args) jobs zcat chicken-install eggs))
      ((|-v|)
       (verbose #t)
       (loop (cdr args) jobs zcat chicken-install eggs))
      ((|-h|)
       (print "Usage: lay [-h] [-c] [-v] [-j jobs] [-z zcat] egg ...")
       (print "    -h         display this help page and quit")
       (print "    -c         clear the local cache before doing anything else")
       (print "    -i prog    use 'prog' as chicken-install program (default: " chicken-install ")")
       (print "    -v         print verbose output")
       (print "    -j jobs    use at most 'jobs' parallel processes (default: " jobs ")")
       (print "    -z prog    use 'prog' as zcat program (default: " zcat ")")
       (exit 0))
      ((|-j|)
       (loop (cddr args) (string->number (cadr args)) zcat chicken-install eggs))
      ((|-z|)
       (loop (cddr args) jobs (cadr args) chicken-install eggs))
      (else
        => (lambda (egg) (loop (cdr args) jobs zcat chicken-install (cons (symbol->string egg) eggs)))))

    ))