#!/bin/sh
#_(
   exec clojure -M "$0" "$@"
   )

;; What to eval is NOT baked in and must be specified. It lives in a tiny
;; per-project `.ait-repl` file, one `key value` per line. See
;; ait-repl-example in ait's scripts/ for information.

(ns repl-listen
  (:require [clojure.string :as str]
            [clojure.java.io :as io])
  (:import [java.net Socket]
           [java.io PushbackInputStream BufferedInputStream BufferedOutputStream]))

(def POLL-MS   75)
(def REQ-FILE  ".ait-eval")
(def CFG-FILE  ".ait-repl")
(def PORT-FILE ".nrepl-port")

;; ---- minimal bencode (nREPL wire format) -----------------------------------
(defn bencode ^String [o]
  (cond
    (string? o)     (let [b (.getBytes ^String o "UTF-8")] (str (count b) ":" o))
    (keyword? o)    (bencode (name o))
    (integer? o)    (str "i" o "e")
    (map? o)        (str "d" (apply str (mapcat (fn [[k v]] [(bencode k) (bencode v)]) o)) "e")
    (sequential? o) (str "l" (apply str (map bencode o)) "e")
    :else           (throw (ex-info "bencode: unsupported" {:o o}))))

(defn- rb ^long [^PushbackInputStream in]
  (let [b (.read in)]
    (when (neg? b) (throw (java.io.EOFException.)))
    b))

(defn- read-until ^String [^PushbackInputStream in end]
  (let [sb (StringBuilder.)]
    (loop []
      (let [c (rb in)]
        (if (= c (int end)) (.toString sb)
            (do (.append sb (char c)) (recur)))))))

(declare bdecode)

(defn- read-bytestring ^String [^PushbackInputStream in first-digit]
  (let [n  (Integer/parseInt (str (char first-digit) (read-until in \:)))
        ba (byte-array n)]
    (loop [off 0]
      (when (< off n)
        (let [r (.read in ba off (- n off))]
          (when (neg? r) (throw (java.io.EOFException.)))
          (recur (+ off r)))))
    (String. ba "UTF-8")))

;; Decode one bencode value. Bytestrings decode to String; nREPL keys are ASCII,
;; so string keys make the returned maps ergonomic to (get m "value") etc.
(defn bdecode [^PushbackInputStream in]
  (let [c (rb in)]
    (case (char c)
      \i (Long/parseLong (read-until in \e))
      \l (loop [acc []]
           (let [b (rb in)]
             (if (= (char b) \e) acc
                 (do (.unread in b) (recur (conj acc (bdecode in)))))))
      \d (loop [m {}]
           (let [b (rb in)]
             (if (= (char b) \e) m
                 (do (.unread in b)
                     (let [k (bdecode in) v (bdecode in)] (recur (assoc m k v)))))))
      (read-bytestring in c))))

;; ---- nREPL conversation -----------------------------------------------------
(defn- send! [^BufferedOutputStream out msg]
  (.write out (.getBytes (bencode msg) "UTF-8"))
  (.flush out))

