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

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

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

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

  (with-input-from-pipe
    (sprintf "~A ~A" zcat eggs-file)
    (cut with-output-to-file dag-file dag-generate))) 

;;
;; DAG utilities
;;
(define (dag-resolve-aliases dag aliases)
  (map
    (lambda (eggs)
      (map
        (lambda (egg)
          (or (alist-ref egg aliases eq?) egg))
        eggs))
    dag))

(define (dag-generate)
  (let ((version (read)))
    (unless (string=? version "2")
      (error "unhandled eggs index version" version))
    (let loop ((form (read))
               (last-name #f)
               (last-version #f)
               (aliases '())
               (dag '()))
      (if (eof-object? form)
        (write (cons aliases (dag-resolve-aliases dag aliases)))
        (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)
                  name
                  (if (or (not last-name)
                          new-name
                          (version>=? version last-version))
                    version
                    last-version)
                  (if new-name
                    (cons (cons last-name
                                (normalize-egg (list 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 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)
  (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)))
                 (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)))
                 (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-map
                 (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)
  (let ((workers (tabulate jobs worker-create)))
    (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 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)))
    (resolve (dag-to-levels dag) jobs)))

(let loop ((args (command-line-arguments))
           (jobs default-jobs)
           (zcat default-zcat)
           (eggs '()))
  (if (null? args)
    (unless (null? eggs)
      (run jobs zcat (reverse (map string->symbol eggs))))
    (case (string->symbol (car args))
      ((-c)
       (cache-clear)
       (loop (cdr args) jobs zcat eggs))
      ((-v)
       (verbose #t)
       (loop (cdr args) jobs zcat 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 "    -v         print verbose output")
       (print "    -j jobs    use at most 'jobs' parallel processes (default: " jobs ")")
       (print "    -z zcat    use the specified 'zcat' program (default: " zcat ")")
       (exit 0))
      ((-j)
       (loop (cddr args) (string->number (cadr args)) zcat eggs) )
      ((-z)
       (loop (cddr args) jobs (cadr args) eggs))
      (else
        => (lambda (egg) (loop (cdr args) jobs zcat (cons (symbol->string egg) eggs)))))))