(ns ajet.chat.auth-gw.proxy
  "Reverse proxy — forwards authenticated requests to internal services.

   Uses babashka.http-client for synchronous proxying and http-kit's
   async channel for SSE pass-through streaming."
  (:require [babashka.http-client :as http]
            [clojure.string :as str]
            [clojure.tools.logging :as log]
            [org.httpkit.server :as hk])
  (:import [java.io InputStream]))

;;; ---------------------------------------------------------------------------
;;; Service resolution
;;; ---------------------------------------------------------------------------

(defn- service-url
  "Build the base URL for a service from config."
  [{:keys [host port] :or {host "localhost"}}]
  (str "http://" host ":" port))

(defn resolve-target
  "Determine the target service and rewritten path for a request URI.

   Returns [service-key path] or nil if no match.

   Service routing:
   - /app/* /sse/* /web/* / -> :web-sm
   - /api/*                 -> :api
   - /tui/* /tui/sse/*      -> :tui-sm"
  [uri]
  (cond
    ;; Web SM routes
    (= uri "/")
    [:web-sm "/"]

    (= uri "/setup")
    [:web-sm "/setup"]

    (= uri "/app")
    [:web-sm "/app"]

    (str/starts-with? uri "/app/")
    [:web-sm uri]

    (str/starts-with? uri "/sse/")
    [:web-sm uri]

    (str/starts-with? uri "/web/")
    [:web-sm uri]

    ;; API routes
    (str/starts-with? uri "/api/")
    [:api uri]

    ;; TUI SM routes
    (str/starts-with? uri "/tui/")
    [:tui-sm uri]

    :else nil))

(defn- build-target-url
  "Build the full target URL from service config and request."
  [services service-key path query-string]
  (let [base (service-url (get services service-key))]
    (if (str/blank? query-string)
      (str base path)
      (str base path "?" query-string))))

;;; ---------------------------------------------------------------------------
;;; Header manipulation
;;; ---------------------------------------------------------------------------

(def ^:private stripped-request-headers
  "Headers to strip from proxied requests (prevent spoofing)."
  #{"x-user-id" "x-user-role" "x-community-id" "x-trace-id"
    "x-forwarded-for" "x-forwarded-proto" "x-forwarded-host"
    "host" "connection" "keep-alive" "transfer-encoding"
    "te" "trailer" "upgrade" "content-length"})

(defn- clean-request-headers
  "Strip auth/hop-by-hop headers and normalize values for java.net.http."
  [headers]
  (into {}
        (comp (remove (fn [[k _]] (contains? stripped-request-headers (str/lower-case k))))
              (map (fn [[k v]] [k (str/replace (str v) #"\r?\n" ", ")])))
        headers))

(defn- inject-proxy-headers
  "Add auth and tracing headers for the upstream service."
  [headers {:keys [user-id user-role community-id trace-id remote-addr]}]
  (cond-> headers
    user-id      (assoc "X-User-Id" user-id)
    user-role    (assoc "X-User-Role" user-role)
    community-id (assoc "X-Community-Id" community-id)
    trace-id     (assoc "X-Trace-Id" trace-id)
    remote-addr  (assoc "X-Forwarded-For" remote-addr)))

(defn- is-sse-request?
  "Check if the request is for an SSE endpoint."
  [uri]
  (or (str/starts-with? uri "/sse/")
      (str/starts-with? uri "/tui/sse/")))

;;; ---------------------------------------------------------------------------
;;; Synchronous proxy (non-SSE)
;;; ---------------------------------------------------------------------------

(def ^:private proxy-client
  "HTTP client that never follows redirects — upstream redirects must be
   passed through to the client untouched."
  (http/client {:follow-redirects :never}))

(defn- proxy-sync
  "Forward a non-SSE request synchronously via babashka.http-client."
  [target-url method headers body]
  (try
    (let [opts (cond-> {:uri     target-url
                        :method  method
                        :headers headers
                        :throw   false
                        :timeout 30000
                        :client  proxy-client}
                 body (assoc :body body))
          resp (http/request opts)]
      {:status  (:status resp)
       :headers (-> (:headers resp)
                    (dissoc "transfer-encoding" "connection")
                    ;; Ensure headers are string->string
                    (->> (into {}
                               (map (fn [[k v]]
                                      [(name k)
                                       (if (sequential? v) (str/join ", " v) (str v))])))))
       :body    (:body resp)})
    (catch java.io.IOException e
      ;; Java 21's HttpClient rejects 204/304 responses that include
      ;; a Content-Length header (strict HTTP compliance). Detect and
      ;; return the correct no-content status directly.
      (let [msg (str (.getMessage e))]
        (if-let [[_ status] (re-find #"content length header with (\d+) response" msg)]
          {:status  (Integer/parseInt status)
           :headers {}
           :body    nil}
          (do (log/error e "Proxy IO error to" target-url)
              {:status  502
               :headers {"Content-Type" "text/plain"}
               :body    "Bad Gateway — upstream service unavailable"}))))
    (catch Exception e
      (log/error e "Proxy error to" target-url)
      {:status  502
       :headers {"Content-Type" "text/plain"}
       :body    "Bad Gateway — upstream service unavailable"})))

;;; ---------------------------------------------------------------------------
;;; SSE streaming proxy
;;; ---------------------------------------------------------------------------

(defn- proxy-sse
  "Forward an SSE request using http-kit async channel for streaming.

   Opens an http-kit async response, makes a streaming request to the
   upstream, and pipes chunks through without buffering."
  [target-url headers hk-channel]
  (try
    (let [resp (http/get target-url
                 {:headers (assoc headers "Accept" "text/event-stream")
                  :throw   false
                  :as      :stream})]
      (if (= 200 (:status resp))
        (do
          ;; Send initial response headers
          (hk/send! hk-channel
                    {:status  200
                     :headers {"Content-Type"  "text/event-stream"
                               "Cache-Control" "no-cache, no-store"
                               "Connection"    "keep-alive"
                               "X-Accel-Buffering" "no"}}
                    false)
          ;; Stream body in a background thread
          (future
            (try
              (let [^InputStream is (:body resp)
                    buf             (byte-array 4096)]
                (loop []
                  (let [n (.read is buf)]
                    (when (and (pos? n) (hk/open? hk-channel))
                      (hk/send! hk-channel
                                (String. buf 0 n "UTF-8")
                                false)
                      (recur))))
                ;; Upstream closed
                (hk/close hk-channel))
              (catch Exception e
                (when-not (instance? java.io.IOException e)
                  (log/warn e "SSE stream error"))
                (hk/close hk-channel)))))
        ;; Upstream error
        (do
          (hk/send! hk-channel
                    {:status  (:status resp)
                     :headers {"Content-Type" "text/plain"}
                     :body    "Upstream SSE connection failed"}
                    true))))
    (catch Exception e
      (log/error e "SSE proxy connection error to" target-url)
      (hk/send! hk-channel
                {:status  502
                 :headers {"Content-Type" "text/plain"}
                 :body    "Bad Gateway — upstream service unavailable"}
                true))))

;;; ---------------------------------------------------------------------------
;;; Public API
;;; ---------------------------------------------------------------------------

(defn proxy-request
  "Proxy a Ring request to the appropriate internal service.

   services: map from config {:api {:host .. :port ..} :web-sm ... :tui-sm ...}
   auth-info: map of {:user-id :user-role :community-id :trace-id :remote-addr}
   request: Ring request map

   Returns a Ring response map (or uses http-kit async for SSE)."
  [services auth-info request]
  (let [uri          (:uri request)
        resolution   (resolve-target uri)]
    (if-not resolution
      {:status  404
       :headers {"Content-Type" "text/plain"}
       :body    "Not Found"}
      (let [[service-key path] resolution
            service-cfg (get services service-key)]
        (if-not service-cfg
          (do (log/error "No service config for" service-key)
              {:status  502
               :headers {"Content-Type" "text/plain"}
               :body    "Bad Gateway — service not configured"})
          (let [target-url (build-target-url services service-key path (:query-string request))
                headers    (-> (:headers request)
                               clean-request-headers
                               (inject-proxy-headers auth-info))]
            (if (is-sse-request? uri)
              ;; SSE: use http-kit async channel
              (hk/with-channel request hk-channel
                (hk/on-close hk-channel
                  (fn [_status]
                    (log/debug "SSE client disconnected from" uri)))
                (proxy-sse target-url headers hk-channel))
              ;; Non-SSE: synchronous proxy
              ;; Use :raw-body (buffered bytes) if available, otherwise :body
              ;; :raw-body is set by wrap-buffer-body so the proxy has the
              ;; original body even after wrap-params consumed the InputStream
              (let [body (or (when-let [raw (:raw-body request)]
                               (when (pos? (alength raw)) raw))
                             (:body request))]
                (proxy-sync target-url
                            (:request-method request)
                            headers
                            body)))))))))
