;;; gpu-worker.lsp -- TCP listener that dispatches bend-forms to local
;;; CUDA daemon binaries.
;;;
;;; Wire protocol: length-prefixed S-exp (see wire.lsp). Same shape
;;; the lumbda fleet workers in `ecdsa/lumbda/fleet/worker.lsp`
;;; already speak -- gpu-worker is the GPU-routing variant.
;;;
;;; Each accepted connection:
;;;   1. read one framed request   (wire-recv)
;;;   2. dispatch by op head       (handle-...)
;;;   3. send one framed reply     (wire-send)
;;;   4. close
;;;
;;; Daemons (long-lived --daemon binaries) are held open across all
;;; connections so CUDA context init pays once per worker startup,
;;; never per request. This is the architecture that makes the bend
;;; pattern actually faster than local CPU for repeated calls.
;;;
;;; Run:
;;;   lumbda gpu-worker.lsp                  ; default port 9091
;;;   lumbda gpu-worker.lsp --port 9001
;;;
;;; Requires the cuda binaries on disk; paths below.

(load "wire.lsp")

(define *worker-port* 9091)
(define *binary-shake-fanout*
  ;; Override via env or per host.
  "./shake256-fanout")

;; demo_ops lives in the foxhop ecdsa repo and runs upstream-format
;; ops.bin against the CPU+GPU simulators, writing (cuda-sim-result …)
;; portals via --portal. No daemon mode: spawn-per-call.
(define *binary-demo-ops*
  (or (get-environment-variable "DEMO_OPS")
      "/home/fox/git/www.foxhop.net/ecdsa/cuda/demo_ops"))

(define (parse-port-arg args)
  (let loop ((rest args))
    (cond
      ((null? rest) *worker-port*)
      ((null? (cdr rest)) *worker-port*)
      ((and (string? (car rest)) (string=? (car rest) "--port"))
       (or (string->number (car (cdr rest))) *worker-port*))
      (else (loop (cdr rest))))))

