;; lib/erlang/lists-ext.sx — extra `lists` module BIFs.
;;
;; Loaded AFTER runtime.sx so the BIF registry + transpile helpers
;; (er-mk-cons, er-lt?, er-equal?, er-bool, er-truthy?, er-apply-fun,
;; er-bif-arg1, er-mk-error-marker, er-mk-atom, er-cons?, er-nil?) and
;; runtime's er-register-pure-bif! are all in scope. Registrations run
;; at load time and persist in the global er-bif-registry cell.
;;
;; Edit-tool note: the sx-tree write tools raise yojson "Expected
;; string, got null" in this worktree, so new lists BIFs land here (a
;; fresh file) rather than as in-place edits to the ~1900-line
;; transpile.sx. Same separate-file pattern the VM dispatcher already
;; uses (lib/erlang/vm/dispatcher.sx).

;; ── cons <-> SX-list bridges ──────────────────────────────────────
(define
  er-cons->sxlist
  (fn (lst)
    (cond
      (er-nil? lst) (list)
      (er-cons? lst) (cons (get lst :head) (er-cons->sxlist (get lst :tail)))
      :else (raise (er-mk-error-marker (er-mk-atom "badarg"))))))

(define
  er-sxlist->cons
  (fn (xs)
    (if (= (len xs) 0)
      (er-mk-nil)
      (er-mk-cons (first xs) (er-sxlist->cons (rest xs))))))

;; ── merge sort over SX lists (stable) ─────────────────────────────
(define
  er-ext-take
  (fn (xs n)
    (if (or (= n 0) (= (len xs) 0))
      (list)
      (cons (first xs) (er-ext-take (rest xs) (- n 1))))))

(define
  er-ext-drop
  (fn (xs n)
    (if (or (= n 0) (= (len xs) 0))
      xs
      (er-ext-drop (rest xs) (- n 1)))))

;; le? returns a truthy value (Erlang bool atom or SX bool) iff a
;; should sort at-or-before b. Taking from the left half first on a
;; true result keeps the sort stable.
(define
  er-ext-merge
  (fn (a b le?)
    (cond
      (= (len a) 0) b
      (= (len b) 0) a
      (er-truthy? (le? (first a) (first b)))
        (cons (first a) (er-ext-merge (rest a) b le?))
      :else (cons (first b) (er-ext-merge a (rest b) le?)))))

(define
  er-ext-msort
  (fn (xs le?)
    (if (<= (len xs) 1)
      xs
      (let ((mid (quotient (len xs) 2)))
        (er-ext-merge
          (er-ext-msort (er-ext-take xs mid) le?)
          (er-ext-msort (er-ext-drop xs mid) le?)
          le?)))))

;; Full Erlang term order. The shared er-lt? (transpile.sx) only
;; deep-compares numbers/atoms/strings and otherwise falls back to a
;; coarse type rank — so any two tuples (or two lists) compare as
;; order-equal there. er-ext-lt? adds the missing structural cases:
;; tuples by arity then elementwise, lists elementwise with a shorter
;; proper prefix sorting first. Cross-type cases delegate to er-lt?.
(define
  er-ext-lt-seq
  (fn (ea eb i)
    (cond
      (>= i (len ea)) false
      (er-ext-lt? (nth ea i) (nth eb i)) true
      (er-ext-lt? (nth eb i) (nth ea i)) false
      :else (er-ext-lt-seq ea eb (+ i 1)))))

(define
  er-ext-lt?
  (fn (a b)
    (cond
      (and (er-tuple? a) (er-tuple? b))
        (let ((ea (get a :elements)) (eb (get b :elements)))
          (cond
            (< (len ea) (len eb)) true
            (> (len ea) (len eb)) false
            :else (er-ext-lt-seq ea eb 0)))
      (and (er-cons? a) (er-cons? b))
        (cond
          (er-ext-lt? (get a :head) (get b :head)) true
          (er-ext-lt? (get b :head) (get a :head)) false
          :else (er-ext-lt? (get a :tail) (get b :tail)))
      (and (er-nil? a) (er-cons? b)) true
      (and (er-cons? a) (er-nil? b)) false
      (and (er-nil? a) (er-nil? b)) false
      :else (er-lt? a b))))

;; Default Erlang term order: a =< b  ==  not (b < a).
(define
  er-ext-term-le
  (fn (a b) (er-bool (not (er-ext-lt? b a)))))

;; ── lists:sort/1, lists:sort/2 ────────────────────────────────────
(define
  er-bif-lists-sort
  (fn (vs)
    (cond
      (= (len vs) 1)
        (er-sxlist->cons
          (er-ext-msort (er-cons->sxlist (nth vs 0)) er-ext-term-le))
      (= (len vs) 2)
        (let ((f (nth vs 0)) (lst (nth vs 1)))
          (er-sxlist->cons
            (er-ext-msort
              (er-cons->sxlist lst)
              (fn (a b) (er-apply-fun f (list a b))))))
      :else (error "Erlang: lists:sort: wrong arity"))))

