(import jnj-primitives)

(def j-engine
  ``
  The default J instance, used by `jnj/j`.
  ``
  (jnj-primitives/init))

(defn- pack
  `Given an ordered list of items and a shape vector, pack the items into
  nested arrays.`
  [shape acc]
  (if (= 1 (length shape))
    # The partitioning should leave (shape 1) items left, so need to partition
    # on the last call
    (if (bytes? acc)
      (string acc)
      (tuple ;acc))
    (let [row (first shape)
          rest (array/slice shape 1)]
      (pack rest (partition row acc)))))

(defn- to-matrix
  [j-array]

  (let [shape (jnj-primitives/shape j-array)
        items (jnj-primitives/values j-array)]
    (pack (reverse shape) items)))

(defn- clear-locale
  "Erase all the variables in the given namespace"
  [je locale]
  (jnj-primitives/do je (string/format "18!:55 <'%s'" locale)))

(defn- jsetm
  [je sym arg]
  (let [response (jnj-primitives/setm je sym arg)]
    (if (not= 0 response)
      (errorf "Got response %q when trying to set %q to %q"
              response sym arg))))

(defn- jdo
  [je cmd-string]
  (let [response (jnj-primitives/do je cmd-string)]
    (if (not= 0 response)
      (errorf "Got response %q when trying to execute %q"
              response cmd-string))))

(defn- make-locale
  []
  (string "l" (string/trim (gensym) "_")))

(defn- intern
  [je arg sym]
  (match arg
    @[(@ :raw) _] :ok
    (a (keyword? a)) :ok
    (jsetm je sym arg)))

(defn- compose
  [res-sym & strings]
  (string res-sym "=: " (string/join strings " ")))

(defn- result-sym
  [locale &opt k]
  (default k "jnjr")
  (string k "_" locale "_"))

(defn- intern-and-do-in
  [locale je verb args &opt kw]
  (let [locale-suffix (string "_" locale "_")
        syms (map |(match $0
                     @[:raw raw-j] raw-j
                     (k (keyword? k)) (string $0 locale-suffix)
                     (string "jnj" (string/trim (gensym) "_") locale-suffix))
                  args)
        arg-syms (map |[$0 $1] args syms)
        res-sym (result-sym locale kw)
        verb (match verb
               @[:raw raw-j] raw-j
               (k (keyword? k)) (string k locale-suffix)
               (string verb))
        cmd-string (match syms
                     @[x y] (compose res-sym x verb y)
                     @[x] (compose res-sym verb x)
                     @[] (compose res-sym verb))]

    (each [arg sym] arg-syms (intern je arg sym))

    (jdo je cmd-string)))

(defn- getm-in
  [locale je res-sym]
  (let [locale-suffix (string "_" locale "_")
        jvalue (jnj-primitives/getm je res-sym)]
    (clear-locale je locale)
    jvalue))

(defn eval*
  ``
  Evaluate `verb` with arguments `args` in the context of `je`. Returns a
  j-array abstract type which can be used for further evaluations, or converted
  into a tuple matrix.
  ``
  [je verb & args]

  (let [locale (make-locale)]
    (intern-and-do-in locale je verb args)
    (getm-in locale je (result-sym locale))))

(defn from-j-value
  ``
  Turn a j-value abstract type result into a Janet term of the appropriate shape:
  - a string
  - a scalar value
  - a matrix of nested tuples
  ``
  [res]
  (if (zero? (jnj-primitives/rank res))
    (res 0)
    (to-matrix res)))

(defn- let-j-impl
  [je bindings body]
  (let [locale (make-locale)
        locale-suffix (string "_" locale "_")
        kws-and-exprs (partition 2 bindings)]

    (each [kw expr] kws-and-exprs
      (let [verb (first expr)
            args (tuple/slice expr 1)]
        (intern-and-do-in locale je verb args kw)))

    (let [res-sym (result-sym locale)
          verb (first body)
          args (tuple/slice body 1)]
      (intern-and-do-in locale je verb args)
      (-> (getm-in locale je res-sym)
          (from-j-value)))))

