(ns jepsen.antietcd.append
  "Tests append/read transactions over lists using Elle."
  (:require
    [clojure.tools.logging :refer [info warn]]
    [jepsen
      [client :as client]
      [checker :as checker]
      [generator :as gen]
      [store :as store]
      [util :as util]
    ]
    [jepsen.tests.cycle :as cycle]
    [jepsen.tests.cycle.append :as append]
    [jepsen.antietcd.client :as anticli]
  )
)

(defn preprocess
  "In this phase, we take a test map, a connection, and an operation, and
  extract its transaction, returning a transaction map with keys:
  :test :client :op :txn :read-only?"
  [test client op]
  {
    :test       test
    :client     client
    :op         op
    :txn        (:value op)
    :read-only? (every? (comp #{:r} first) (:value op))
  }
)

(defn encode-put
  "Takes a transaction map and a value to write, and transforms the value
  somehow. This is identity normally, but when we put something to antietcd in
  debug mode we also include extra information about the txn that generated
  that value."
  [t v]
  (if (:debug (:test t))
    {:time (str (util/local-time))
     :dir  (.getName (store/path (:test t)))
     :txn (:txn t),
     :value v}
    v
  )
)

(defn decode-get
  "Takes a transaction map and a value read from etc, and transforms the value
  back into what the Jepsen test expects--stripping off debugging information,
  for instance."
  [t v]
  (if (:debug (:test t))
    (:value v)
    v
  )
)

(defn written-keys
  "Takes a transaction and returns the set of keys being written."
  [t]
  (->>
    (:txn t)
    (filter (comp #{:append} first)) ; Take only appends
    (map (comp str second))          ; [f k v] -> k
    set
  )
)

(defn apply-read
  "In the read phase, we take a txn map and perform a transaction which reads
  all the keys we need to write, adding new keys to our txn map:

    :read-values    is a key-value map
    :read-revision  is the revision we read at
    :read-res       Contains the raw results of the read txn, used for debugging"
  [t all]
  (let
    [
      ops (mapv
        (fn [x] {:request_range {:key x}})
        (if all
          (set (map (comp str second) (:txn t)))
          (written-keys t)
        )
      )
      res (anticli/throw #(anticli/antietcd_request (str (:url (:client t)) "/v3/kv/txn") {:success ops}))
    ]
    (assoc t
      :read-values (->>
        (:responses (:body res))
        (map (fn [x] (:kvs (:response_range x))))
        (reduce concat)
        (map (fn [x] {(str (:key x)) (decode-get t (:value x))}))
        (reduce merge)
      )
      :read-res res
      :read-revision (:revision (:header (res :body)))
    )
  )
)

(defn guards
  "Takes a transaction with reads, and constructs a collection of guards
  verifying each read's revision is intact."
  [t]
  (mapv
    (fn [k] {:key k :target :MOD :result :LESS :mod_revision (+ (:read-revision t) 1)})
    (written-keys t)
  )
)

(defn write-txn
  "We take a transaction with reads, and iterate through all its :append
  operations, converting each to an antietcd txn op like {:request_put}. To do
  this, we use the value initially read, and conj on the appended element.
  Returns a sequence of etd gets and puts, rather than reads and appends."
  [t]
  (->>
    (:txn t)
    ; `state` maps keys to currently expected values; we play forward
    ; simulated effects of multiple appends in a single txn, turning them
    ; into multiple puts. This lets us catch intermediate read or other txn
    ; anomalies involving multiple writes.
    (reduce
      (fn [[state txn] [f k v :as mop]]
        (case f
          :r
            [state (conj txn {:request_range {:key k}})]
          :append
            (let
              [ v' (conj (get state (str k) []) v) ]
              [
                (assoc state (str k) v')
                (conj txn {:request_put {:key k :value (encode-put t v')}})
              ]
            )
        )
      )
      [
        (:read-values t)
        []
      ]
    )
    second
  )
)

(defn apply-rw!
  "Takes a transaction map and applies reads and writes, rewriting :txn to
  include the results of reads. Returns a completed operation, including extra
  debug information under a :debug key:

    :write-txn      The write txn we executed, as a vector of [guards true-branch]
    :write-res      The result of the write txn
  "
  [t]
  (let
    [
      t (apply-read t false)
      read-values (:read-values t)
      guards (guards t)
      write-txn (write-txn t)
      res (anticli/antietcd_txn (:url (:client t)) {:compare guards :success write-txn})

      ; Rewrite Jepsen txn with read results
      txn' (
        mapv (fn [[f k v :as mop] r]
          (case f
            :append mop
            :r [f k (decode-get t (-> r :response_range :kvs first :value))]
          )
        )
        (:txn t)
        (:responses (:body res))
      )

      ; Add debugging info
      op (if (:debug (:test t))
        (assoc (:op t)
          :debug
          {
            :read-values    read-values
            :read-res       (:read-res t)
            :read-revision  (:read-revision t)
            :write-txn      [guards write-txn]
            :write-txn-res  res
          }
        )
        (:op t)
      )
    ]
    (cond
      (= (:type res) :ok) (assoc op :type :ok, :value txn')
      (= (:error res) :mismatch) (assoc op :type :fail, :error "optimistic lock broken")
      true (throw (java.util.concurrent.TimeoutException. (str (:error res))))
    )
  )
)

(defn apply-ro!
  "Takes a transaction map for a read only transaction, and returns the
  completed operation, with extra information under a :debug key."
  [t]
  (let
    [
      t (apply-read t true)

      ; Rewrite Jepsen txn with read results
      txn' (
        mapv (fn [[f k v :as mop]]
          (case f
            :append mop
            :r [f k (decode-get t (get (:read-values t) (str k)))]
          )
        )
        (:txn t)
      )

      ; Add debugging info
      op (if (:debug (:test t))
        (assoc (:op t)
          :debug
          {
            :read-values    (:read-values t)
            :read-res       (:read-res t)
            :read-revision  (:read-revision t)
          }
        )
        (:op t)
      )
    ]
    (assoc op :type :ok, :value txn')
  )
)

(defn apply-singleton-ro!
  "Takes a transaction map for a read-only transaction with a single operation,
  and returns the completed operation."
  [{:keys [txn client op test] :as t}]
  (let
    [
      [f k _ :as mop] (first txn)
      res (anticli/throw #(anticli/antietcd_request (str (:url (:client t)) "/v3/kv/range") {:key k}))
      txn' [[f k (decode-get t (:value (first (:kvs (:body res)))))]]
    ]
    (assoc op :type :ok, :value txn')
  )
)

(defrecord TxnClient [url]
  client/Client
  (open! [this test node] (assoc this :url (anticli/node-url node)))

  (setup! [_ test])

  (invoke! [this test op]
    (when
      (:debug test)
      (gen/log (str "=================================================\nOp: " (pr-str op)))
    )

    (let [t (preprocess test this op)]
      (if (:read-only? t)
        (if (= 1 (count (:txn t)))
          (apply-singleton-ro! t)
          (apply-ro! t))
        (apply-rw! t)
      )
    )
  )

  (teardown! [_ test])

  (close! [_ test])
)

(defn workload
  "A generator, client, and checker for a list-append test using Elle."
  [opts]
  (-> (append/test
    {
      :key-count          3
      :max-txn-length     4
      :consistency-models [(if (get opts :stale-read false) :serializable :strict-serializable)]
    })
    (assoc :client (TxnClient. nil))
  )
)