;; ── lists:usort/1 (sort then drop adjacent term-equal dups) ───────
(define
  er-ext-dedup
  (fn (xs)
    (cond
      (= (len xs) 0) (list)
      (= (len xs) 1) xs
      (er-equal? (first xs) (nth xs 1)) (er-ext-dedup (rest xs))
      :else (cons (first xs) (er-ext-dedup (rest xs))))))

(define
  er-bif-lists-usort
  (fn (vs)
    (let ((lst (er-bif-arg1 vs "lists:usort")))
      (er-sxlist->cons
        (er-ext-dedup
          (er-ext-msort (er-cons->sxlist lst) er-ext-term-le))))))

;; ── keylists (lists of tuples keyed on element N, 1-indexed) ──────
;; keyfind/keymember/keydelete/keyreplace/keystore/keytake/keysort.
;; Key comparison is == (er-equal?), matching the standard lib. Only
;; the FIRST matching tuple is acted on. Non-tuples / tuples shorter
;; than N never match and are passed through unchanged.
(define
  er-ext-tup-elem
  (fn (tup n)
    (if (er-tuple? tup)
      (let ((es (get tup :elements)))
        (if (and (>= n 1) (<= n (len es))) (nth es (- n 1)) nil))
      nil)))

(define
  er-ext-key-match?
  (fn (key n tup)
    (and
      (er-tuple? tup)
      (>= n 1)
      (<= n (len (get tup :elements)))
      (er-equal? key (nth (get tup :elements) (- n 1))))))

(define
  er-ext-keyfind
  (fn (key n lst)
    (cond
      (er-nil? lst) (er-mk-atom "false")
      (er-cons? lst)
        (if (er-ext-key-match? key n (get lst :head))
          (get lst :head)
          (er-ext-keyfind key n (get lst :tail)))
      :else (er-mk-atom "false"))))

(define
  er-ext-keydelete
  (fn (key n lst)
    (cond
      (er-nil? lst) (er-mk-nil)
      (er-cons? lst)
        (if (er-ext-key-match? key n (get lst :head))
          (get lst :tail)
          (er-mk-cons (get lst :head) (er-ext-keydelete key n (get lst :tail))))
      :else lst)))

(define
  er-ext-keyreplace
  (fn (key n lst new)
    (cond
      (er-nil? lst) (er-mk-nil)
      (er-cons? lst)
        (if (er-ext-key-match? key n (get lst :head))
          (er-mk-cons new (get lst :tail))
          (er-mk-cons (get lst :head) (er-ext-keyreplace key n (get lst :tail) new)))
      :else lst)))

(define
  er-ext-keystore
  (fn (key n lst new)
    (cond
      (er-nil? lst) (er-mk-cons new (er-mk-nil))
      (er-cons? lst)
        (if (er-ext-key-match? key n (get lst :head))
          (er-mk-cons new (get lst :tail))
          (er-mk-cons (get lst :head) (er-ext-keystore key n (get lst :tail) new)))
      :else lst)))

(define
  er-bif-lists-keyfind
  (fn (vs) (er-ext-keyfind (nth vs 0) (nth vs 1) (nth vs 2))))

(define
  er-bif-lists-keymember
  (fn (vs)
    (er-bool (not (er-atom? (er-ext-keyfind (nth vs 0) (nth vs 1) (nth vs 2)))))))

(define
  er-bif-lists-keydelete
  (fn (vs) (er-ext-keydelete (nth vs 0) (nth vs 1) (nth vs 2))))

(define
  er-bif-lists-keyreplace
  (fn (vs) (er-ext-keyreplace (nth vs 0) (nth vs 1) (nth vs 2) (nth vs 3))))

(define
  er-bif-lists-keystore
  (fn (vs) (er-ext-keystore (nth vs 0) (nth vs 1) (nth vs 2) (nth vs 3))))

(define
  er-bif-lists-keytake
  (fn (vs)
    (let ((key (nth vs 0)) (n (nth vs 1)) (lst (nth vs 2)))
      (let ((hit (er-ext-keyfind key n lst)))
        (if (er-atom? hit)
          (er-mk-atom "false")
          (er-mk-tuple
            (list (er-mk-atom "value") hit (er-ext-keydelete key n lst))))))))

(define
  er-bif-lists-keysort
  (fn (vs)
    (let ((n (nth vs 0)) (lst (nth vs 1)))
      (er-sxlist->cons
        (er-ext-msort
          (er-cons->sxlist lst)
          (fn (a b)
            (er-bool
              (not (er-ext-lt? (er-ext-tup-elem b n) (er-ext-tup-elem a n))))))))))

