(ns jepsen.antietcd.watch
  "Tests for watches. We have processes watch a changing key for a while, and
  return a list of the versions it went through. At the end of the test, we
  confirm that every process observed an arbitrary sub-sequence of system states
  (not all of them because antietcd doesn't keep all history) in the same order."
  (:require
    [clojure.data :refer [diff]]
    [clojure.data.json :as json]
    [clojure.tools.logging :refer [info warn]]
    [hato.websocket :as ws]
    [jepsen
      [checker :as checker]
      [client :as client]
      [generator :as gen]
      [util :as util :refer [meh map-vals]]
    ]
    [jepsen.checker.timeline :as timeline]
    [jepsen.antietcd.client :as anticli]
    [knossos.op :as op]
    [slingshot.slingshot :refer [throw+ try+]]
  )
  (:import
    (java.util.concurrent BrokenBarrierException)
    (java.util.concurrent.locks LockSupport)
  )
)

(defn converger
  "Generates a convergence context for n threads, where values are converged
  when (converged? values) returns true."
  [n converged?]
  (atom {
    ; The convergence function
    :converged? converged?
    ; What threads are involved?
    :threads []
    ; And what values did they most recently come to?
    :values (vec (repeat n ::init))
  })
)

(defn await-converger-change
  "Wait for a converger's state to change. Might return for no reason."
  [c]
  (LockSupport/park c)
)

(defn signal-converger-change!
  "Let a converger know that its state has changed."
  [c]
  (doseq [t (:threads @c)]
    (LockSupport/unpark t)
  )
)

(defn stable?
  "A converger state is stable when no value is initial or evolving."
  [c]
  (not-any? #{::init ::evolving} (:values c))
)

(defn divergent?
  "A converger state is divergent when any values are initial, or some
  non-evolving values don't pass the convergence test."
  [c]
  (or (some #{::init} (:values c))
    (let [vs (remove #{::evolving} (:values c))]
      (and (seq vs)
        (not ((:converged? c) vs))
      )
    )
  )
)

(defn converged?
  "A converger state is converged when it is stable and non-divergent."
  [c]
  (and (stable? c) (not (divergent? c)))
)

(defn evolve!
  "Takes a converger, an evolve function, an initial value, and a thread index
  i. Takes the previous value for this thread (or the initial value), clears it
  from the converger, and calls (evolve value) to generate value', finally
  updating the converger with value'."
  [converger evolve init i]
  (let [value (atom nil)]
    ; Get the current value and clear it from the converger
    (swap! converger
      (fn [c]
        (let
          [
            v (nth (:values c) i)
            v (if (= ::init v) init v)
          ]
          (reset! value v)
        )
        (assoc-in c [:values i] ::evolving)
      )
    )

    ; Evolve
    (try
      (let [value' (evolve @value)]
        (swap! converger #(assoc-in % [:values i] value'))
      )
      (catch Throwable t
        ; If we throw here, all bets are off. We can't safely return, since we
        ; haven't converged. We set :crashed? to true in the converger, and let
        ; all threads wake up.
        (swap! converger assoc :crashed? true)
      )
      (finally
        ; Let other threads know we've done some work and they can wake up
        (signal-converger-change! converger)
      )
    )
    nil
  )
)

(defn converge!
  "Takes a timeout in ms, a converger, an initial value, and a function which
  evolves values over time.

  When `converge` is called, evaluates `evolve` repeatedly, starting with the
  initial value; each successive invocation receives the result of the previous
  invocation. Returns a value only when (converged? [v1 v2 ...]) returns
  truthy.

  Always invokes evolve at least once; the initial state doesn't count as
  converged.

  If evolve throws, horrible things will happen. Probably deadlock or livelock.
  I should fix this later."
  [timeout converger init evolve]
  ; Acquire our unique thread index
  (let
    [
      deadline (+ (util/linear-time-nanos) (util/ms->nanos timeout))
      i (->
        (swap! converger update :threads conj (Thread/currentThread))
        :threads
        count
        dec
      )
    ]
    (loop []
      (let [c @converger]
        (cond
          ; If we've crashed, throw a BrokenBarrierException. We're not
          ; really using a cyclic barrier, but this is *basically* a
          ; concurrency barrier, and it *is* broken.
          (:crashed? c)
          (throw (BrokenBarrierException. "Convergence failed"))

          ; If the deadline has passed, throw a timeout with our value.
          (< deadline (util/linear-time-nanos))
          (throw+ {
            :type  ::converge-timeout
            :value (-> converger deref :values (get i))
          })

          ; If we're converged, return our value
          (converged? c)
          (-> converger deref :values (get i))

          ; If we're divergent, evolve
          (divergent? c)
          (do (evolve! converger evolve init i) (recur))

          ; Otherwise, wait for something to change
          true
          (do (await-converger-change converger) (recur))
        )
      )
    )
  )
)

(defn watch
  "Starts a watch `req` from revision on, returning a deref-able which, when
  dereferenced, stops watching and returns results."
  [url process req revision]
  ; State is a map of {:revision x, :log y}.
  (let
    [
      state (atom {:revision revision :log []})
      ws @(ws/websocket (str "ws" (subs url 4) "/v3/watch") {
        :connect-timeout 5000
        :on-open (fn [ws]
          (ws/send! ws (json/write-str {
            :create_request (merge req {
              ; Revision is inclusive, so we want to start just after we left off
              ; Definitely don't pass 0, or you'll get wherever it is currently
              :start_revision (inc revision)
              :progress_notify true
            })
          }))
        )
        :on-message (fn [ws msg last?]
          (let
            [
              result (:result (json/read-str (str msg) :key-fn keyword))
              events (->> (:events result) (map (comp :value :kv)))
              rev' (->> (:events result) (map (comp :mod_revision :kv)) (reduce max))
              error' (cond
                (not result)
                (do
                  (info process (str "invalid json in websocket message: " msg))
                  {:type :invalid_json}
                )
                (:compact_revision result)
                {:type :watch_canceled :msg (str "watch canceled, reason: revisions before " (:compact_revision result) " were compacted")}
                (:canceled result)
                {:type :watch_canceled :msg (str "watch canceled, reason: " (:cancel_reason result))}
                (>= revision rev')
                {:type :nonmonotonic :msg (str "got event with revision " rev' " but we last saw " revision)}
                true nil
              )
            ]
            (info process (str msg))
            (swap! state (fn advance [{:keys [revision log error] :as s}]
              {
                :revision rev'
                :log      (conj log result)
                :error    (or error error')
              }
            ))
          )
        )
        :on-close (fn [ws status reason]
          (info process (str "websocket closed, reason: " reason ", status: " status))
        )
        :on-error (fn [ws err]
          (info process (str "websocket aborted, reason: " err))
        )
      })
    ]
    ; When deref'ed, we kill the watch and return results.
    (delay
      (meh (ws/abort! ws))
      (info process (if (:error @state)
        (str "error watching from" revision "to" (:revision @state) ": " (:error @state))
        (str "completed up to revision " (:revision @state))
      ))
      @state
    )
  )
)

(defn watch-for
  "Starts a watch request `req` from `revision` for `ms` milliseconds."
  [url process req revision ms]
  (let [w (watch url process req revision)]
    (Thread/sleep (long ms))
    @w
  )
)

(defrecord WatchClient [url max-revision revision converger]
  client/Client
  (open! [this test node]
    (assoc this
      :url (anticli/node-url node)
      :revision (atom 0)
    )
  )

  (setup! [this test])

  (invoke! [_ test op]
    (case (:f op)
      :write (let
        [res (anticli/throw #(anticli/antietcd_put url (:key op) (:value op)))]
        (->> res :body :header :revision (swap! max-revision max))
        (assoc op :type :ok :revision (->> res :body :header :revision))
      )

      :watch
      (let
        [res (watch-for url (:process op) {:key "/" :range_end "0"} @revision (rand-int 5000))]
        ; Advance our revision counter for the next watch
        (reset! revision (:revision res))
        ; And the global revision counter
        (swap! max-revision max (:revision res))
        (assoc op :type :ok, :value res)
      )

      :final-watch
      (try+
        (let
          [
            v (converge!
              60000 converger {:revision @revision, :log []}
              (fn [v] (let
                [
                  rev (:revision v)
                  log (:log v)
                ]
                (let
                  [
                    _ (info "at rev" rev "catching up to" @max-revision)
                    w (watch-for url (:process op) {:key "/" :range_end "0"} rev (rand-int 5000))
                  ]
                  (if (instance? Throwable w)
                    (do
                      (warn w "caught during final-watch; retrying")
                      (Thread/sleep 1000)
                      v
                    )
                    (do
                      ; Advance our revisions
                      (reset! revision (:revision w))
                      (swap! max-revision max (:revision w))
                      (assoc w :log (into (:log v) (:log w)))
                    )
                  )
                )
              ))
            )
          ]
          (assoc op :type :ok, :value v)
        )
        (catch [:type ::converge-timeout] e
          (assoc op :type :ok, :value (:value e), :error [:converge-timeout])
        )
      )
    )
  )

  (teardown! [this test])

  (close! [this test])
)

(defn log-to-states
  "Transforms a log into a series of system states observed at different points of history.
  The idea is that each client may see a subsequence of system states, but he must see each
  state consistently even if his watches break and restart during operation. Returns map
  with revisions => data."
  [log]
  (:states (reduce
    (fn [acc val] (let
      [
        revision' (:revision (:header val))
        st' (apply dissoc
          (merge (get (:states acc) (:revision acc)) (apply hash-map
            (apply concat
              (map
                (fn [x] [(:key (:kv x)) (:value (:kv x))])
                (filter #(= (:type %) "PUT") (:events val))
              )
            )
          ))
          (map
            #(:key (:kv %))
            (filter #(= (:type %) "DELETE") (:events val))
          )
        )
      ]
      {:states (assoc (:states acc) revision' st') :revision revision'}
    ))
    {:states {0 {}} :revision 0}
    log
  ))
)

(defn canonical-log
  "Merges all events from raw logs to get the canonical log."
  [logs]
  (->>
    logs
    (apply concat)
    (map :events)
    (apply concat)
    (sort-by (comp :mod_revision :kv))
    (map (fn [ev] {:header {:revision (:mod_revision (:kv ev))} :events [ev]}))
    (log-to-states)
  )
)

(defn per-thread-watches
  "Takes a test and a history, and computes a map of threads to ok watch operations."
  [test history]
  (let [concurrency (:concurrency test)]
    (->>
      history
      (filter op/ok?)
      (filter (comp #{:watch :final-watch} :f))
      (group-by (fn [op] (mod (:process op) concurrency)))
    )
  )
)

(defn per-thread-logs
  "Takes a test and a history, and computes a map of threads to the sequence of values they observed."
  [test history]
  (map-vals
    (fn [ops] (mapcat (comp :log :value) ops))
    (per-thread-watches test history)
  )
)

(defn per-thread-revisions
  "Takes a test and a history, and computes a map of threads to the highest
  revision they observed."
  [test history]
  (map-vals
    (fn [ops] (reduce max 0 (map (comp :revision :value) ops)))
    (per-thread-watches test history)
  )
)

(defn nonmonotonic-errors
  "Takes a history and extracts any nonmonotonic error messages."
  [history]
  (->> history
    (keep :error)
    (filter #(= (:type %) :nonmonotonic))
  )
)

(defn checker
  "We reconstruct the order of values as seen by each individual process, then
  compare them for equality."
  []
  (reify checker/Checker
    (check [_ test history opts]
      (let
        [
          logs (per-thread-logs test history)
          revisions (per-thread-revisions test history)
          _         (info :logs logs)
          canonical (canonical-log (vals logs))
          sorted    (vec (sort-by #(get % 0) canonical))
          phantoms  (vec (filter
            (fn [x] (< 1 (count (second x))))
            (map
              (fn [x] [
                (first (get sorted (inc x)))
                (second (diff
                  (second (get sorted x))
                  (second (get sorted (inc x)))
                ))
              ])
              (range (dec (count sorted)))
            )
          ))
          diffs (apply hash-map (apply concat (->> logs
            (keep (fn [[thread log]] (let
              [bad (filter
                (fn [[rev state]]
                  (not (= state (get canonical rev)))
                )
                (log-to-states log)
              )]
              (if (> (count bad) 0) [thread bad] nil)
            )))
          )))
          nm-errors (nonmonotonic-errors history)
          valid? (cond
            (seq nm-errors) false
            (seq diffs)     false
            (seq phantoms)  false
            :else           true
          )
        ]
        (cond->
          {
            :valid?    valid?
            :revisions revisions
          }
          (not valid?) (assoc
            :logs       logs
            :canonical  sorted
            :phantoms   phantoms
            :diffs      diffs
            :nonmonotonic-errors nm-errors
          )
        )
      )
    )
  )
)

(defn workload
  "We allocate the first node-count threads as writers, and the remaining
  threads as watchers."
  [opts]
  (let
    [
      node-count  (count (:nodes opts))
      watch-count (- (:concurrency opts) node-count)
      converger   (converger watch-count (fn converged? [ms] (apply = (map :revision ms))))
    ]
    {
      :client    (WatchClient. nil (atom 0) nil converger)
      :checker   (checker)
      :generator (let
        [
          write (->>
            (range)
            (map (fn [x] {:type :invoke :f :write :key (str "/k" (rand-int 5)) :value x}))
          )
          watch (repeat {:type :invoke, :f :watch})
        ]
        (gen/reserve node-count write watch)
      )
      :final-generator (
        gen/reserve node-count nil
        (gen/each-thread {:type :invoke :f :final-watch})
      )
    }
  )
)
