(import 
  (chicken file posix)
  (chicken process)
  (chicken process-context)

  typed-records
  )

; Kindergarten is a Process Pool abstraction. You spawn child processes into
; the Kindergarten and at the end wait for them.
;
; Usage example with chain from SRFI 197:
;
;(chain empty-kindergarten
;       (kindergarten-spawn _ "ls" '() name: 'ls-cwd)
;       (kindergarten-spawn _ "ls" '("/") name: 'ls-root)
;       (kindergarten-spawn _ "echo" '("hello"))
;       (kindergarten-spawn _ "cat" '("kinder.scm") name: 'a-kind-cat)
;       (kindergarten-wait _))

(define (i-am-luke cmd args cwd env stdin stdout stderr)
  (when cwd (change-directory cwd))
  (when stdin (duplicate-fileno stdin fileno/stdin))
  (when stdout (duplicate-fileno stdout fileno/stdout))
  (when stderr (duplicate-fileno stderr fileno/stderr))
  (process-execute cmd args env))

(define (i-am-your-father pid stdin stdout stderr)
  (let ((stdin (and stdin (open-output-file* stdin)))
        (stdout (and stdout (open-input-file* stdout)))
        (stderr (and stderr (open-input-file* stderr))))
    (values pid stdin stdout stderr)))

(define (maybe-setup-pipe std*)
  (if std* (create-pipe) (values #f #f)))

(define (spawn cmd args
               #!key
               (stdin #t) (stdout #t) (stderr #t)
               cwd env)

  (let-values
    (((stdin/r stdin/w) (maybe-setup-pipe stdin))
     ((stdout/r stdout/w) (maybe-setup-pipe stdout))
     ((stderr/r stderr/w) (maybe-setup-pipe stderr)))

    (let ((pid (process-fork)))
      (cond
        ((zero? pid) (i-am-luke cmd args cwd env stdin/r stdout/w stderr/w)) ; noret
        ((positive? pid) (i-am-your-father pid stdin/w stdout/r stderr/r))
        (else (values #f #f #f #f)))))) ; error


(defstruct kind
  name ; Anything, chosen by the user

  (pid : fixnum)

  (stdin : (or output-port false))
  (stdout : (or input-port false))
  (stderr : (or input-port false))

  (exited? : boolean)
  (exited-normally? : boolean)
  ((exit-status 0) : fixnum)
  )

(define exit/ok? zero?)

(define (exited-ok? exited-normally? exit-status)
  (and exited-normally? (exit/ok? exit-status)))

(define (kind-exited-ok? kind)
  (and (kind-exited? kind)
       (exited-ok? (kind-exited-normally? kind) (kind-exit-status kind))))

(define (kind-spawn cmd args
                    #!key
                    (stdin #t) (stdout #t) (stderr #t)
                    cwd env name)
  (receive (pid stdin stdout stderr) (spawn cmd args
                                            stdin: stdin stdout: stdout stderr: stderr
                                            cwd: cwd env: env)
    (make-kind pid: pid name: name stdin: stdin stdout: stdout stderr: stderr)))

; TODO: Update kind
(define (kind-wait kind #!optional nohang)
  (process-wait (kind-pid kind) nohang))


(defstruct kindergarten
  ((running '()) : (list-of (struct kind)))
  ((successfully-exited '()) : (list-of (struct kind)))
  ((unsuccessfully-exited '()) : (list-of (struct kind)))
  )

(define empty-kindergarten (make-kindergarten))

(: kindergarten-enroll ((struct kindergarten) (struct kind) --> (struct kindergarten)))
; Assume kind was just spawn and is still running
(define (kindergarten-enroll kg kind)
  (kindergarten-update kg running: (cons kind (kindergarten-running kg))))

(define (kindergarten-spawn kg cmd args
                            #!key
                            (stdin #t) (stdout #t) (stderr #t)
                            cwd env name)
  (let ((kind (kind-spawn cmd args
                          stdin: stdin stdout: stdout stderr: stderr
                          cwd: cwd env: env name: name)))
    (kindergarten-enroll kg kind)))


; TODO: Update to be compatible with the kindergarten struct
(define (kindergarten-wait1 kg)
  (let loop ((next-kg '())
             (kg kg))
    (if (null? kg)
      (if (null? next-kg)
        (values '() #f)
        (loop '() next-kg))

      (let ((kind (car kg))
            (kg (cdr kg)))
        (receive (pid exited-normally? exit-status) (kind-wait kind #t)
          (if (zero? pid) ; process hasn't exited yet
            (loop (cons kind next-kg) kg)
            (let ((kind (kind-update kind exited?: #t exited-normally?: exited-normally? exit-status: exit-status))
                  (kg (delete pid kg (lambda (pid k) (= pid (kind-pid k))))))
              (values kg kind))))))))


(define (kindergarten-wait1/any kg)
  (if (null? kg)
    (values '() #f)
    (receive (pid exited-normally? exit-status) (process-wait)
      (let ((kind (and-let* ((kind (find (lambda (k) (= (kind-pid k) pid)) kg)))
                    (kind-update kind exited?: #t exited-normally?: exited-normally? exit-status: exit-status))))
        (if kind
          (let ((kg (delete pid kg (lambda (pid k) (= pid (kind-pid k))))))
            (values kg kind))
          (kindergarten-wait1/any kg))))))


(define ((while func) kg)
  (unless (null? kg)
    (receive (kg . _) (func kg)
      ((while func) kg))))

(define kindergarten-wait! (while kindergarten-wait1))
(define kindergarten-wait!/any (while kindergarten-wait1/any))


(define ((kindergarten-wait-acc-loop func) kg)
  (let loop ((ret '())
             (kg kg))
    (if (null? kg)
      ret
      (receive (kg kind) (func kg)
        (let ((ret (if (or (not kind) (kind-exited-ok? kind))
                     ret
                     (cons kind ret))))
          (loop ret kg))))))

;; @brief Wait for all children of a kindergarten. Busy poll.
(define kindergarten-wait (kindergarten-wait-acc-loop kindergarten-wait1))

;; @brief Wait for all children of a kindergarten. Blocking waitpid().
;; NOTE: May only be used if you're sure there are no child processes other
;;       than those in @a kg!
(define kindergarten-wait/any (kindergarten-wait-acc-loop kindergarten-wait1/any))