;; CljElixir PersistentVector — bit-partitioned trie
;; Ported from ClojureScript's PersistentVector
;;
;; BEAM adaptations:
;; - "Arrays" are BEAM tuples (immutable, O(1) indexed access via elem)
;; - aset = put-elem (returns new tuple)
;; - Tail grows via :erlang/append-element
;;
;; NOTE: Uses Map/get for struct field access (not keyword-as-function)
;; because ILookup is not yet implemented for these struct types.

;; ---------------------------------------------------------------------------
;; VectorNode — trie node containing up to 32 children
;; ---------------------------------------------------------------------------

(defrecord CljElixir.VectorNode [edit arr])

;; ---------------------------------------------------------------------------
;; PersistentVector — bit-partitioned trie with O(log32 n) indexed access
;; ---------------------------------------------------------------------------

(defrecord CljElixir.PersistentVector [meta cnt shift root tail]

  ;; --- Constants as zero-arity functions ---

  (defn empty-node []
    (CljElixir.VectorNode/new nil (erlang/make-tuple 32 nil)))

  (defn empty-vec []
    (CljElixir.PersistentVector/new nil 0 5 (empty-node) (erlang/make-tuple 0 nil)))

  ;; --- Internal helpers ---

  (defn- tail-off [cnt]
    (if (< cnt 32)
      0
      (erlang/bsl (erlang/bsr (dec cnt) 5) 5)))

  (defn- array-for [pv i]
    (let [cnt (Map/get pv :cnt)]
      (if (and (>= i 0) (< i cnt))
        (if (>= i (tail-off cnt))
          (Map/get pv :tail)
          (loop [node (Map/get pv :root)
                 level (Map/get pv :shift)]
            (if (> level 0)
              (let [child-idx (erlang/band (erlang/bsr i level) 31)]
                (recur (elem (Map/get node :arr) child-idx)
                       (- level 5)))
              (Map/get node :arr))))
        (throw (str "Index " i " out of bounds for vector of size " cnt)))))

  ;; --- Nth ---

  (defn pv-nth
    ([pv i]
      (let [node (array-for pv i)]
        (elem node (erlang/band i 31))))
    ([pv i not-found]
      (if (and (>= i 0) (< i (Map/get pv :cnt)))
        (pv-nth pv i)
        not-found)))

  ;; --- Path operations ---

  (defn- new-path [level node]
    (if (= level 0)
      node
      (let [new-arr (put-elem (erlang/make-tuple 32 nil) 0 (new-path (- level 5) node))]
        (CljElixir.VectorNode/new nil new-arr))))

  (defn- push-tail [cnt level parent tail-node]
    (let [subidx (erlang/band (erlang/bsr (dec cnt) level) 31)
          parent-arr (Map/get parent :arr)
          node-to-insert
            (if (= level 5)
              tail-node
              (let [child (elem parent-arr subidx)]
                (if (not (nil? child))
                  (push-tail cnt (- level 5) child tail-node)
                  (new-path (- level 5) tail-node))))
          new-arr (put-elem parent-arr subidx node-to-insert)]
      (CljElixir.VectorNode/new nil new-arr)))

  ;; --- Conj (append) ---

  (defn pv-conj [pv val]
    (let [cnt (Map/get pv :cnt)
          tail (Map/get pv :tail)
          tail-len (tuple-size tail)
          meta (Map/get pv :meta)
          shift (Map/get pv :shift)
          root (Map/get pv :root)]
      (if (< tail-len 32)
        ;; Room in tail
        (CljElixir.PersistentVector/new
          meta (inc cnt) shift root
          (erlang/append-element tail val))
        ;; Tail full — push into trie
        (let [tail-node (CljElixir.VectorNode/new nil tail)
              overflow? (> (erlang/bsr cnt 5) (erlang/bsl 1 shift))]
          (if overflow?
            ;; New root level
            (let [new-arr (put-elem
                            (put-elem (erlang/make-tuple 32 nil) 0 root)
                            1 (new-path shift tail-node))]
              (CljElixir.PersistentVector/new
                meta (inc cnt) (+ shift 5)
                (CljElixir.VectorNode/new nil new-arr)
                (erlang/make-tuple 1 val)))
            ;; Room at current depth
            (CljElixir.PersistentVector/new
              meta (inc cnt) shift
              (push-tail cnt shift root tail-node)
              (erlang/make-tuple 1 val)))))))

  ;; --- Assoc (update at index) ---

  (defn- do-assoc [level node i val]
    (let [node-arr (Map/get node :arr)
          node-edit (Map/get node :edit)]
      (if (= level 0)
        (CljElixir.VectorNode/new node-edit
          (put-elem node-arr (erlang/band i 31) val))
        (let [subidx (erlang/band (erlang/bsr i level) 31)
              new-child (do-assoc (- level 5) (elem node-arr subidx) i val)]
          (CljElixir.VectorNode/new node-edit
            (put-elem node-arr subidx new-child))))))

  (defn pv-assoc [pv i val]
    (let [cnt (Map/get pv :cnt)]
      (cond
        (and (>= i 0) (< i cnt))
          (if (>= i (tail-off cnt))
            (CljElixir.PersistentVector/new
              (Map/get pv :meta) cnt (Map/get pv :shift) (Map/get pv :root)
              (put-elem (Map/get pv :tail) (erlang/band i 31) val))
            (CljElixir.PersistentVector/new
              (Map/get pv :meta) cnt (Map/get pv :shift)
              (do-assoc (Map/get pv :shift) (Map/get pv :root) i val)
              (Map/get pv :tail)))
        (= i cnt)
          (pv-conj pv val)
        true
          (throw (str "Index " i " out of bounds for assoc on vector of size " cnt)))))

  ;; --- Pop (remove last) ---

  (defn- pop-tail [cnt level node]
    (let [subidx (erlang/band (erlang/bsr (dec cnt) level) 31)
          node-arr (Map/get node :arr)
          node-edit (Map/get node :edit)]
      (cond
        (> level 5)
          (let [new-child (pop-tail cnt (- level 5) (elem node-arr subidx))]
            (if (and (nil? new-child) (= subidx 0))
              nil
              (CljElixir.VectorNode/new node-edit
                (put-elem node-arr subidx new-child))))
        (= subidx 0)
          nil
        true
          (CljElixir.VectorNode/new node-edit
            (put-elem node-arr subidx nil)))))

  (defn pv-pop [pv]
    (let [cnt (Map/get pv :cnt)]
      (cond
        (= cnt 0)
          (throw "Can't pop empty vector")
        (= cnt 1)
          (empty-vec)
        true
          (let [tail (Map/get pv :tail)
                tail-len (tuple-size tail)]
            (if (> tail-len 1)
              ;; Shrink tail
              (let [new-tail (List/to-tuple (lists/droplast (Tuple/to-list tail)))]
                (CljElixir.PersistentVector/new
                  (Map/get pv :meta) (dec cnt) (Map/get pv :shift) (Map/get pv :root) new-tail))
              ;; Pull last leaf from trie
              (let [new-tail (array-for pv (- cnt 2))
                    shift (Map/get pv :shift)
                    new-root (pop-tail cnt shift (Map/get pv :root))
                    new-root (if (nil? new-root) (empty-node) new-root)
                    squish? (and (> shift 5) (nil? (elem (Map/get new-root :arr) 1)))
                    new-root (if squish? (elem (Map/get new-root :arr) 0) new-root)
                    new-shift (if squish? (- shift 5) shift)]
                (CljElixir.PersistentVector/new
                  (Map/get pv :meta) (dec cnt) new-shift new-root new-tail)))))))

  ;; --- Construction ---

  (defn from-list [xs]
    (Enum/reduce xs (empty-vec) (fn [x acc] (CljElixir.PersistentVector/pv-conj acc x))))

  (defn to-list [pv]
    (let [cnt (Map/get pv :cnt)]
      (if (= cnt 0)
        (list)
        (loop [i 0
               acc (list)]
          (if (< i cnt)
            (recur (inc i) (++ acc (list (pv-nth pv i))))
            acc)))))

  ;; --- Utility ---

  (defn pv-count [pv]
    (Map/get pv :cnt))

  (defn pv-with-meta [pv new-meta]
    (CljElixir.PersistentVector/new
      new-meta (Map/get pv :cnt) (Map/get pv :shift) (Map/get pv :root) (Map/get pv :tail))))

