#!/usr/bin/env bb

(def version "2.0.0")

(require '[babashka.cli :as cli]
         '[clojure.java.io :as io]
         '[clojure.string :as str]
         '[bencode.core :as bencode]
         '[clojure.pprint :as pp]
         '[clojure.walk :as walk]
         '[cheshire.core :as json])

;; Load hook modules (loaded conditionally when needed)
(def hook-modules-loaded (atom false))

(import '[java.net Socket]
        '[java.io PushbackInputStream])

(def cli-spec
  {:e {:desc "Expression to evaluate (everything after -e is treated as code)"
       :ref "<expr>"}
   :f {:desc "File to load and execute"
       :ref "<file>"}
   :m {:desc "Raw nREPL message (EDN format)"
       :ref "<message>"
       :alias :message}
   :h {:desc "nREPL host"
       :ref "<host>"
       :default "localhost"
       :default-desc "localhost or BREPL_HOST"}
   :p {:desc "nREPL port (required - auto-detects from .nrepl-port or BREPL_PORT)"
       :ref "<port>"}
   :verbose {:desc "Show raw nREPL messages instead of parsed output"}
   :version {:desc "Show brepl version"}
   :help {:desc "Show this help message"
          :alias :?}
   :hook {:desc "Output Claude Code hook-compatible JSON format"}})

(defn print-help []
  (println "brepl - Fast Babashka nREPL client for one-shot code evaluation")
  (println)
  (println "USAGE:")
  (println "    brepl [OPTIONS] -e <expr>")
  (println "    brepl [OPTIONS] -f <file>")
  (println "    brepl [OPTIONS] -m <message>")
  (println)
  (println "OPTIONS:")
  (println (cli/format-opts {:spec cli-spec :order [:e :f :m :h :p :verbose :version :help]}))
  (println)
  (println "PORT RESOLUTION:")
  (println "    Port is resolved in the following order:")
  (println "    1. -p <port> command line argument")
  (println "    2. .nrepl-port file:")
  (println "       - For -f: searches from file's directory upward to find project-specific port")
  (println "       - For -e/-m: uses .nrepl-port in current directory")
  (println "    3. BREPL_PORT environment variable")
  (println "    4. Error if none found")
  (println)
  (println "EXAMPLES:")
  (println "    brepl -e '(+ 1 2 3)'")
  (println "    brepl -f script.clj")
  (println "    brepl -m '{\"op\" \"describe\"}'")
  (println "    brepl -p 7888 -e '(println \"Hello\")'")
  (println "    BREPL_PORT=7888 brepl -e '(+ 1 2)'"))

(defn read-nrepl-port []
  (when (.exists (io/file ".nrepl-port"))
    (-> (slurp ".nrepl-port")
        str/trim
        Integer/parseInt)))

(defn find-nrepl-port-in-parents
  "Search for .nrepl-port file starting from the given directory and walking up the directory tree.
  Returns the port number from the first .nrepl-port file found, or nil if none found."
  [start-dir]
  (loop [dir (io/file start-dir)]
    (when dir
      (let [port-file (io/file dir ".nrepl-port")]
        (if (.exists port-file)
          (let [port (try
                       (-> (slurp port-file)
                           str/trim
                           Integer/parseInt)
                       (catch Exception e
                         ;; If we can't parse the port file, continue searching
                         nil))]
            (if port
              port
              (recur (.getParentFile dir))))
          (recur (.getParentFile dir)))))))

(defn get-env-var [var-name]
  (System/getenv var-name))

(defn resolve-host [cli-host]
  (or cli-host (get-env-var "BREPL_HOST") "localhost"))

(defn resolve-port 
  "Resolve the nREPL port from various sources.
  Priority: CLI arg > .nrepl-port file > BREPL_PORT env var
  When file-path is provided (for -f option), searches for .nrepl-port 
  starting from the file's directory and walking up the tree."
  ([cli-port] (resolve-port cli-port nil))
  ([cli-port file-path]
   (or cli-port
       (if file-path
         ;; For -f option: search from file's directory upward
         (let [file (io/file file-path)
               parent-dir (.getParentFile file)]
           (when parent-dir
             (find-nrepl-port-in-parents parent-dir)))
         ;; For -e and -m options: use current directory
         (read-nrepl-port))
       (when-let [env-port (get-env-var "BREPL_PORT")]
         (Integer/parseInt env-port)))))

(defn validate-args [opts]
  (when (:help opts)
    (print-help)
    (System/exit 0))

  (when (:version opts)
    (println (str "brepl " version))
    (System/exit 0))

  (let [has-expr (contains? opts :e)
        has-file (contains? opts :f)
        has-message (contains? opts :m)
        option-count (count (filter identity [has-expr has-file has-message]))]
    (cond
      (> option-count 1)
      (do (println "Error: Cannot specify multiple options (-e, -f, -m) together")
          (println)
          (print-help)
          (System/exit 1))

      (= option-count 0)
      (do (println "Error: Must specify one of -e EXPR, -f FILE, or -m MESSAGE")
          (println)
          (print-help)
          (System/exit 1))

      (and has-file (not (.exists (io/file (:f opts)))))
      (do (println "Error: File does not exist:" (:f opts))
          (System/exit 1)))))

;; nREPL client implementation
(defn bytes->str [x]
  (if (bytes? x) (String. x) x))

(defn convert-bytes-in-map [m]
  (walk/postwalk
   (fn [x]
     (if (bytes? x)
       (String. x)
       x))
   m))

(defn send-eval-message [host port code opts]
  (let [socket (Socket. host port)
        out (.getOutputStream socket)
        in (PushbackInputStream. (.getInputStream socket))
        msg-id (str (System/currentTimeMillis))
        msg {"op" "eval"
             "code" (str code)
             "id" msg-id}]

    ;; Print client message in verbose mode
    (when (:verbose opts)
      (pp/pprint msg))

    (bencode/write-bencode out msg)
    (.flush out)

    ;; Collect all response messages until we get "done" status
    (loop [responses []]
      (let [response (bencode/read-bencode in)
            ;; Print each response immediately in verbose mode
            _ (when (:verbose opts)
                (pp/pprint (convert-bytes-in-map response)))
            status (get response "status")
            status-strs (when status
                          (if (coll? status)
                            (map #(if (bytes? %) (String. %) %) status)
                            [(if (bytes? status) (String. status) status)]))]
        (if (and status-strs (some #(= % "done") status-strs))
          (do
            (.close socket)
            (if (:verbose opts)
              responses ; Return empty since we already printed
              (conj responses response)))
          (recur (if (:verbose opts)
                   responses ; Don't accumulate in verbose mode
                   (conj responses response))))))))

(defn process-eval-responses [responses]
  (let [combined {:out []
                  :err []
                  :values []
                  :ex nil
                  :status []}]
    (reduce (fn [acc resp]
              (cond-> acc
                (get resp "out") (update :out conj (bytes->str (get resp "out")))
                (get resp "err") (update :err conj (bytes->str (get resp "err")))
                (get resp "value") (update :values conj (bytes->str (get resp "value")))
                (get resp "ex") (assoc :ex (bytes->str (get resp "ex")))
                (get resp "status") (update :status into
                                            (map bytes->str
                                                 (if (coll? (get resp "status"))
                                                   (get resp "status")
                                                   [(get resp "status")])))))
            combined
            responses)))

(defn format-hook-response [processed has-error?]
  (if has-error?
    ;; Format error response with detailed information
    (let [;; Build comprehensive error details
          error-parts (cond-> []
                        ;; Include exception info
                        (:ex processed)
                        (conj (str "Exception: " (:ex processed)))

                        ;; Include stderr output
                        (seq (:err processed))
                        (concat (map #(str "Error: " %) (:err processed))))

          ;; Include stdout if no other error info
          error-parts (if (and (empty? error-parts) (seq (:out processed)))
                        (concat error-parts (map #(str "Output: " %) (:out processed)))
                        error-parts)

          ;; Create main error message
          error-msg (if (seq error-parts)
                      (str/join " | " error-parts)
                      "Evaluation error occurred")

          ;; Build detailed reason with all available info
          reason-parts (cond-> []
                         (:ex processed)
                         (conj (:ex processed))

                         (seq (:err processed))
                         (concat (:err processed))

                         (seq (:out processed))
                         (concat (:out processed)))

          reason-msg (if (seq reason-parts)
                       (str "Code evaluation failed:\n" (str/join "\n" reason-parts))
                       "Code contains errors that must be fixed")]

      {:continue true
       :stopReason error-msg
       :suppressOutput true
       :decision "block"
       :reason reason-msg})

    ;; Success response
    {:continue true
     :suppressOutput true}))

(defn eval-expression [host port code opts]
  (try
    (let [responses (send-eval-message host port code opts)]

      ;; Process responses to check for errors even in verbose mode
      (let [processed (process-eval-responses responses)
            has-error? (or (:ex processed)
                           (some #{"eval-error"} (flatten (:status processed))))]

        ;; If hook mode, return processed data
        (if (:hook opts)
          {:processed processed
           :has-error? has-error?}

          ;; Otherwise handle output normally
          (do
            ;; If verbose mode, responses were already printed
            (if (:verbose opts)
              has-error? ; Just return error status

              ;; Otherwise print output normally
              (do
                ;; Print stdout output
                (doseq [out (:out processed)]
                  (print out))

                ;; Print stderr output to stderr
                (doseq [err (:err processed)]
                  (binding [*out* *err*]
                    (print err)
                    (flush)))

                ;; Print evaluation results
                (doseq [val (:values processed)]
                  (println val))

                ;; Handle exceptions
                (when (:ex processed)
                  (binding [*out* *err*]
                    (println "Exception:" (:ex processed))))

                ;; Handle empty response
                (when (and (empty? (:values processed))
                           (empty? (:out processed))
                           (empty? (:err processed))
                           (not (:ex processed)))
                  (when (some #{"eval-error"} (flatten (:status processed)))
                    (binding [*out* *err*]
                      (println "Evaluation error occurred"))))

                ;; Return error status
                has-error?))))))

    (catch Exception e
      (if (:hook opts)
        {:processed {:ex (str "Connection error: " (.getMessage e))}
         :has-error? true}
        (do
          (println "Error connecting to nREPL server at" (str host ":" port))
          (println (.getMessage e))
          (System/exit 1))))))

(defn eval-file [host port file-path opts]
  (let [code (str "(load-file \"" file-path "\")")]
    (eval-expression host port code opts)))

(defn send-raw-message [host port message-str opts]
  (try
    (let [socket (Socket. host port)
          out (.getOutputStream socket)
          in (PushbackInputStream. (.getInputStream socket))
          ;; Parse the EDN message and add an ID if not present
          msg (read-string message-str)
          msg-with-id (if (get msg "id")
                        msg
                        (assoc msg "id" (str (System/currentTimeMillis))))]

      ;; Print client message in verbose mode
      (when (:verbose opts)
        (pp/pprint msg-with-id))

      (bencode/write-bencode out msg-with-id)
      (.flush out)

      ;; Collect all response messages until we get "done" status or no status
      (loop [responses []
             timeout-count 0]
        (let [response (bencode/read-bencode in)
              ;; Print each response immediately in verbose mode
              _ (when (:verbose opts)
                  (pp/pprint (convert-bytes-in-map response)))
              status (get response "status")
              status-strs (when status
                            (if (coll? status)
                              (map #(if (bytes? %) (String. %) %) status)
                              [(if (bytes? status) (String. status) status)]))]
          (cond
            ;; If we have a done status, we're finished
            (and status-strs (some #(= % "done") status-strs))
            (do
              (.close socket)
              (if (:verbose opts)
                responses ; Return empty since we already printed
                (conj responses response)))

            ;; For operations that don't return "done", stop after reasonable response count
            (> timeout-count 10)
            (do
              (.close socket)
              responses)

            ;; Otherwise keep reading
            :else
            (recur (if (:verbose opts)
                     responses ; Don't accumulate in verbose mode
                     (conj responses response))
                   (inc timeout-count))))))

    (catch Exception e
      (println "Error sending message to nREPL server at" (str host ":" port))
      (println (.getMessage e))
      (System/exit 1))))

(defn process-raw-message [host port message-str opts]
  (let [responses (send-raw-message host port message-str opts)]
    ;; In verbose mode, everything was already printed
    (when-not (:verbose opts)
      ;; Pretty print each response
      (doseq [resp responses]
        (pp/pprint (convert-bytes-in-map resp))))))

(defn -main [& args]
  (let [opts (cli/parse-opts args {:spec cli-spec})
        host (resolve-host (:h opts))
        port (resolve-port (:p opts) (:f opts))]

    (validate-args opts)

    (when-not port
      (println "Error: No port specified, no .nrepl-port file found, and BREPL_PORT not set")
      (System/exit 1))

    (let [result (cond
                   (:e opts) (eval-expression host port (:e opts) opts)
                   (:f opts) (eval-file host port (:f opts) opts)
                   (:m opts) (do (process-raw-message host port (:m opts) opts)
                                 false))] ; raw messages don't track errors

      ;; Handle hook mode
      (if (:hook opts)
        (let [{:keys [processed has-error?]} result
              hook-response (format-hook-response processed has-error?)]
          (println (json/generate-string hook-response))
          (System/exit (if has-error? 2 0)))

        ;; Regular mode - for backward compatibility, use exit code 2 for eval errors
        (when result
          (System/exit 2))))))

;; Lazy load hook modules
(defn ensure-hook-modules-loaded []
  (when (not @hook-modules-loaded)
    (let [script-dir (-> (System/getProperty "babashka.file")
                         (clojure.java.io/file)
                         (.getParentFile)
                         (.getAbsolutePath))]
      (load-file (str script-dir "/lib/hook_utils.clj"))
      (require '[brepl.lib.hook-utils :as hook-utils])
      (load-file (str script-dir "/lib/validator.clj"))
      (require '[brepl.lib.validator :as validator])
      (load-file (str script-dir "/lib/backup.clj"))
      (require '[brepl.lib.backup :as backup])
      (load-file (str script-dir "/lib/installer.clj"))
      (require '[brepl.lib.installer :as installer]))
    (reset! hook-modules-loaded true)))

;; Hook subcommand handlers
(defn handle-hook-subcommand [subcommand args]
  (case subcommand
    "validate"
    (if (< (count args) 2)
      (do (println (json/generate-string {:continue true :decision "block" :stopReason "invalid_arguments" :reason "Usage: brepl hook validate <file-path> <new-content>"}))
          (System/exit 1))
      (do (ensure-hook-modules-loaded)
          (let [file-path (first args)
                content (second args)
                clojure-file? (ns-resolve (the-ns 'brepl.lib.validator) 'clojure-file?)
                delimiter-error? (ns-resolve (the-ns 'brepl.lib.validator) 'delimiter-error?)]
            (if-not (clojure-file? file-path)
              (do (println (json/generate-string {:continue true :decision "allow" :suppressOutput true}))
                  (System/exit 0)))
            (let [error (delimiter-error? content)
                  format-msg (ns-resolve (the-ns 'brepl.lib.validator) 'format-error-message)
                  auto-fix (ns-resolve (the-ns 'brepl.lib.validator) 'auto-fix-brackets)
                  create-backup (ns-resolve (the-ns 'brepl.lib.backup) 'create-backup)
                  session-id (System/getenv "SESSION_ID")]
              (if error
                (let [fixed (auto-fix content)]
                  (if fixed
                    (do (when (and session-id create-backup)
                          (create-backup file-path session-id))
                        (println (json/generate-string {:continue true :decision "allow" :suppressOutput true :correction fixed}))
                        (System/exit 0))
                    (do (println (json/generate-string {:continue true :decision "block" :stopReason "syntax_error" :reason (str "Syntax error in " file-path ": " (format-msg error file-path))}))
                        (System/exit 1))))
                (do (when (and session-id create-backup)
                      (create-backup file-path session-id))
                    (println (json/generate-string {:continue true :decision "allow" :suppressOutput true}))
                    (System/exit 0)))))))

    "eval"
    (if (empty? args)
      (do (println (json/generate-string {:continue true :decision "block" :stopReason "invalid_arguments" :reason "Usage: brepl hook eval <file-path>"}))
          (System/exit 1))
      (do (ensure-hook-modules-loaded)
          (let [file-path (first args)
                delimiter-error? (ns-resolve (the-ns 'brepl.lib.validator) 'delimiter-error?)
                restore-backup (ns-resolve (the-ns 'brepl.lib.backup) 'restore-backup)
                session-id (System/getenv "SESSION_ID")
                host (resolve-host nil)
                port (resolve-port nil file-path)]
            ;; Load file and validate syntax first
            (if (not (.exists (clojure.java.io/file file-path)))
              (do (println (json/generate-string {:continue true :decision "block" :stopReason "file_not_found" :reason (str "File not found: " file-path)}))
                  (System/exit 1))
              ;; File exists, validate syntax
              (let [file-content (slurp file-path)
                    error (delimiter-error? file-content)]
                (if error
                  ;; Syntax error - restore backup
                  (do (when (and session-id restore-backup)
                        (restore-backup session-id file-path))
                      (println (json/generate-string {:continue true :decision "block" :stopReason "syntax_error" :reason (str "Syntax error in " file-path ": " (:message error))}))
                      (System/exit 1))
                  ;; Valid syntax - try nREPL evaluation if port available
                  (if port
                    ;; Evaluate via nREPL
                    (let [result (eval-file host port file-path {:hook true})]
                      (if (:has-error? result)
                        ;; Eval error - warn but allow by default
                        (let [ex-msg (get-in result [:processed :ex] "Evaluation error")]
                          (println (json/generate-string {:continue true :decision "allow" :suppressOutput true :warning (str "Evaluation warning in " file-path ": " ex-msg)}))
                          (System/exit 0))
                        ;; Success
                        (do (println (json/generate-string {:continue true :decision "allow" :suppressOutput true}))
                            (System/exit 0))))
                    ;; No nREPL port - gracefully degrade, just allow
                    (do (println (json/generate-string {:continue true :decision "allow" :suppressOutput true}))
                        (System/exit 0)))))))))

    "install"
    (do (ensure-hook-modules-loaded)
        (let [opts (cli/parse-opts args {:spec {:strict-eval {:coerce :boolean}
                                                 :parinfer {:coerce :boolean}}})
              install-result (ns-resolve (the-ns 'brepl.lib.installer) 'install-hooks)
              check-result (ns-resolve (the-ns 'brepl.lib.installer) 'check-status)]
          (install-result opts)
          (let [status (check-result)]
            (println "Installing Claude Code hooks...")
            (println "Hooks installed successfully.")
            (println)
            (when (:installed status)
              (println "Settings file:" (:settings-path status)))
            (System/exit 0))))

    "uninstall"
    (do (ensure-hook-modules-loaded)
        (let [uninstall-result (ns-resolve (the-ns 'brepl.lib.installer) 'uninstall-hooks)]
          (uninstall-result)
          (println "Removing Claude Code hooks...")
          (println "Hooks uninstalled successfully.")
          (System/exit 0)))

    "session-end"
    (if (empty? args)
      (do (println "Error: session-id required")
          (System/exit 1))
      (do (ensure-hook-modules-loaded)
          (let [session-id (first args)
                cleanup (ns-resolve (the-ns 'brepl.lib.backup) 'cleanup-session)]
            (cleanup session-id)
            (System/exit 0))))

    ;; Help for hook subcommands
    (do
      (println "brepl hook - Claude Code integration commands")
      (println)
      (println "USAGE:")
      (println "    brepl hook validate <file> <content>")
      (println "    brepl hook eval <file>")
      (println "    brepl hook install [--strict-eval] [--parinfer]")
      (println "    brepl hook uninstall")
      (println "    brepl hook session-end <session-id>")
      (println)
      (println "SUBCOMMANDS:")
      (println "    validate       Validate Clojure file syntax before edit")
      (println "    eval           Evaluate file and check for runtime errors")
      (println "    install        Install hooks in .claude/settings.local.json")
      (println "    uninstall      Remove hooks from .claude/settings.local.json")
      (println "    session-end    Clean up session backup files")
      (System/exit (if (and subcommand (not (contains? #{"validate" "eval" "install" "uninstall" "session-end"} subcommand))) 1 0)))))

(defn -main-hook [args]
  "Handle brepl hook subcommands"
  (if (empty? args)
    (handle-hook-subcommand nil nil)
    (let [subcommand (first args)
          remaining-args (rest args)]
      (handle-hook-subcommand subcommand remaining-args))))

(when (= *file* (System/getProperty "babashka.file"))
  (if (and (> (count *command-line-args*) 0)
           (= "hook" (first *command-line-args*)))
    (-main-hook (rest *command-line-args*))
    (apply -main *command-line-args*)))
