(ns ajet.chat.auth-gw.rate-limit
  "In-memory token bucket rate limiter.

   Each bucket is keyed by [category identity-key] and tracks remaining tokens
   plus the last refill timestamp. Expired/stale buckets are cleaned up
   periodically by a background thread."
  (:require [clojure.tools.logging :as log])
  (:import [java.time Instant Duration]))

;;; ---------------------------------------------------------------------------
;;; Configuration
;;; ---------------------------------------------------------------------------

(def ^:private rate-rules
  "Rate limit rules keyed by category keyword.
   :capacity  — max tokens (requests) in the window
   :refill-ms — window size in milliseconds"
  {:auth-callback  {:capacity 10  :refill-ms 60000}   ; 10/min per IP
   :api-write      {:capacity 60  :refill-ms 60000}   ; 60/min per user
   :api-read       {:capacity 120 :refill-ms 60000}   ; 120/min per user
   :webhook        {:capacity 30  :refill-ms 60000}   ; 30/min per webhook
   :sse            {:capacity 5   :refill-ms 60000}})  ; 5/min per user

(def ^:private stale-threshold-ms
  "Remove buckets that haven't been touched in 10 minutes."
  600000)

;;; ---------------------------------------------------------------------------
;;; Bucket state
;;; ---------------------------------------------------------------------------

(defn make-limiter
  "Create a new rate limiter state atom. Returns the atom."
  []
  (atom {}))

(defn- now-ms []
  (System/currentTimeMillis))

(defn- refill-bucket
  "Refill tokens based on elapsed time since last refill."
  [{:keys [tokens last-refill-ms] :as bucket} {:keys [capacity refill-ms]}]
  (let [elapsed (- (now-ms) last-refill-ms)]
    (if (>= elapsed refill-ms)
      ;; Full window elapsed — reset to capacity
      (assoc bucket
        :tokens capacity
        :last-refill-ms (now-ms))
      ;; Partial refill — add proportional tokens
      (let [fraction     (/ (double elapsed) (double refill-ms))
            added        (* fraction capacity)
            new-tokens   (min capacity (+ tokens added))]
        (assoc bucket
          :tokens new-tokens
          :last-refill-ms (now-ms))))))

(defn- get-or-create-bucket
  "Get an existing bucket or create a fresh one at capacity."
  [buckets bucket-key rule]
  (or (get buckets bucket-key)
      {:tokens         (:capacity rule)
       :last-refill-ms (now-ms)}))

(defn check-rate-limit!
  "Check and consume a token from the bucket for [category identity-key].

   Returns {:allowed? true} if the request is permitted, or
   {:allowed? false :retry-after-ms N} if rate limited.

   category: one of :auth-callback :api-write :api-read :webhook :sse
   identity-key: string (IP address, user-id, or webhook-id)"
  [limiter-atom category identity-key]
  (let [rule (get rate-rules category)]
    (if-not rule
      ;; Unknown category — allow (no rule defined)
      {:allowed? true}
      (let [bucket-key [category identity-key]
            result     (atom nil)]
        (swap! limiter-atom
               (fn [buckets]
                 (let [bucket   (get-or-create-bucket buckets bucket-key rule)
                       refilled (refill-bucket bucket rule)
                       tokens   (:tokens refilled)]
                   (if (>= tokens 1.0)
                     ;; Allow — consume a token
                     (do (reset! result {:allowed? true})
                         (assoc buckets bucket-key
                                (assoc refilled
                                  :tokens (dec tokens)
                                  :last-access-ms (now-ms))))
                     ;; Deny — calculate retry-after
                     (let [deficit       (- 1.0 tokens)
                           refill-rate   (/ (double (:capacity rule))
                                           (double (:refill-ms rule)))
                           retry-after   (long (Math/ceil (/ deficit refill-rate)))]
                       (reset! result {:allowed?       false
                                       :retry-after-ms retry-after})
                       (assoc buckets bucket-key
                              (assoc refilled :last-access-ms (now-ms))))))))
        @result))))

;;; ---------------------------------------------------------------------------
;;; Cleanup
;;; ---------------------------------------------------------------------------

(defn cleanup-stale!
  "Remove buckets that haven't been accessed in over `stale-threshold-ms`."
  [limiter-atom]
  (let [cutoff (- (now-ms) stale-threshold-ms)
        removed (atom 0)]
    (swap! limiter-atom
           (fn [buckets]
             (let [active (into {}
                                (filter (fn [[_ b]]
                                          (> (get b :last-access-ms 0) cutoff)))
                                buckets)]
               (reset! removed (- (count buckets) (count active)))
               active)))
    (when (pos? @removed)
      (log/debug "Rate limiter cleanup: removed" @removed "stale buckets"))))

(defn start-cleanup-task!
  "Start a background thread that periodically cleans up stale buckets.
   Returns the future (cancel with `future-cancel`)."
  [limiter-atom]
  (future
    (log/info "Rate limiter cleanup task started")
    (try
      (loop []
        (Thread/sleep 60000) ;; every minute
        (cleanup-stale! limiter-atom)
        (recur))
      (catch InterruptedException _
        (log/info "Rate limiter cleanup task stopped"))
      (catch Exception e
        (log/error e "Rate limiter cleanup task error")))))

;;; ---------------------------------------------------------------------------
;;; Route classification
;;; ---------------------------------------------------------------------------

(defn classify-request
  "Determine the rate limit category and identity key for a request.

   Returns [category identity-key] or nil if no rate limit applies.

   request: Ring request map (expects :uri, :request-method, and optionally
            ::user-id, ::remote-addr, ::webhook-id keys set by auth middleware)."
  [request]
  (let [uri    (:uri request)
        method (:request-method request)]
    (cond
      ;; OAuth callback — rate limit by IP
      (and (= method :post)
           (re-matches #"/auth/callback/.*" uri))
      [:auth-callback (or (::remote-addr request)
                          (get-in request [:headers "x-forwarded-for"])
                          (:remote-addr request)
                          "unknown")]

      ;; Webhook incoming — rate limit by webhook ID
      (and (= method :post)
           (re-matches #"/api/webhooks/.*/incoming" uri))
      [:webhook (or (::webhook-id request) "unknown")]

      ;; SSE connections — rate limit by user
      (and (= method :get)
           (or (re-matches #"/sse/.*" uri)
               (re-matches #"/tui/sse/.*" uri)))
      [:sse (or (::user-id request) "anonymous")]

      ;; API writes — rate limit by user
      (and (= method :post)
           (re-matches #"/api/.*" uri))
      [:api-write (or (::user-id request) "anonymous")]

      ;; API reads — rate limit by user
      (and (= method :get)
           (re-matches #"/api/.*" uri))
      [:api-read (or (::user-id request) "anonymous")]

      ;; No rate limit for other routes
      :else nil)))
