(ns ird.p2p.envelope
  "The bytes inside the transport (PROTOCOL.md): sealed envelopes — the
  JSON envelope {\"v\":1,\"iv\":…,\"ct\":…} of the relay path and data-channel
  text messages, the binary envelope 0x01‖iv‖ct‖tag of data-channel bytes,
  the at-rest envelope 0x02‖iv‖ct‖tag of the mailbox — with their AAD
  strings; the inner messages they carry (hello, offer, answer, ice,
  restart, msg, bin, chunk, ping); and chunking of large messages.

  Opening never throws and never distinguishes tampering from
  \"not for me\": every failure is nil, and the caller drops silently.
  Bytes are standard base64 strings; envelopes and inner messages are maps
  whose key order is the wire order."
  (:require [clojure.string :as str]
            [cljs.core.async :as async :refer [go <!]]
            [ird.p2p.crypto :as crypto]
            [ird.p2p.impl.bytes :as bytes]
            [ird.p2p.impl.json :as json]))

;; ---- AADs -----------------------------------------------------------------------

(defn sig-aad [room-id from to] (str "v1|" room-id "|" from "|" to "|sig"))
(defn msg-aad [room-id from] (str "v1|" room-id "|" from "|msg"))
(defn at-rest-aad [room-id suffix] (str "v2|" room-id "|" suffix))

;; ---- JSON envelope -----------------------------------------------------------------

(def version 1)

(defn envelope
  "The wire map {:v 1 :iv <b64> :ct <b64>}, keys in exactly that order."
  [iv-b64 ct-b64]
  {:v version :iv iv-b64 :ct ct-b64})

(defn parse-envelope
  "Strict decoding of a parsed JSON value: `v` must be exactly 1 and `iv`,
  `ct` well-formed base64 strings; nil otherwise."
  [v]
  (when (and (map? v) (= 1 (:v v))
             (bytes/valid-b64? (:iv v)) (bytes/valid-b64? (:ct v)))
    (envelope (:iv v) (:ct v))))

(defn envelope->json
  "The envelope's wire text."
  [env]
  (json/stringify (select-keys env [:v :iv :ct])))

(defn json->envelope
  "Envelope from wire text; nil when not an envelope."
  [text]
  (parse-envelope (json/parse text)))

;; ---- binary envelopes -------------------------------------------------------------

(def prefix-session 1)
(def prefix-at-rest 2)
(def minimum-length (+ 1 12 16))

(defn- binary [prefix iv-b64 ct-b64]
  (bytes/concat-b64 (bytes/bytes->b64 [prefix]) iv-b64 ct-b64))

(defn split-binary
  "[iv-b64 ct-b64] of a binary envelope with PREFIX, or nil when too short
  or not that prefix."
  [b64 prefix]
  (when (and (bytes/valid-b64? b64)
             (>= (bytes/byte-length b64) minimum-length)
             (= prefix (bytes/byte-at b64 0)))
    [(bytes/slice-b64 b64 1 13) (bytes/slice-b64 b64 13)]))

;; ---- sealing as the local peer --------------------------------------------------------

(defn- utf8-b64 [s] (bytes/utf8->b64 s))

(defn seal-signal
  "Seal INNER (a map) for TO under this peer's signal key. Channel of the
  envelope map."
  [{:keys [room-id peer-id sig-sealer] :as rc} inner to]
  (go
    (let [iv (crypto/take-iv! sig-sealer)
          key (<! (crypto/sig-key rc peer-id))
          ct (<! (crypto/aes-seal key iv (sig-aad room-id peer-id to) (utf8-b64 (json/stringify inner))))]
      (envelope iv ct))))

(defn seal-message
  "Seal INNER under this peer's message key (same ciphertext for every
  recipient). Channel of the envelope map."
  [{:keys [room-id peer-id msg-sealer] :as rc} inner]
  (go
    (let [iv (crypto/take-iv! msg-sealer)
          key (<! (crypto/msg-key rc peer-id))
          ct (<! (crypto/aes-seal key iv (msg-aad room-id peer-id) (utf8-b64 (json/stringify inner))))]
      (envelope iv ct))))

(defn seal-message-binary
  "Seal DATA-B64 under this peer's message key: 0x01‖iv‖ct‖tag as base64. Channel."
  [{:keys [room-id peer-id msg-sealer] :as rc} data-b64]
  (go
    (let [iv (crypto/take-iv! msg-sealer)
          key (<! (crypto/msg-key rc peer-id))
          ct (<! (crypto/aes-seal key iv (msg-aad room-id peer-id) data-b64))]
      (binary prefix-session iv ct))))

;; ---- opening from a remote peer: nil means drop silently ---------------------------------

(defn- open-json [key iv aad ct]
  (go
    (when-let [pt (<! (crypto/aes-open key iv aad ct))]
      (when-let [text (bytes/b64->utf8 pt)]
        (let [v (json/parse text)]
          (when (map? v) v))))))