(defn eval-once
  "One eval op; collect out/err/value/status until 'done'."
  [in out code session]
  (send! out {"op" "eval" "code" code "session" session})
  (loop [acc {:out [] :err [] :value [] :status []}]
    (let [m   (bdecode in)
          acc (cond-> acc
                (get m "out")                 (update :out conj (get m "out"))
                (get m "err")                 (update :err conj (get m "err"))
                (some? (get m "value"))       (update :value conj (get m "value")))
          st  (get m "status")
          acc (if (sequential? st) (update acc :status into st) acc)]
      (if (some #{"done"} (:status acc)) acc (recur acc)))))

(defn clj-string
  "Render s as a Clojure string literal."
  [s]
  (str \" (-> s (str/replace "\\" "\\\\") (str/replace "\"" "\\\"")) \"))

;; ---- config -----------------------------------------------------------------
(defn load-cfg [path]
  (reduce
    (fn [cfg line]
      (let [t (str/trim line)]
        (if (or (empty? t) (str/starts-with? t "#"))
          cfg
          (let [[k v] (str/split line #"\s+" 2)]
            (assoc cfg (str/trim k) (some-> v str/trim))))))
    {}
    (str/split-lines (slurp path))))

;; ---- connection (with reconnect) -------------------------------------------
(defn port-file
  "Which file holds the nREPL port. Defaults to .nrepl-port (figwheel/lein/etc.);
   shadow-cljs writes .shadow-cljs/nrepl.port, so set `port-file` in .ait-repl."
  [cfg]
  (get cfg "port-file" PORT-FILE))

(defn connect
  "Open a socket to the nREPL port, clone a fresh session, run any `upgrade`
   form once. Returns {:sock :in :out :session}."
  [cfg]
  (let [port (Integer/parseInt (str/trim (slurp (port-file cfg))))
        sock (doto (Socket. "127.0.0.1" port) (.setTcpNoDelay true))
        in   (PushbackInputStream. (BufferedInputStream. (.getInputStream sock)) 8)
        out  (BufferedOutputStream. (.getOutputStream sock))]
    (send! out {"op" "clone"})
    (let [session (get (bdecode in) "new-session")]
      (when-let [up (get cfg "upgrade")]
        (eval-once in out up session))
      {:sock sock :in in :out out :session session})))

(defn render-form [cfg code ns]
  (-> (get cfg "eval")
      (str/replace "%CODE%" (clj-string code))
      (str/replace "%NS%" (if (str/blank? ns) "cljs.user" ns))))

(defn process
  "Eval the user's form, print out/err, then optionally repaint."
  [cfg {:keys [in out session]} ns code]
  (let [r (eval-once in out (render-form cfg code ns) session)]
    (when-let [o (seq (:out r))]
      (let [s (apply str o)]
        (print (if (str/ends-with? s "\n") s (str s "\n")))))
    (when-let [e (seq (:err r))] (binding [*out* *err*] (print (apply str e))))
    ;; repaint after the user's form (its value already printed above)
    (when-let [rr (get cfg "rerender")]
      (let [r2 (eval-once in out (render-form cfg rr ns) session)]
        (when-let [e (seq (:err r2))] (binding [*out* *err*] (print (apply str e))))))
    (flush)))

;; ---- request file (nonce\nns\ncode) ----------------------------------------
(defn read-request
  "Returns [nonce ns code] or nil if the file isn't there yet."
  []
  (when (.exists (io/file REQ-FILE))
    (let [parts (str/split (slurp REQ-FILE) #"\n" 3)]
      [(nth parts 0 "") (nth parts 1 "") (nth parts 2 "")])))

(defn -main [& _]
  ;; Config/request/port files resolve against the CURRENT DIRECTORY: run this
  ;; from the project root (the dir holding .ait-repl). This tool lives on PATH
  ;; (dotfiles), not in the project — cwd is the anchor, not *file*'s location.
  (when-not (.exists (io/file CFG-FILE))
    (binding [*out* *err*]
      (println (str "repl-listen: no " CFG-FILE " in " (System/getProperty "user.dir")
                    " — run me from the project root")))
    (System/exit 1))
  (let [cfg (load-cfg CFG-FILE)]
    (when-not (get cfg "eval")
      (binding [*out* *err*] (println (str "repl-listen: " CFG-FILE " needs an `eval` line")))
      (System/exit 1))

    (let [conn (atom nil)
          ensure! (fn []
                    (loop []
                      (let [c (try (connect cfg) (catch java.io.IOException _ nil))]
                        (if c
                          (do (reset! conn c)
                              (println (str "repl-listen: connected (session "
                                            (subs (:session c) 0 8) "), watching " REQ-FILE)))
                          (do (binding [*out* *err*]
                                (println (str "repl-listen: no nREPL yet (" (port-file cfg) ") — retrying...")))
                              (Thread/sleep 1000) (recur))))))]
      (ensure!)
      ;; Don't replay a stale request left over from a previous run.
      (loop [last-nonce (first (read-request))]
        (let [[nonce ns code] (or (read-request) ["" "" ""])]
          (if (and (seq nonce) (not= nonce last-nonce))
            (do (try
                  (process cfg @conn ns code)
                  (catch java.io.IOException _
                    (binding [*out* *err*]
                      (println "repl-listen: connection lost — reconnecting..."))
                    (ensure!)))
                (Thread/sleep POLL-MS)
                (recur nonce))
            (do (Thread/sleep POLL-MS)
                (recur last-nonce))))))))

(apply -main *command-line-args*)