;; ---------------------------------------------------------------------------
;; SubVector — efficient view into an existing PersistentVector
;; ---------------------------------------------------------------------------

(defrecord CljElixir.SubVector [meta v start end]

  (defn sv-new
    ([v start] (CljElixir.SubVector/new nil v start (Map/get v :cnt)))
    ([v start end-idx]
      ;; If v is already a SubVector, flatten
      (if (Kernel/is-struct v CljElixir.SubVector)
        (CljElixir.SubVector/new nil (Map/get v :v) (+ (Map/get v :start) start) (+ (Map/get v :start) end-idx))
        (CljElixir.SubVector/new nil v start end-idx))))

  (defn sv-count [sv] (- (Map/get sv :end) (Map/get sv :start)))

  (defn sv-nth
    ([sv i]
      (let [actual-i (+ (Map/get sv :start) i)]
        (if (and (>= i 0) (< actual-i (Map/get sv :end)))
          (CljElixir.PersistentVector/pv-nth (Map/get sv :v) actual-i)
          (throw (str "Index " i " out of bounds for subvec of size " (sv-count sv))))))
    ([sv i not-found]
      (let [actual-i (+ (Map/get sv :start) i)]
        (if (and (>= i 0) (< actual-i (Map/get sv :end)))
          (CljElixir.PersistentVector/pv-nth (Map/get sv :v) actual-i)
          not-found))))

  (defn sv-to-list [sv]
    (let [start (Map/get sv :start)
          end-idx (Map/get sv :end)
          v (Map/get sv :v)]
      (loop [i start acc (list)]
        (if (< i end-idx)
          (recur (inc i) (++ acc (list (CljElixir.PersistentVector/pv-nth v i))))
          acc)))))