(ns ird.p2p.crypto
  "Every derivation of one room (PROTOCOL.md): the room id, per-sender
  session keys, the mailbox room id and the room-wide at-rest key — all
  from HKDF-SHA256(ikm: secret, salt: UTF8(appSalt), info: label) — plus
  the session sealer (IV = 4 random epoch bytes ‖ 8-byte big-endian
  counter) and the receiver's anti-replay window.

  Values at this surface are strings and plain data: secrets, ids and keys
  as URL-safe base64, IVs and ciphertexts as standard base64, hex where
  vectors speak hex. Derivations are asynchronous (WebCrypto) and deliver
  on a promise-chan; sealers and windows are immutable maps threaded by
  the caller (`room-crypto` wraps its own in atoms)."
  (:require [clojure.string :as str]
            [cljs.core.async :as async :refer [go <!]]
            [ird.p2p.async :as a :refer [go-try <?]]
            [ird.p2p.errors :as err]
            [ird.p2p.impl.bytes :as bytes]
            [ird.p2p.impl.webcrypto :as wc]))

;; ---- identifiers -------------------------------------------------------------

(def room-id-length 43)
(def peer-id-length 32)

(defn valid-peer-id?
  "16 bytes as 32 lowercase hex characters."
  [s]
  (and (string? s) (boolean (re-matches #"[0-9a-f]{32}" s))))

(defn valid-room-id?
  "43 URL-safe base64 characters (32 bytes)."
  [s]
  (and (string? s) (boolean (re-matches #"[A-Za-z0-9_-]{43}" s))))

(defn secret-bytes
  "Number of bytes SECRET (URL-safe base64) decodes to; nil when malformed."
  [secret]
  (bytes/b64url-byte-length secret))

(defn valid-secret?
  "A room secret: URL-safe base64 of at least 16 bytes."
  [secret]
  (let [n (secret-bytes secret)]
    (and n (>= n 16))))

(defn check-secret
  "Return SECRET, or throw a config error when it is not a valid room secret."
  [secret]
  (cond
    (not (bytes/valid-b64url? secret))
    (throw (err/config-error "Room secret is not URL-safe base64" {:secret-length (count (str secret))}))
    (< (secret-bytes secret) 16)
    (throw (err/config-error "room secret too short" {:bytes (secret-bytes secret)}))
    :else secret))

(defn random-secret
  "32 fresh random bytes as URL-safe base64 (43 characters)."
  []
  (bytes/random-b64url 32))

(defn random-peer-id
  "16 fresh random bytes as 32 lowercase hex characters."
  []
  (bytes/random-hex 16))

;; ---- HKDF derivations ---------------------------------------------------------

(defn hkdf
  "HKDF(info): 32 bytes as URL-safe base64. Channel."
  [secret app-salt info]
  (wc/hkdf secret app-salt info))

(defn room-id
  "`b64url(HKDF(\"roomid\"))` — the relay room id (43 characters). Channel."
  [secret app-salt]
  (hkdf secret app-salt "roomid"))

(defn mailbox-room-id
  "`b64url(HKDF(\"mailbox|room\"))` — the mailbox room id; never equal to the
  relay room id, so the node cannot correlate them. Channel."
  [secret app-salt]
  (hkdf secret app-salt "mailbox|room"))

(defn room-topic
  "Gossip topic of the reference implementation (vector parity). Channel."
  [secret app-salt]
  (go (str "sueta/2/" (<! (hkdf secret app-salt "topic|room")))))

(defn db-name
  "Database name of the reference implementation (vector parity). Channel."
  [secret app-salt]
  (hkdf secret app-salt "dbname"))

(defn sig-key-info [peer-id] (str "signal|" peer-id))
(defn msg-key-info [peer-id] (str "msg|" peer-id))
(def mailbox-key-info "mailbox-blob")

;; ---- AES-GCM ---------------------------------------------------------------------

(defn aes-seal
  "ciphertext‖tag (standard base64) of PT-B64 under KEY (URL-safe base64),
  IV-B64 and AAD (string). Channel."
  [key iv-b64 aad pt-b64]
  (wc/aes-gcm-seal key iv-b64 aad pt-b64))

(defn aes-open
  "Plaintext (standard base64) or nil on any failure. Channel."
  [key iv-b64 aad ct-b64]
  (wc/aes-gcm-open key iv-b64 aad ct-b64))

;; ---- sealer ------------------------------------------------------------------------

(defn make-sealer
  "A sealer state: {:epoch <8 hex chars> :counter 0}. The epoch is random
  unless given (tests and vector reproduction only)."
  ([] (make-sealer (bytes/random-hex 4)))
  ([epoch-hex]
   (when-not (and (string? epoch-hex) (= 8 (count epoch-hex)))
     (throw (err/config-error "sealer epoch is 4 bytes" {:epoch epoch-hex})))
   {:epoch epoch-hex :counter 0}))

(defn iv-hex
  "The IV of SEALER's current counter: epoch ‖ counter (big-endian, 8 bytes)."
  [{:keys [epoch counter]}]
  (str epoch (bytes/int->hex counter 16)))

(defn next-iv
  "[sealer' iv-b64]: the IV to use now and the advanced sealer."
  [sealer]
  [(update sealer :counter inc) (bytes/hex->b64 (iv-hex sealer))])

(defn take-iv!
  "Advance the sealer held in ATOM and return its IV (standard base64)."
  [sealer-atom]
  (let [iv (bytes/hex->b64 (iv-hex @sealer-atom))]
    (swap! sealer-atom update :counter inc)
    iv))

;; ---- replay window ---------------------------------------------------------------

(def window-size 64)
(def retired-epochs-kept 8)

(defn make-replay-window
  "Per remote sender and key: the current epoch, its highest counter, the
  counters seen within the last 64, and the retired epochs."
  []
  {:epoch nil :highest 0 :seen #{} :retired []})

(defn- split-iv [iv-b64]
  (when (and (bytes/valid-b64? iv-b64) (= 12 (bytes/byte-length iv-b64)))
    (let [h (bytes/b64->hex iv-b64)]
      [(subs h 0 8) (bytes/hex->int (subs h 8))])))

(defn replay-accept
  "[window' accepted?]: accepts an IV that is fresh — new epoch (unless
  retired), or a counter not yet seen and not more than 64 behind the
  highest — and records it. As in DTLS, out-of-order frames within the
  window pass; duplicates and stale counters do not."
  [{:keys [epoch highest seen retired] :as w} iv-b64]
  (if-let [[e c] (split-iv iv-b64)]
    (cond
      (nil? epoch) [(assoc w :epoch e :highest c :seen #{c}) true]
      (not= e epoch)
      (if (some #{e} retired)
        [w false]
        [(assoc w :epoch e :highest c :seen #{c}
                :retired (vec (take-last retired-epochs-kept (conj retired epoch))))
         true])
      (> c highest)
      (let [floor (- c window-size)]
        [(assoc w :highest c :seen (conj (set (filter #(> % floor) seen)) c)) true])
      (>= (- highest c) window-size) [w false]
      (contains? seen c) [w false]
      :else [(update w :seen conj c) true])
    [w false]))

;; ---- room crypto -------------------------------------------------------------------

(defn room-crypto
  "The crypto context of one session: derives the room id and this peer's
  own signal/message keys, creates the two sealers and an empty key cache.
  OPTS: :secret (URL-safe base64, ≥16 bytes), :app-salt, :peer-id (32 hex),
  :epoch (8 hex chars; tests and vectors only). Channel of the context map
  {:secret :app-salt :peer-id :room-id :keys <atom> :sig-sealer <atom>
  :msg-sealer <atom>} or a config error."
  [{:keys [secret app-salt peer-id epoch]}]
  (go-try
    (check-secret secret)
    (when-not (valid-peer-id? peer-id)
      (throw (err/config-error "peer id must be 32 lowercase hex characters" {:peer-id peer-id})))
    (when-not (string? app-salt)
      (throw (err/config-error "app salt must be a string" {:app-salt app-salt})))
    (let [rid (<! (room-id secret app-salt))
          k-sig (<! (hkdf secret app-salt (sig-key-info peer-id)))
          k-msg (<! (hkdf secret app-salt (msg-key-info peer-id)))]
      {:secret secret
       :app-salt app-salt
       :peer-id peer-id
       :room-id rid
       :keys (atom {(sig-key-info peer-id) k-sig (msg-key-info peer-id) k-msg})
       :sig-sealer (atom (if epoch (make-sealer epoch) (make-sealer)))
       :msg-sealer (atom (if epoch (make-sealer epoch) (make-sealer)))})))

(def key-cache-limit 512)

(defn key-for
  "The derived key (URL-safe base64) for INFO in context RC, cached; remote
  sender ids are attacker-mintable, so the cache is bounded. Channel."
  [{:keys [secret app-salt keys]} info]
  (go
    (or (get @keys info)
        (let [k (<! (hkdf secret app-salt info))]
          (swap! keys (fn [m] (assoc (if (>= (count m) key-cache-limit) {} m) info k)))
          k))))

(defn sig-key
  "K_sig(P): the key for envelopes sent by PEER over the relay. Channel."
  [rc peer]
  (key-for rc (sig-key-info peer)))

(defn msg-key
  "K_msg(P): the key for envelopes sent by PEER over a data channel. Channel."
  [rc peer]
  (key-for rc (msg-key-info peer)))

(defn mailbox-key
  "K_mbx: the room-wide key for mailbox payloads. Channel."
  [rc]
  (key-for rc mailbox-key-info))

(defn random-iv
  "A random 12-byte IV (standard base64) for at-rest envelopes."
  []
  (bytes/random-b64 12))