(ns clj-branca.core
  (:require [byte-streams :as bs]
            [clj-base62.core :as base62]
            [clj-branca.crypto :as crypto])
  (:import (java.nio BufferUnderflowException ByteBuffer)))

#_(set! *warn-on-reflection* true)

(def +version+
  "The supported Branca token version."
  (.byteValue (Long. 0xBA)))

(defn- key->bytes [key]
  (let [key (bs/to-byte-array key)]
    (when-not (= crypto/+key-length+ (count key))
      (throw (ex-info (format "Key must be %d bytes long; got %d bytes"
                              crypto/+key-length+
                              (count key))
                      {:type ::invalid-key})))
    key))

(defn- to-ms [timestamp]
  (if (inst? timestamp)
    (quot (inst-ms timestamp) 1000)
    timestamp))

(defn- get-now [options]
  (or (some-> (:now options) (to-ms))
      (quot (System/currentTimeMillis) 1000)))

(defn- ^ByteBuffer get-header
  [version timestamp ^bytes nonce]
  (doto (ByteBuffer/allocate 29)
    (.put (byte version))
    (.putInt timestamp)
    (.put nonce)
    (.rewind)))

(defn encode
  "Encode `payload` as a Branca token using the encryption key `key`.

  `key` should be a 32-byte byte array. `payload` should be a byte array, a
  string, or something that can be converted into a byte array using the
  [byte-streams] library.

  [byte-streams]: https://github.com/aleph-io/byte-streams

  Returns the token as a string, or throws an ExceptionInfo with `:type` set to
  `:clj-branca.core/encode-failure` in exception data.

  The following options are available:

  | key | description |
  | --- | ------------|
  | now | The token timestamp as Inst or seconds since the UNIX epoch. Default: system time. |
  "
  ([key payload] (encode key payload nil))
  ([key payload options]
   (let [now (get-now options)
         nonce (crypto/random-bytes 24)
         header (get-header +version+ now nonce)
         ciphertext (crypto/encrypt (bs/to-byte-array payload)
                                    (.array header)
                                    nonce
                                    (key->bytes key))
         token (doto (ByteBuffer/allocate (+ 29 (count ciphertext)))
                 (.put header)
                 (.put ^bytes ciphertext))]
     (if ciphertext
       (base62/encode (.array token))
       (throw (ex-info "Encryption failed" {:type ::encode-failure}))))))

(defn decode*
  "Decode `token` as a Branca token using encryption key `key`.

  Returns a map representing the token. If you just want the payload, use
  [[decode]]. If there's a problem, throws an ExceptionInfo with `:type` set to
  `:clj-branca.core/invalid-token` in exception data.

  Takes the same options as [[decode]]."
  ([key token] (decode* key token nil))
  ([key token options]
   (try
     (let [key (key->bytes key)
           buffer (ByteBuffer/wrap (base62/decode token))
           version (.get buffer)
           timestamp (Integer/toUnsignedLong (.getInt buffer))
           nonce (doto (byte-array 24) (->> (.get buffer)))
           ciphertext (doto (byte-array (.remaining buffer))
                        (->> (.get buffer)))
           header (get-header version timestamp nonce)
           payload (crypto/decrypt ciphertext
                                   (.array header)
                                   nonce
                                   key)]
       (when-not (= +version+ version)
         (throw (ex-info (format "Token version must be %d, got %d" +version+ version)
                         {:type ::invalid-token})))
       (when-not payload
         (throw (ex-info "Decryption failed" {:type ::invalid-token})))
       (when-let [ttl (:ttl options)]
         (when (< (+ timestamp ttl) (get-now options))
           (throw (ex-info "Expired token"
                           {:type ::invalid-token
                            :timestamp timestamp
                            :ttl ttl}))))
       {:version version
        :timestamp timestamp
        :nonce nonce
        :payload payload})
     (catch BufferUnderflowException _
       (throw (ex-info "Malformed token" {:type ::invalid-token}))))))

(defn decode
  "Decode `token` as a Branca token using encryption key `key`.

  Returns the token payload as a byte array, or throws an ExceptionInfo with
  `:type` set to `:clj-branca.core/invalid-token` in exception data.

  The following options are available:

  | key | description |
  | --- | ------------|
  | now | The current time as Inst or seconds since the UNIX epoch. Default: system time. |
  | ttl | Token time-to-live in seconds. If set, throws if the token has expired. Default: nil. |
  "
  ([key token] (:payload (decode* key token)))
  ([key token options] (:payload (decode* key token options))))