;; ── higher-order traversal (foldr / partition / *while) ───────────
(define
  er-ext-foldr
  (fn (f acc lst)
    (cond
      (er-nil? lst) acc
      (er-cons? lst)
        (er-apply-fun f (list (get lst :head) (er-ext-foldr f acc (get lst :tail))))
      :else (raise (er-mk-error-marker (er-mk-atom "badarg"))))))

(define
  er-bif-lists-foldr
  (fn (vs) (er-ext-foldr (nth vs 0) (nth vs 1) (nth vs 2))))

(define
  er-ext-partition
  (fn (pred lst yes no)
    (cond
      (er-nil? lst)
        (er-mk-tuple
          (list
            (er-list-reverse-iter yes (er-mk-nil))
            (er-list-reverse-iter no (er-mk-nil))))
      (er-cons? lst)
        (if (er-truthy? (er-apply-fun pred (list (get lst :head))))
          (er-ext-partition pred (get lst :tail) (er-mk-cons (get lst :head) yes) no)
          (er-ext-partition pred (get lst :tail) yes (er-mk-cons (get lst :head) no)))
      :else (raise (er-mk-error-marker (er-mk-atom "badarg"))))))

(define
  er-bif-lists-partition
  (fn (vs) (er-ext-partition (nth vs 0) (nth vs 1) (er-mk-nil) (er-mk-nil))))

(define
  er-ext-takewhile
  (fn (pred lst)
    (cond
      (er-nil? lst) (er-mk-nil)
      (er-cons? lst)
        (if (er-truthy? (er-apply-fun pred (list (get lst :head))))
          (er-mk-cons (get lst :head) (er-ext-takewhile pred (get lst :tail)))
          (er-mk-nil))
      :else (er-mk-nil))))

(define
  er-bif-lists-takewhile
  (fn (vs) (er-ext-takewhile (nth vs 0) (nth vs 1))))

(define
  er-ext-dropwhile
  (fn (pred lst)
    (cond
      (er-nil? lst) (er-mk-nil)
      (er-cons? lst)
        (if (er-truthy? (er-apply-fun pred (list (get lst :head))))
          (er-ext-dropwhile pred (get lst :tail))
          lst)
      :else lst)))

(define
  er-bif-lists-dropwhile
  (fn (vs) (er-ext-dropwhile (nth vs 0) (nth vs 1))))

(define
  er-bif-lists-splitwith
  (fn (vs)
    (let ((pred (nth vs 0)) (lst (nth vs 1)))
      (er-mk-tuple
        (list (er-ext-takewhile pred lst) (er-ext-dropwhile pred lst))))))

;; ── register ──────────────────────────────────────────────────────
;; Hook into er-register-builtin-bifs! rather than registering once:
;; the registry can be reset + rebuilt mid-run (tests/runtime.sx does
;; this), and a plain one-shot registration would be wiped. Wrapping
;; the rebuild fn means these BIFs are re-added on every reset.
(define er-ext-lists-register!
  (fn ()
    (er-register-pure-bif! "lists" "sort" 1 er-bif-lists-sort)
    (er-register-pure-bif! "lists" "sort" 2 er-bif-lists-sort)
    (er-register-pure-bif! "lists" "usort" 1 er-bif-lists-usort)
    (er-register-pure-bif! "lists" "keyfind" 3 er-bif-lists-keyfind)
    (er-register-pure-bif! "lists" "keymember" 3 er-bif-lists-keymember)
    (er-register-pure-bif! "lists" "keydelete" 3 er-bif-lists-keydelete)
    (er-register-pure-bif! "lists" "keyreplace" 4 er-bif-lists-keyreplace)
    (er-register-pure-bif! "lists" "keystore" 4 er-bif-lists-keystore)
    (er-register-pure-bif! "lists" "keytake" 3 er-bif-lists-keytake)
    (er-register-pure-bif! "lists" "keysort" 2 er-bif-lists-keysort)
    (er-register-pure-bif! "lists" "foldr" 3 er-bif-lists-foldr)
    (er-register-pure-bif! "lists" "partition" 2 er-bif-lists-partition)
    (er-register-pure-bif! "lists" "takewhile" 2 er-bif-lists-takewhile)
    (er-register-pure-bif! "lists" "dropwhile" 2 er-bif-lists-dropwhile)
    (er-register-pure-bif! "lists" "splitwith" 2 er-bif-lists-splitwith)))

(define er-ext-prev-register-builtins er-register-builtin-bifs!)
(define er-register-builtin-bifs!
  (fn ()
    (er-ext-prev-register-builtins)
    (er-ext-lists-register!)))

;; register into the currently-live registry too
(er-ext-lists-register!)