;;;; pty.scm -- Easy Pseudo-Terminal Interface
;;
;; Copyright (c) 2006-2012 Alex Shinn. All rights reserved.
;; BSD-style license: http://synthcode.com/license.txt

(declare
  (usual-integrations)
  (fixnum-arithmetic)
  (no-bound-checks)
  (no-procedure-checks))

(require-library posix)

(cond-expand
  (linux   (foreign-declare "#include <pty.h>\n")
           (foreign-declare "#include <utmp.h>\n"))

  (freebsd (foreign-declare "#include <libutil.h>\n"))

  (else    (foreign-declare "#include <util.h>\n")))

(foreign-declare "#include <fcntl.h>\n")
(foreign-declare "#include <errno.h>\n")
(foreign-declare "#include <termios.h>\n")
(foreign-declare "#include <sys/types.h>\n")
(foreign-declare "#include <sys/wait.h>\n")

(module pty
  (fcntl-ref fcntl-set! file-select-one file-read/maybe
   open-file-io/non-blocking process-alive?
   open-pty open-pty-process login-tty
   with-pty-process-io call-with-pty-process-io)

(import scheme chicken data-structures ports posix foreign)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Create non-blocking ports from a general file descriptor.

(define (yield)
  (##sys#call-with-current-continuation
   (lambda (return)
     (let ((ct ##sys#current-thread))
       (##sys#setslot ct 1 (lambda () (return (##core#undefined))))
       (##sys#schedule) ) ) ) )

(define fcntl-ref
  (foreign-lambda* int ((int fd))
   "return(fcntl(fd, F_GETFL));"))

(define fcntl-set!
  (foreign-lambda* bool ((int fd) (int arg))
    "int val = fcntl(fd, F_GETFL, 0);"
    "if (val == -1) return(0);"
    "return(fcntl(fd, F_SETFL, val | arg) != -1);") )

;; Identical to ##NET#SELECT, could also be done less efficiently with
;; FILE-SELECT from posix.
(define file-select-one
  (foreign-lambda* int ((int fd))
    "fd_set in;
     struct timeval tm;
     int rv;
     FD_ZERO(&in);
     FD_SET(fd, &in);
     tm.tv_sec = tm.tv_usec = 0;
     rv = select(fd + 1, &in, NULL, NULL, &tm);
     if(rv > 0) { rv = FD_ISSET(fd, &in) ? 1 : 0; }
     return(rv);") )

;; Alternately use FILE-READ and trap errors?
(define file-read/maybe
  (foreign-lambda int "read" int scheme-pointer int))

(define-constant +input-buffer-size+ 1024)

(define-foreign-variable error-number int "errno")
(define-foreign-variable strerror c-string "strerror(errno)")

(define-foreign-variable _ewouldblock int "EWOULDBLOCK")
(define-foreign-variable _eagain int "EAGAIN")

(define-inline (file-error msg fd)
  (##sys#update-errno)
  (##sys#signal-hook #:file-error (##sys#string-append msg strerror) fd))

(define (open-file-io/non-blocking fd . o)
  (let* ((buf (make-string +input-buffer-size+))
         (buflen 0)
         (bufindex 0)
         (iclosed #f)
         (oclosed #f)
         (more? (if (pair? o) (car o) (lambda () #t))))
    (fcntl-set! fd open/nonblock)
    (list

     ;; INPUT
     (make-input-port
      (lambda () ; read-char
        (if (>= bufindex buflen) ; clear and refill buffer
          (let ((n (let loop ()
                     (let ((n (file-read/maybe fd buf +input-buffer-size+)))
                       (if (< n 1)
                         (if (or (eq? n 0)
                                 (eq? error-number _ewouldblock)
                                 (eq? error-number _eagain))
                           (if (more?)
                             (begin
;                                (##sys#thread-block-for-i/o!
;                                       ##sys#current-thread fd #t)
                               (yield)
                               (loop))
                             ;; no more, try to read one last time to
                             ;; guard against race conditions
                             (let ((n (file-read/maybe fd buf +input-buffer-size+)))
                               (if (= n -1)
                                 (if (or (eq? error-number _ewouldblock)
                                         (eq? error-number _eagain))
                                   0
                                   (file-error "can't read from FD - " fd))
                                 n)))
                           (file-error "can't read from FD - " fd))
                         n)))))
            (cond ((zero? n)
                   #!eof)
                  (else
                   (set! buflen n)
                   (set! bufindex 1)
                   (##core#inline "C_subchar" buf 0))))
          (let ((c (##core#inline "C_subchar" buf bufindex)))
            (set! bufindex (+ bufindex 1))
            c)))
      (lambda () ; char-ready?
        (or (< bufindex buflen)
            (let ((f (file-select-one fd)))
              (when (eq? f -1)
                (file-error "can't select from FD - " fd))
              (eq? f 1))))
      (lambda () ; close-input-port
        (unless iclosed
          (set! iclosed #t)
          (file-close fd))))

     ;; OUTPUT (simple, unbuffered)
     (make-output-port
      (lambda (s) 
        (when (> (##sys#size s) 0)
          (file-write fd s)))
      (lambda ()
        (unless oclosed
          (set! oclosed #t)
          (when (and iclosed (eq? -1 (file-close fd)))
            (file-error "can't close FD output port - " fd))))))))

(define process-alive?
  (foreign-lambda* bool ((int pid))
     "int status;
      return(! waitpid(pid, &status, WNOHANG)
               || ! (WIFEXITED(status) || WIFSIGNALED(status)));"))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Low-level utilities

;; create a new pty and return the master and slave file descriptors
(define raw-open-pty
  (foreign-lambda* int ((c-string name) (int width) (int height))
   "int amaster, aslave, err;
    struct winsize ws;
    if (width && height) {
      ws.ws_row = width;
      ws.ws_col = height;
      ws.ws_xpixel = 1;
      ws.ws_ypixel = 1;
      err = openpty(&amaster, &aslave, name, NULL, &ws);
    } else {
      err = openpty(&amaster, &aslave, name, NULL, NULL);
    }
    if (err) {
      return(0);
    } else {
      return(amaster*1024 + aslave);
    }"))

;; Friendlier open-pty which allows setting the width and height of the
;; pty.  Consider using a keyword-based API and adding termios options.
(define (open-pty . o)
  (let-optionals* o ((name #f)
                     (width 0)
                     (height 0))
    (let ((res (raw-open-pty name width height)))
      (and-let* (((integer? res))
                 ((positive? res))
                 (master-fd (quotient res 1024))
                 (slave-fd (remainder res 1024)))
        (list master-fd slave-fd)))))

;; Login to a tty, setting current in/out/err ports accordingly.  Used
;; by the slave process.
(define login-tty
  (foreign-lambda int "login_tty" int))

;; Run COMMAND as a subprocess in a new pty, and return a list of two
;; values, the master FD of the new pty, and the slave PID.
(define (open-pty-process command . o)
  (let ((res (apply open-pty o))
        (command (if (and (string? command) (substring-index " " command))
                   (string-split command)
                   command)))
    (when (and (pair? res) (integer? (car res)) (integer? (cadr res)))
      ;; run slave process
      (let* ((master (car res))
             (slave (cadr res))
             (pid (process-fork
                  (lambda ()
                    (file-close master)
                    (login-tty slave)
                    (if (pair? command)
                      (process-execute (car command) (cdr command))
                      (process-execute command))))))
        ;; make non-blocking
        (fcntl-set! master open/nonblock)
        ;; return master I/O and child PID
        (list master pid)))))


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; High-level interface

;; Call and return the result of PROC with three arguments: the input,
;; output and PID of the sub-process COMMAND running in a new PTY.
;; Ensures the sub-process is terminated on completion.  Any additional
;; arguments are passed to OPEN-PTY-PROCESS.
(define (call-with-pty-process-io command proc . o)
  (unless (procedure? proc)
    (error 'call-with-pty-process-io "invalid procedure" proc))
  (let ((pty (apply open-pty-process command o)))
    (if (and (pair? pty)
             (integer? (car pty)) (not (negative? (car pty)))
             (integer? (cadr pty)) (not (zero? (cadr pty))))
      (let* ((fd (car pty))
             (pid (cadr pty))
             (ports (open-file-io/non-blocking
                     fd
                     (lambda () (process-alive? pid))))
             (in (car ports))
             (out (cadr ports))
             (res (proc in out pid)))
        (file-close fd)
        ;; Probably a bad idea
        ;;(if (process-alive? pid)
        ;;  (process-signal pid))
        res)
      (error "couldn't open-pty-process" command o pty))))

;; As above but bind to current input/output ports and only pass PROC
;; one argument, the child PID.
(define (with-pty-process-io command proc . o)
  (unless (procedure? proc)
    (error 'with-pty-process-io "invalid procedure" proc))
  (apply call-with-pty-process-io
         command
         (lambda (in out pid)
           (with-input-from-port in
             (lambda () (with-output-to-port out (proc pid)))))
         o))

)