;;; bend.lsp -- Lisp-smart GPU dispatch primitive.
;;;
;;; (bend expr)
;;;   inspects expr at runtime, routes to a registered GPU worker
;;;   when the cost estimator says it's worth the trip, else
;;;   evaluates locally in the original lexical scope.
;;;
;;;   The decision is per-call, based on what's in the args -- not
;;;   just the op name. Tiny SHAKE on 5 inputs stays local; big
;;;   SHAKE on 5,000,000 inputs bends to a GPU worker.
;;;
;;; (bend! expr)
;;;   forces GPU dispatch; raises an error if no worker available.
;;;
;;; Wire protocol: length-prefixed S-exp (see wire.lsp). Same shape
;;; lumbda's fleet workers already speak, so bend reuses the same
;;; tcp-* + read-from-string + write-to-string primitives every
;;; lumbda tier already provides.
;;;
;;; Per-tier integration is now ZERO new primitives -- bend.lsp
;;; works on Python, C, and asm tiers as-is.

(load "wire.lsp")

;;; -- portable error reporter -----------------------------------
;;;
;;; asm tier doesn't have `error` as a Scheme-callable builtin; the
;;; Python and C tiers do. This wraps both so the rest of the file
;;; stays portable. Displays the message, then either raises (on tiers
;;; that have it) or returns a tagged failure value.

(define (bend-error msg . rest)
  (display ";;; bend ERROR: ") (display msg)
  (for-each (lambda (x) (display " ") (display x)) rest)
  (newline)
  'bend-failure)

;;; -- tuning constants ------------------------------------------

;; Per-spawn cuda init ~200 ms; daemon-warm RPC ~100 us. Set high
;; if your workers haven't started their daemon pools yet.
(define *bend-overhead-ns* 100000)        ; 100 us once daemons are warm

;; Host throughput for SHAKE256 (cycles ? ns/byte). Comparing
;; (cost-fn args) against *bend-overhead-ns* picks the route.
(define *bend-host-ns-per-byte* 2)        ; ~500 MB/s on a modern x86

;;; -- registry -- what `bend` knows how to ship ------------------
;;;
;;; alist of (op-name . cost-fn) where cost-fn takes the literal
;;; arg list and returns estimated host-side runtime in nanoseconds.

(define *bend-gpu-ops* '())

(define (bend-register-op! name cost-fn)
  (set! *bend-gpu-ops* (cons (cons name cost-fn) *bend-gpu-ops*)))

;; Cost estimator for (cuda-shake-fanout inputs out-bytes).

;; Sum of input byte lengths x ns-per-byte.

(define (bend-shake-cost args)
  (let* ((inputs (car args))
         (n (length inputs))
         (avg-len (if (= n 0) 0
                      (quotient (string-length (car inputs)) 2))))
    (* n avg-len *bend-host-ns-per-byte*)))

(bend-register-op! 'cuda-shake-fanout bend-shake-cost)

;;; -- worker endpoint -------------------------------------------

(define *bend-worker-host* "127.0.0.1")
(define *bend-worker-port* 9091)

;; Override the default localhost:9091 endpoint.

(define (bend-set-worker! host port)
  (set! *bend-worker-host* host)
  (set! *bend-worker-port* port))

;; Probe TCP connect; return #t/#f without raising.

(define (bend-worker-available?)
  (let ((sock (tcp-connect *bend-worker-host* *bend-worker-port*)))
    (cond
      ((eq? sock #f) #f)
      (else (tcp-close sock) #t))))

;; Open a fresh connection, send the form framed, read framed reply,

;; close. Returns the result from the worker, or raises if the worker

;; responded with (bend-error...).

(define (bend-dispatch-to-gpu quoted-form)
  (let ((sock (tcp-connect *bend-worker-host* *bend-worker-port*)))
    (cond
      ((eq? sock #f)
       (bend-error "bend: tcp-connect failed to"
              (list *bend-worker-host* *bend-worker-port*)))
      (else
       (wire-send sock quoted-form)
       (let ((reply (wire-recv sock)))
         (tcp-close sock)
         (cond
           ((eq? reply #f) (bend-error "bend: worker closed connection"))
           ((not (pair? reply)) (bend-error "bend: malformed reply" reply))
           ((eq? (car reply) 'ok) (car (cdr reply)))
           ((eq? (car reply) 'error)
            (bend-error "bend worker error:" (cdr reply)))
           (else (bend-error "bend: unexpected reply" reply))))))))

;;; -- core dispatcher -------------------------------------------

;; thunk: 0-arg lambda that evaluates the form in its original scope.

;; quoted-form: literal sexp for inspection / serialization.

;; force-gpu?: when #t, must dispatch to GPU or error.

(define (bend-dispatch thunk quoted-form force-gpu?)
  (let* ((head (if (pair? quoted-form) (car quoted-form) #f))
         (entry (assoc head *bend-gpu-ops*)))
    (cond
      (force-gpu?
       (if (and entry (bend-worker-available?))
           (bend-dispatch-to-gpu quoted-form)
           (bend-error "bend!: no GPU worker available for op" head)))
      ((not entry) (thunk))
      (else
       (let* ((cost-fn (cdr entry))
              (args (cdr quoted-form))
              (est-host-ns (cost-fn args)))
         (cond
           ((< est-host-ns *bend-overhead-ns*) (thunk))
           ((bend-worker-available?) (bend-dispatch-to-gpu quoted-form))
           (else
            (display "bend: worker unreachable, running local: ")
            (display head) (newline)
            (thunk))))))))

;;; -- function form (works on every tier -- no macros needed) ----
;;;
;;; (bend-call '(cuda-shake-fanout '("a" "b") 32))
;;;
;;; Useful from lumbda's asm tier which doesn't have define-syntax.
;;; The macro form below is a thin wrapper over this.

(define (bend-call quoted-form)
  ;; The thunk fallback returns 'no-local-fallback so callers can
  ;; pattern-match without needing the asm tier to support (bend-error...).
  (bend-dispatch (lambda () 'no-local-fallback) quoted-form #f))

(define (bend!-call quoted-form)
  (bend-dispatch (lambda () 'gpu-required-but-no-worker) quoted-form #t))

;;; -- macros live in bend-macros.lsp (Python/C tiers) ----------
;;; asm tier doesn't have define-syntax; load bend-macros.lsp only on
;;; tiers that support it. The function-form bend-call / bend!-call
;;; above is the asm-portable path.

;;; -- demo ------------------------------------------------------

(define (make-input n)
  (let loop ((i 0) (acc '()))
    (if (= i n) acc
        (loop (+ i 1) (cons "deadbeefcafebabe1234567890abcdef" acc)))))

(define (demo)
  ;; 3 inputs: cost is tiny (96 bytes), stays local.
  (display "tiny (3 inputs): bend chooses ")
  (let ((r (bend (cuda-shake-fanout '("00" "01" "deadbeef") 32))))
    (display "-> ") (display (length r)) (display " result(s)\n"))
  ;; 100k inputs x 16 bytes = 1.6 MB; should bend if worker is up.
  (display "heavy (100k inputs): bend chooses ")
  (let ((r (bend (cuda-shake-fanout (make-input 100000) 32))))
    (display "-> ") (display (length r)) (display " result(s)\n")))

;; (demo)  ; uncomment after launching gpu-worker.lsp on port 9091