(defn open-signal
  "The inner map of a relay envelope from FROM addressed to us, or nil. Channel."
  [{:keys [room-id peer-id] :as rc} env from]
  (go
    (when-let [{:keys [iv ct]} (parse-envelope env)]
      (when (crypto/valid-peer-id? from)
        (let [key (<! (crypto/sig-key rc from))]
          (<! (open-json key iv (sig-aad room-id from peer-id) ct)))))))

(defn open-message
  "The inner map of a data-channel text envelope from FROM, or nil. Channel."
  [{:keys [room-id] :as rc} env from]
  (go
    (when-let [{:keys [iv ct]} (parse-envelope env)]
      (when (crypto/valid-peer-id? from)
        (let [key (<! (crypto/msg-key rc from))]
          (<! (open-json key iv (msg-aad room-id from) ct)))))))

(defn open-message-binary
  "The plaintext (base64) of a binary data-channel envelope from FROM, or nil. Channel."
  [{:keys [room-id] :as rc} b64 from]
  (go
    (when-let [[iv ct] (split-binary b64 prefix-session)]
      (when (crypto/valid-peer-id? from)
        (let [key (<! (crypto/msg-key rc from))]
          (<! (crypto/aes-open key iv (msg-aad room-id from) ct)))))))

;; ---- at rest (mailbox): room-wide key, random IV -----------------------------------------

(defn at-rest-seal
  "0x02‖iv‖ct‖tag (base64) of DATA-B64 under the key for INFO with AAD
  suffix SUFFIX; IV random unless given (vectors only). Channel."
  ([rc info suffix data-b64] (at-rest-seal rc info suffix data-b64 (crypto/random-iv)))
  ([{:keys [room-id] :as rc} info suffix data-b64 iv-b64]
   (go
     (let [key (<! (crypto/key-for rc info))
           ct (<! (crypto/aes-seal key iv-b64 (at-rest-aad room-id suffix) data-b64))]
       (binary prefix-at-rest iv-b64 ct)))))

(defn at-rest-open
  "The plaintext (base64) of an at-rest envelope, or nil. Channel."
  [{:keys [room-id] :as rc} info suffix b64]
  (go
    (when-let [[iv ct] (split-binary b64 prefix-at-rest)]
      (let [key (<! (crypto/key-for rc info))]
        (<! (crypto/aes-open key iv (at-rest-aad room-id suffix) ct))))))

(defn mailbox-seal
  "Seal DATA-B64 for the mailbox (K_mbx, AAD v2|room|mbx). Channel."
  ([rc data-b64] (at-rest-seal rc crypto/mailbox-key-info "mbx" data-b64))
  ([rc data-b64 iv-b64] (at-rest-seal rc crypto/mailbox-key-info "mbx" data-b64 iv-b64)))

(defn mailbox-open
  "Open a mailbox payload; nil on any failure. Channel."
  [rc b64]
  (at-rest-open rc crypto/mailbox-key-info "mbx" b64))

;; ---- inner messages -------------------------------------------------------------------------

(defn identity-block
  "The identity block of a hello: {:pub :sig :name?} in wire order. Only
  these three keys are ever emitted; :name is omitted when nil."
  [{:keys [pub sig name]}]
  (cond-> {:pub pub :sig sig}
    (some? name) (assoc :name name)))

(defn hello
  "A hello inner message; ID is an identity block or nil."
  ([] {:kind "hello"})
  ([id] (if id {:kind "hello" :id (identity-block id)} {:kind "hello"})))

(defn offer [sdp] {:kind "offer" :sdp sdp})
(defn answer [sdp] {:kind "answer" :sdp sdp})
(defn ice
  "An ICE candidate; CANDIDATE nil means end of candidates."
  [candidate sdp-mid sdp-mline-index]
  {:kind "ice" :candidate candidate :sdpMid sdp-mid :sdpMLineIndex sdp-mline-index})
(def restart {:kind "restart"})
(defn msg [payload] {:kind "msg" :payload payload})
(defn bin [data-b64] {:kind "bin" :d data-b64})
(defn chunk-message [id i n enc data-b64] {:kind "chunk" :id id :i i :n n :enc enc :d data-b64})
(def ping {:kind "ping"})

(defn- parse-identity-block
  "{:pub :sig :name} of a hello's `id` (name nil when absent or not a
  string), nil when pub/sig are missing. Any other field a peer announces
  (e.g. the :hue/:glyph display opinions of other SDKs) is dropped silently."
  [v]
  (when (and (map? v) (string? (:pub v)) (string? (:sig v)))
    {:pub (:pub v) :sig (:sig v)
     :name (when (string? (:name v)) (:name v))}))