(define *daemons* '())    ; alist (op-name . (stdin-port . stdout-port))

;;; -- daemon pool (per-tier stubs marked) -----------------------
;;;
;;; spawn-process-stdio returns (stdin-port . stdout-port) for a
;;; long-running subprocess. Implementation varies by tier:
;;;   Python: subprocess.Popen with stdin/stdout=PIPE
;;;   C:      fork + pipe + exec + dup2
;;;   asm:    syscall fork + pipe + execve
;;;
;;; lumbda's fleet doesn't have it today; the first place it's
;;; needed is here. Should land in each tier's primitive table
;;; alongside `tcp-*`.

(define (start-daemon binary-path)
  (let* ((pair (spawn-process-stdio binary-path '("--daemon")))
         (ready (read-line (cdr pair))))
    (if (string=? ready "ready")
        pair
        (error "daemon failed to ready:" ready))))

;; Send `process[-bin] in out` to daemon, await `done` or `error`.

(define (daemon-process daemon-pair in-portal out-portal use-bin?)
  (let ((cmd (if use-bin? "process-bin " "process ")))
    (display cmd (car daemon-pair))
    (display in-portal (car daemon-pair))
    (display " " (car daemon-pair))
    (display out-portal (car daemon-pair))
    (newline (car daemon-pair))
    (flush-port (car daemon-pair))
    (let ((line (read-line (cdr daemon-pair))))
      (cond
        ((eq? line #f) (cons 'error "daemon closed"))
        ((>= (string-length line) 5)
         (cond
           ((string=? (substring line 0 4) "done") 'ok)
           ((string=? (substring line 0 5) "error") (cons 'error line))
           (else (cons 'error (string-append "?: " line)))))
        (else (cons 'error (string-append "?: " line)))))))

(define (register-daemon! op-name binary-path)
  (set! *daemons*
        (cons (cons op-name (start-daemon binary-path)) *daemons*))
  (display "gpu-worker: ready ") (display op-name)
  (display " <- ") (display binary-path) (newline))

;;; -- op handler -- cuda-shake-fanout ----------------------------

(define (gensym-path prefix suffix)
  ;; Each tier already has current-time-ms or gensym; if neither is
  ;; present, fall back to a counter. Simple uniqueness only.
  (string-append prefix "-" (number->string (current-time-ms)) suffix))

(define (write-shake-input-portal! path inputs out-bytes)
  (let ((port (open-output-file path)))
    (display "(cuda-shake-fanout\n" port)
    (display "  (output-bytes " port) (display out-bytes port) (display ")\n" port)
    (display "  (inputs\n" port)
    (for-each (lambda (h)
                (display "    \"" port) (display h port) (display "\"\n" port))
              inputs)
    (display "))\n" port)
    (close-port port)))

(define (read-shake-output-portal path)
  (let ((sexp (read-from-string (file->string path))))
    (let loop ((children (cdr sexp)))
      (cond
        ((null? children) '())
        ((and (pair? (car children)) (eq? (car (car children)) 'hashes))
         (cdr (car children)))
        (else (loop (cdr children)))))))

;; If x looks like (quote (...)) -- bend's serialized form -- return

;; the inner list. Otherwise return x unchanged.

(define (unquote-list x)
  (cond
    ((and (pair? x) (eq? (car x) 'quote) (pair? (cdr x)))
     (car (cdr x)))
    (else x)))

(define (handle-cuda-shake-fanout args)
  (let* ((inputs (unquote-list (car args)))
         (out-bytes (car (cdr args)))
         (in-path (gensym-path "/tmp/bend-in" ".portal"))
         (out-path (gensym-path "/tmp/bend-out" ".portal"))
         (daemon (cdr (assoc 'cuda-shake-fanout *daemons*))))
    (write-shake-input-portal! in-path inputs out-bytes)
    (let ((status (daemon-process daemon in-path out-path #f)))
      (cond
        ((eq? status 'ok)
         (let ((result (read-shake-output-portal out-path)))
           (delete-file in-path)
           (delete-file out-path)
           (list 'ok result)))
        (else
         (delete-file in-path) (delete-file out-path)
         (list 'error (cdr status)))))))

;;; -- op handler -- cuda-sim-ops-bin -----------------------------
;;;
;;; Request: (cuda-sim-ops-bin <ops-bin-path> <n-batches>)
;;; Spawns demo_ops <path> <n-batches> --portal <tmp>, drains its
;;; stdout to EOF (process exit), reads the portal, returns the
;;; parsed (cuda-sim-result ...) form unchanged so callers can pull
;;; out gates / phase / timings.

(define (drain-to-eof port)
  ;; Read until read-line returns #f (or eof-object). Discard lines.
  (let loop ()
    (let ((line (read-line port)))
      (cond
        ((eq? line #f) #t)
        ((eof-object? line) #t)
        (else (loop))))))

;; Pull a numeric value out of a parsed portal expression for telemetry.
;; Walks (a b c ...) looking for (name <number>); returns #f if missing.
(define (portal-find-number form name)
  (let loop ((rest (cdr form)))
    (cond
      ((null? rest) #f)
      ((and (pair? (car rest)) (eq? (car (car rest)) name)
            (pair? (cdr (car rest))) (number? (car (cdr (car rest)))))
       (car (cdr (car rest))))
      (else (loop (cdr rest))))))

;; Walk through (timing-ms (cpu-total X) (gpu-kernel Y)) shape.
(define (portal-timing form which)
  (let loop ((rest (cdr form)))
    (cond
      ((null? rest) #f)
      ((and (pair? (car rest)) (eq? (car (car rest)) 'timing-ms))
       (portal-find-number (car rest) which))
      (else (loop (cdr rest))))))

(define (handle-cuda-sim-ops-bin args)
  (let* ((ops-path (car args))
         (n-batches (car (cdr args)))
         (portal-path (gensym-path "/tmp/bend-sim-ops" ".portal"))
         (t-start (current-time-ms)))
    (cond
      ((not (file-exists? ops-path))
       (display ";;; bend ERROR cuda-sim-ops-bin ops-bin-missing ")
       (display ops-path) (newline)
       (list 'error (list 'ops-bin-missing ops-path)))
      (else
       (display ";;; bend RECV cuda-sim-ops-bin ops=")
       (display ops-path)
       (display " n-batches=") (display n-batches)
       (display " t-ms=") (display t-start) (newline)
       (let* ((argv (list ops-path
                          (number->string n-batches)
                          "--portal" portal-path))
              (pair (spawn-process-stdio *binary-demo-ops* argv)))
         ;; demo_ops does not read stdin; drain stdout until exit.
         (drain-to-eof (cdr pair))
         (close-port (car pair))
         (close-port (cdr pair))
         (let ((wall-ms (- (current-time-ms) t-start)))
           (cond
             ((file-exists? portal-path)
              (let* ((result (read-from-string (file->string portal-path)))
                     (cpu-ms (portal-timing result 'cpu-total))
                     (gpu-ms (portal-timing result 'gpu-kernel))
                     (mismatches (portal-find-number result 'mismatches)))
                (delete-file portal-path)
                (display ";;; bend DONE cuda-sim-ops-bin")
                (display " n-batches=") (display n-batches)
                (display " wall-ms=") (display wall-ms)
                (display " cpu-ms=") (display (or cpu-ms 'NA))
                (display " gpu-ms=") (display (or gpu-ms 'NA))
                (display " mismatches=") (display (or mismatches 'NA))
                (cond
                  ((and (number? cpu-ms) (number? gpu-ms) (> gpu-ms 0))
                   (display " gpu/cpu=")
                   (display (/ cpu-ms gpu-ms))))
                (newline)
                (list 'ok result)))
             (else
              (display ";;; bend FAIL cuda-sim-ops-bin no-portal wall-ms=")
              (display wall-ms) (newline)
              (list 'error (list 'no-portal portal-path))))))))))

;;; -- dispatch --------------------------------------------------

(define (handle-request sexp)
  (cond
    ((not (pair? sexp)) (list 'error "not a form"))
    (else
     (let ((op (car sexp)) (args (cdr sexp)))
       (cond
         ((eq? op 'cuda-shake-fanout) (handle-cuda-shake-fanout args))
         ((eq? op 'cuda-sim-ops-bin) (handle-cuda-sim-ops-bin args))
         ((eq? op 'ping) (list 'ok 'pong))
         (else (list 'error (list 'unknown-op op))))))))

;; Binary wire mode: payload starts with magic "BSHK" then a
;; daemon-binary-portal blob (u32 out_bytes | u32 n | n x (u32 len + bytes)).
;; Worker writes the blob to disk, calls daemon process-bin, reads the
;; binary result, prepends "BSHR" magic, wire-send-raws it back.

(define (handle-binary-shake client payload)
  (let* ((daemon (cdr (assoc 'cuda-shake-fanout *daemons*)))
         (in-path (gensym-path "/tmp/bend-bin-in" ".bin"))
         (out-path (gensym-path "/tmp/bend-bin-out" ".bin"))
         (portal-blob (substring payload 4 (string-length payload))))
    (write-binary-file in-path portal-blob)
    (let ((status (daemon-process daemon in-path out-path #t)))
      (cond
        ((eq? status 'ok)
         (let ((result-blob (read-binary-file out-path)))
           (delete-file in-path) (delete-file out-path)
           (wire-send-raw client (string-append "BSHR" result-blob))))
        (else
         (delete-file in-path) (delete-file out-path)
         (wire-send-raw client (string-append "BERR" (cdr status))))))))

;; Accept one client, handle one request, close. Returns #t to keep
;; serving, #f when the server should stop.
(define (handle-one server)
  (let ((client (tcp-accept server)))
    (cond
      ((eq? client #f) #t)
      (else
       (let ((payload (wire-recv-raw client)))
         (cond
           ((eq? payload #f) (tcp-close client) #t)
           ((and (>= (string-length payload) 4)
                 (string=? (substring payload 0 4) "BSHK"))
            (handle-binary-shake client payload)
            (tcp-close client) #t)
           (else
            (let* ((req (read-from-string payload))
                   (resp (handle-request req)))
              (wire-send client resp)
              (tcp-close client)
              #t))))))))

;;; -- entry -----------------------------------------------------

(define (run-loop server)
  (handle-one server)
  (run-loop server))

(define (main)
  (let ((port (parse-port-arg *argv*)))
    (set! *worker-port* port)
    (register-daemon! 'cuda-shake-fanout *binary-shake-fanout*)
    (let ((server (tcp-listen port)))
      (cond
        ((eq? server #f)
         (display ";;; ERROR -- tcp-listen failed on port ")
         (display port) (newline))
        (else
         (display "gpu-worker listening on port ")
         (display port) (newline)
         (run-loop server))))))

;; (main)  ; uncomment to run; needs the cuda-shake-fanout binary
;;            on disk + spawn-process-stdio primitive per tier