(defmacro let-j*
  ````
  Use a series of intermediate bindings to compute a complex J value and
  convert it back to Janet.

  `bindings` should be, like with `let`, alternating binding keywords and J expressions.

  A binding keyword is any keyword, which can then be referred to in subsequent
  expressions.

  A J expression is a tuple of the form that can be passed to `eval*` or `j*` -
  an arbitrary symbol-or-string sentence followed by 0, 1 or 2 args. In
  addition to the arg types understood by `eval*` and `j*`, an arg can also be
  a keyword, in which case it will refer to the result bound to that keyword
  earlier in the `bindings` form.

  Finally, `body` is the J expression to be evaluated - again, a sentence
  followed by 0, 1 or 2 arguments, where the arguments can include any keywords
  specified in the bindings.

  Here's an example that creates a 4x3 matrix, gets its shape, and then sums
  the shape:

  ```
  > (let-j* je
  >  [x ("$" [3 4] [0 1])
  >   y ("$" x)]
  >   ('+/ y))
  7
  ```

  Because the intermediate expressions in let-j only consist in the J engine,
  it can also be used to handle datatypes not yet understood by JNJ, like
  boxes:

  ```
  > (let-j* je
  >  [box ("<" "foo")
  >   arglist (";" box "bar")]
  >   ("#" arglist))
  2
  ```

  This same fact makes let-j operations somewhat more efficient than multiple
  `eval` calls, as the intermediate values don't need to be copied back into
  the J runtime.
  ````
  [je bindings body]
  (let [bindings-vals @[]
        bound-symbols @{}]

    (defn walker
      [node]
      (match node
        @[(@ '@) raw] @[:raw (string raw)]
        (n (bound-symbols n)) (keyword node)
        node))

    (each [binding expr] (partition 2 bindings)
      (put bound-symbols binding true)
      (array/push bindings-vals [(keyword binding)
                                 (array/slice (walk walker expr))]))

    ~(,let-j-impl ,je ,(mapcat identity bindings-vals) ,(array/slice (walk walker body)))))

(defn j*
  ``
  Evaluate `verb` with `args` in the context of `je`, returning a native Janet
  datatype:
  - a matrix of nested tuples for anything with rank > 0
  - a Janet atom (string, number) otherwise
  ``
  [je verb & args]

  (let [res (eval* je verb ;args)]
    (from-j-value res)))

(defn to-j-array
  ```
  Turn an arbitrarily nested array/tuple of numbers into a J Array.

  NB: This function performs some validation on its input to assert that it can
  be transformed into a J Array. So if the data is not already in Janet form,
  it might be more efficient to generate it using `jnj/eval`.
  ```
  [matrix]
  (defn build-shape
    [row acc]
    (let [hd (first row)
          hd-type (type hd)
          hd-length (if (indexed? hd) (length hd))]

      (each elem row
        (unless (or (indexed? elem) (= :number (type elem)))
          (errorf "Could not convert to numeric j-array: %q" matrix))
        (unless (or (not hd-length) (= hd-length (length elem)))
          (errorf "Matrix has inconsistent shape: %q" matrix)))

      (array/push acc (length row))
      (if (indexed? hd) (build-shape hd acc) acc)))

  (let [shape (build-shape matrix @[])
        rank (length shape)
        data (flatten matrix)
        count (length data)]
    (jnj-primitives/to-j-value rank count shape data)))

(defn eval
  ``
  Evaluate `verb` and `args` with the default J engine instance (see `eval*`
  for details).
  ``
  [verb & args]
  (eval* j-engine verb ;args))

(defn j
  ``
  Evaluate `verb` with `args` in the default J engine instance (see `j*` for
  details).
  ``
  [verb & args]
  (j* j-engine verb ;args))

(defmacro let-j
  ``
  Evaluate let-j `bindings` and `expr` in the default J engine instance (see
  `let-j*` for details).
  ``
  [bindings expr]
  (apply let-j* [j-engine bindings expr]))