(defn parse-inner
  "Classify a decrypted plaintext map: {:kind <keyword> ...} with the fields
  of PROTOCOL.md, {:kind :unknown :name <string>} for kinds this SDK does not
  know, nil when the value is not an inner message at all (not a map, no
  kind, required field missing/malformed)."
  [v]
  (when (and (map? v) (string? (:kind v)))
    (case (:kind v)
      "hello" {:kind :hello :id (parse-identity-block (:id v))}
      "offer" (when (string? (:sdp v)) {:kind :offer :sdp (:sdp v)})
      "answer" (when (string? (:sdp v)) {:kind :answer :sdp (:sdp v)})
      "ice" {:kind :ice
             :candidate (when (string? (:candidate v)) (:candidate v))
             :sdp-mid (when (string? (:sdpMid v)) (:sdpMid v))
             :sdp-mline-index (when (integer? (:sdpMLineIndex v)) (:sdpMLineIndex v))}
      "restart" {:kind :restart}
      "msg" (when (contains? v :payload) {:kind :msg :payload (:payload v)})
      "bin" (when (bytes/valid-b64? (:d v)) {:kind :bin :data (:d v)})
      "chunk" (let [{:keys [id i n enc d]} v]
                (when (and (string? id) (integer? i) (integer? n) (contains? #{"json" "bin"} enc)
                           (bytes/valid-b64? d) (>= i 0) (pos? n) (< i n))
                  {:kind :chunk :id id :index i :count n :enc (keyword enc) :data d}))
      "ping" {:kind :ping}
      {:kind :unknown :name (:kind v)})))

;; ---- chunking ---------------------------------------------------------------------------------

(def data-channel-slice-size 8192)

(defn relay-slice-size
  "Raw bytes per relayed slice: base64 (×4/3) plus JSON and envelope
  overhead must stay under the frame cap: min(32768, (max_frame_bytes − 512) / 1.85)."
  [max-frame-bytes]
  (min 32768 (max 1024 (quot (* 100 (- max-frame-bytes 512)) 185))))

(defn new-chunk-id
  "8 random bytes as URL-safe base64."
  []
  (bytes/random-b64url 8))

(defn split
  "The `chunk` inner messages carrying DATA-B64 in slices of SLICE-SIZE raw
  bytes; ENC is \"json\" (a serialized msg inner message) or \"bin\"."
  [data-b64 enc slice-size]
  (let [id (new-chunk-id)
        total (bytes/byte-length data-b64)
        n (max 1 (quot (+ total slice-size -1) slice-size))]
    (mapv (fn [i]
            (let [start (* i slice-size)
                  end (min total (+ start slice-size))]
              (chunk-message id i n enc (bytes/slice-b64 data-b64 start end))))
          (range n))))

(def max-concurrent-assemblies 8)
(def max-assembly-age-ms 60000)
(def default-max-inflight-bytes (* 8 1024 1024))

(defn make-assembler
  "Reassembly state for one peer: {:partials {id {...}} :total 0 :max-bytes n}."
  ([] (make-assembler default-max-inflight-bytes))
  ([max-bytes] {:partials {} :total 0 :max-bytes max-bytes}))

(defn- drop-partial [asm id]
  (if-let [p (get-in asm [:partials id])]
    (-> asm (update :partials dissoc id) (update :total - (:bytes p)))
    asm))

(defn sweep
  "Drop assemblies older than 60 s."
  [asm now-ms]
  (reduce (fn [a [id p]] (if (> (- now-ms (:started p)) max-assembly-age-ms) (drop-partial a id) a))
          asm (:partials asm)))

(defn assemble
  "Add one chunk. Returns [asm' result] where result is nil (incomplete),
  :dropped (a bound was violated: 8 concurrent ids, the byte budget, an
  inconsistent chunk) or {:enc :json/:bin :data <b64>} when the message is
  complete."
  [asm {:keys [id index count enc data]} now-ms]
  (let [asm (sweep asm now-ms)
        existing (get-in asm [:partials id])
        p (or existing {:parts {} :count count :enc enc :bytes 0 :started now-ms})
        size (bytes/byte-length data)]
    (cond
      (or (not= (:count p) count) (not= (:enc p) enc) (>= index count) (contains? (:parts p) index))
      [(drop-partial asm id) :dropped]

      (and (nil? existing) (>= (clojure.core/count (:partials asm)) max-concurrent-assemblies))
      [asm :dropped]

      (> (+ (:total asm) size) (:max-bytes asm))
      [(drop-partial asm id) :dropped]

      :else
      (let [p (-> p (assoc-in [:parts index] data) (update :bytes + size))
            asm (-> asm (assoc-in [:partials id] p) (update :total + size))]
        (if (= (clojure.core/count (:parts p)) count)
          [(drop-partial asm id)
           {:enc enc :data (apply bytes/concat-b64 (map (:parts p) (range count)))}]
          [asm nil])))))