plum

module std

import std/Option
import std/Buffer
import std/Number
import std/Bool
import std/Str

# A node stores the data in a list and contains pointers to the previous and next sibling nodes
enum Node[T] =
  | Node(value: T, prev: Option[Node[T]], next: Option[Node[T]])

# A list is a data structure describing a contiguous section of an array stored separately from the slice variable itself.
# It contains the pointers to the start and end nodes (head, tail) and maintains the size as well
enum List[T: ToStr](ToStr) =
  | List(head: Option[Node[T]], tail: Option[Node[T]], size: Int)

  # `List[T](1, 2, 3)` (or, for an empty list, `List[T]()` — explicit
  # generics required either way, since nothing else in a call this shape
  # pins down `T`) constructs a list via this variadic `init` — see
  # `plum-checker`/`plum-wasm-codegen`'s `ClassCall`/`FnCall` handling, which
  # desugars both call shapes to `List.init(...)`.
  fun init(values: ...T) -> List[T] =
    l := List(head: None, tail: None, size: 0)
    for v := range values
      l.add(v)
    return l

  # gets the element at i'th index of the list
  fun get(self, i: Int) -> Option[T] =
    current := self.head
    index := 0
    while current != None
      match current
        Some(Node(value, _, next)) =>
          if index == i
            return Some(value)
          current = next
          index = index + 1
        None =>
          break
    None

  # sets the element at i'th index of the list, returning the old value (or None if i is out of bounds)
  fun set(self, i: Int, v: T) -> Option[T] =
    current := self.head
    index := 0
    while current != None
      match current
        Some(node) =>
          if index == i
            old := node.value
            node.value = v
            return Some(old)
          current = node.next
          index = index + 1
        None =>
          break
    None

  # returns the no of elements in the list
  fun length(self) -> Int =
    self.size

  # Elementwise comparison via `==` — see the caveat on `expectEq` in
  # `libs/std/testing.plum` for what that means when `T` is a class/enum
  # without its own `equals` (reference identity, not structural equality).
  fun equals(self, other: List[T]) -> Bool =
    if self.length() != other.length()
      return False
    for i := range self.length()
      if self.get(i).unwrap() != other.get(i).unwrap()
        return False
    return True

  # adds the specified elements to the end of the list
  fun add(self, values: ...T) =
    for v := range values
      node := Node(value: v, prev: self.tail, next: None)
      match self.tail
        Some(t) =>
          t.next = Some(node)
        None =>
          self.head = Some(node)
      self.tail = Some(node)
      self.size = self.size + 1

  # unlinks node from the list, patching its neighbors (or head/tail) to close the gap
  fun unlink(self, node: Node[T]) =
    match node.prev
      Some(p) =>
        p.next = node.next
      None =>
        self.head = node.next
    match node.next
      Some(n) =>
        n.prev = node.prev
      None =>
        self.tail = node.prev
    self.size = self.size - 1

  # removes the element at i'th index of the list
  fun removeAt(self, i: Int) =
    current := self.head
    index := 0
    while current != None
      match current
        Some(node) =>
          if index == i
            self.unlink(node)
            return
          current = node.next
          index = index + 1
        None =>
          break

  # removes the first element equal to v from list
  fun remove(self, v: T) =
    current := self.head
    while current != None
      match current
        Some(node) =>
          if node.value == v
            self.unlink(node)
            return
          current = node.next
        None =>
          break

  # removes all objects from this list
  fun clear(self) =
    self.head = None
    self.tail = None
    self.size = 0

  # returns the list with the elements relinked in reverse order.
  fun reverse(self) -> List =
    current := self.head
    while current != None
      match current
        Some(node) =>
          next := node.next
          node.next = node.prev
          node.prev = next
          current = next
        None =>
          break
    old_head := self.head
    self.head = self.tail
    self.tail = old_head
    self

  # Inserts `value` immediately before `node` (which must belong to `self`),
  # patching neighbors (or `head`) to make room — the insertion-point half of
  # `sort`'s pointer surgery, split out since `unlink`'s counterpart already
  # gets its own method.
  fun insertBefore(self, node: Node[T], value: T) -> Unit =
    new_node := Node(value: value, prev: node.prev, next: Some(node))
    match node.prev
      Some(p) =>
        p.next = Some(new_node)
      None =>
        self.head = Some(new_node)
    node.prev = Some(new_node)
    self.size = self.size + 1

  # Inserts `value` into `self` (assumed already sorted by `before`) at the
  # position that keeps it sorted — the first position whose existing element
  # `before(value, existing)` holds for, or the end if none does.
  fun insertSorted(self, value: T, before: fn(T, T) -> Bool) -> Unit =
    current := self.head
    while current != None
      match current
        Some(node) =>
          if before(value, node.value)
            self.insertBefore(node, value)
            return
          current = node.next
        None =>
          break
    self.add(value)

  # Returns a NEW list with `self`'s elements sorted, using `before(a, b)` as
  # the "does a belong before b" comparator. Insertion sort (O(n^2)) — simple
  # and adequate for the list sizes a linked list is already suited to.
  fun sort(self, before: fn(T, T) -> Bool) -> List[T] =
    result := List(head: None, tail: None, size: 0)
    current := self.head
    while current != None
      match current
        Some(Node(value, _, next)) =>
          result.insertSorted(value, before)
          current = next
        None =>
          break
    return result

  # Returns the first element equal to `search`, or `None`.
  fun find(self, search: T) -> Option[T] =
    current := self.head
    while current != None
      match current
        Some(Node(value, _, next)) =>
          if value == search
            return Some(value)
          current = next
        None =>
          break
    return None

  fun contains(self, v: T) -> Bool =
    return self.find(v) != None

  # calls f for each elem in the list
  fun each(self, cb: fn(T)) -> Unit =
    current := self.head
    while current != None
      match current
        Some(Node(value, _, next)) =>
          cb(value)
          current = next
        None =>
          break

  # Returns a new list with each element transformed by `cb`.
  #
  # `U` here is a generic param of `map` ITSELF, separate from `List[T]`'s own
  # `T` — resolved per CALL SITE (from `cb`'s own inferred return type), not
  # when `List[T]` itself is specialized, via `resolveMethodOwnGenerics` in
  # `plum-checker/src/monomorphize.rs`.
  fun map(self, cb: fn(T) -> U) -> List[U] =
    result := List(head: None, tail: None, size: 0)
    current := self.head
    while current != None
      match current
        Some(Node(value, _, next)) =>
          result.add(cb(value))
          current = next
        None =>
          break
    return result

  # returns a new list with each element flat-mapped
  fun flatMap(self) =
    todo

  # Returns a new list with only the elements `predicate` holds for.
  fun retain(self, predicate: fn(T) -> Bool) -> List[T] =
    result := List(head: None, tail: None, size: 0)
    current := self.head
    while current != None
      match current
        Some(Node(value, _, next)) =>
          if predicate(value)
            result.add(value)
          current = next
        None =>
          break
    return result

  # Returns a new list with the elements `predicate` holds for removed.
  fun reject(self, predicate: fn(T) -> Bool) -> List[T] =
    result := List(head: None, tail: None, size: 0)
    current := self.head
    while current != None
      match current
        Some(Node(value, _, next)) =>
          if !predicate(value)
            result.add(value)
          current = next
        None =>
          break
    return result

  fun any(self, predicate: fn(T) -> Bool) -> Bool =
    current := self.head
    while current != None
      match current
        Some(Node(value, _, next)) =>
          if predicate(value)
            return True
          current = next
        None =>
          break
    return False

  fun every(self, predicate: fn(T) -> Bool) -> Bool =
    current := self.head
    while current != None
      match current
        Some(Node(value, _, next)) =>
          if !predicate(value)
            return False
          current = next
        None =>
          break
    return True

  # Folds `self`'s elements into a single value, starting from `acc` and
  # combining each element in with `cb(accumulatedSoFar, element)`.
  #
  # `U` is `reduce`'s OWN generic param (the accumulator's type), separate
  # from `List[T]`'s own `T` — same per-CALL-SITE resolution as `map`'s `U`.
  fun reduce(self, acc: U, cb: fn(U, T) -> U) -> U =
    result := acc
    current := self.head
    while current != None
      match current
        Some(Node(value, _, next)) =>
          result = cb(result, value)
          current = next
        None =>
          break
    return result

  # returns the first element in the list
  fun first(self) -> Option[T] =
    match self.head
      Some(Node(value, _, _)) => Some(value)
      None => None

  # returns the last element in the list
  fun last(self) -> Option[T] =
    match self.tail
      Some(Node(value, _, _)) => Some(value)
      None => None

  # Returns a new list of the elements at index `[start, end)`, clamped into
  # range (a `start`/`end` outside `[0, length]` is clamped, not an error).
  fun sublist(self, start: Int, end: Int) -> List[T] =
    result := List(head: None, tail: None, size: 0)
    s := start < 0 ? 0 : start
    e := end > self.size ? self.size : end
    current := self.head
    i := 0
    while current != None && i < e
      match current
        Some(Node(value, _, next)) =>
          if i >= s
            result.add(value)
          current = next
          i = i + 1
        None =>
          break
    return result

  # The first `n` elements.
  fun take(self, n: Int) -> List[T] =
    return self.sublist(0, n)

  # Every element AFTER the first `n`.
  fun skip(self, n: Int) -> List[T] =
    return self.sublist(n, self.size)

  fun drop(self, n: Int) -> List[T] =
    return self.skip(n)

  # Returns one uniformly random element, or `None` if `self` is empty.
  fun sample(self) -> Option[T] =
    if self.size == 0
      return None
    return self.get(randomInt(self.size))

  # Returns a new list with `self`'s elements in random order (Fisher-Yates).
  fun shuffle(self) -> List[T] =
    result := self.sublist(0, self.size)
    i := result.size - 1
    while i > 0
      j := randomInt(i + 1)
      vi := result.get(i).unwrap()
      vj := result.get(j).unwrap()
      result.set(i, vj)
      result.set(j, vi)
      i = i - 1
    return result

  # `chunk`/`partition` both return `List[List[T]]` — the SAME generic class
  # (`List`) nested inside itself as a method's own return type. This used to
  # be an unconditional stack-overflow landmine for the WHOLE COMPILER (see
  # `plum-checker/src/monomorphize.rs`'s `isSelfNested`/`methodMentionsTemplate`
  # doc comments for the full fix) — that danger is now FIXED (2026-09-02).
  #
  # A second, more contained bug surfaced once past that: two SEPARATE bare
  # `List(head: None, tail: None, size: 0)` constructions in the SAME method
  # body — one for the OUTER `result: List[List[T]]` (matching this
  # method's own declared return type) and one for an INNER per-chunk
  # `piece: List[T]` — are ambiguous to `resolveClassInstantiation`'s
  # `current_return_type` fallback: since `List[T]` and `List[List[T]]` both
  # have exactly ONE generic argument, the fallback (which only ever
  # compares ARITY, not identity) can't tell "this bare construction IS the
  # function's own return value" apart from "this is some OTHER, unrelated
  # nested construction of the same template that happens to need the same
  # arity" — so `piece` was silently resolved to the SAME type as `result`
  # (`List$List$Int` instead of `List$Int`), corrupting its own `add` calls.
  # Worked around by building `piece` via `self.sublist(...)` instead of a
  # second bare construction — `sublist` is an ordinary method already
  # unambiguously resolved through the RECEIVER's own (already-correct)
  # binding, no `current_return_type` guessing involved at all.
  fun chunk(self, size: Int) -> List[List[T]] =
    result := List(head: None, tail: None, size: 0)
    if size <= 0
      return result
    i := 0
    n := self.length()
    while i < n
      result.add(self.sublist(i, i + size))
      i = i + size
    return result

  # Splits into two lists: elements `predicate` holds for, then the rest —
  # in that order. Returns them as a 2-element `List[List[T]]` (`[matched,
  # unmatched]`) since this language has no tuple type. Deliberately built
  # from `self.retain(predicate)`/`self.reject(predicate)` rather than two
  # local `matched`/`unmatched := List(head: None, ...)` bare constructions
  # — see `chunk`'s doc comment above for exactly why that would be
  # ambiguous (`matched`/`unmatched` need `List[T]`, but a bare
  # construction here would resolve via `current_return_type` fallback
  # against THIS method's own `List[List[T]]` return instead, same arity
  # coincidence). `retain`/`reject` are each already unambiguous (single
  # bare construction per method body, matching their own single-level
  # return type).
  fun partition(self, predicate: fn(T) -> Bool) -> List[List[T]] =
    result := List(head: None, tail: None, size: 0)
    result.add(self.retain(predicate))
    result.add(self.reject(predicate))
    return result

  # Grouping by an arbitrary key type would need its own generic param
  # (`groupBy<K>(keyFn: fn(T) -> K) -> Map[K, List[T]]`) resolved from a
  # closure's return type — the same method-level-generics gap documented on
  # `List.map` above, PLUS `Map[K, List[T]]` would make `list.plum` depend on
  # `map.plum`, which already depends on `list.plum`. Not implemented.
  fun groupBy(self) =
    todo

  # Note: NOT written with `each` + a closure-captured "is this the first
  # element" flag — a closure's captured variables are snapshotted at
  # CREATION time (see `libs/std/list.plum`'s own `add`/`each` usage
  # elsewhere), so an assignment to a captured var INSIDE the closure body
  # does not persist across that closure's own repeated invocations from one
  # `each` call. A plain `while`/`match` traversal (same idiom as `unlink`)
  # sidesteps that entirely.
  fun join(self, sep: Str = ",") -> Str =
    res := Buffer()
    current := self.head
    while current != None
      match current
        Some(Node(value, _, next)) =>
          res.write(value.toStr())
          if next != None
            res.write(sep)
          current = next
        None =>
          break
    res.toStr()

  # `List[T: ToStr](ToStr)` declared implementing `ToStr`
  # from the start but never actually defined `toStr` — harmless for a
  # PLAIN `List[Int]`/`List[Str]`/etc (nothing needs to call `.toStr()` on
  # the LIST ITSELF, just on its elements, which `join` above already
  # does), but a genuine gap the moment a `List` is used AS another
  # `List`'s own ELEMENT type (`List[List[T]]`) — the OUTER list's eagerly-
  # compiled `join`/`toStr` calls `.toStr()` on each element, and an inner
  # `List[T]` element had no such method to call, failing to compile with
  # "unknown method 'List$T.toStr'" even for a program that never actually
  # calls `.toStr()`/`.join()` on the outer list (per the "eager compile
  # every method for every specialization" architecture noted elsewhere in
  # this codebase's history).
  fun toStr(self) -> Str =
    return "[" + self.join(",") + "]"

# ---- regression tests ----

fun optSumForListTest(o: Option[Int]) -> Int =
  match o
    Some(v) => v
    None => -1000

fun exerciseList() -> Int =
  l = List(1, 2, 3, 4, 5)
  a = l.length()
  b = optSumForListTest(l.get(0))
  c = optSumForListTest(l.get(4))
  oldVal = optSumForListTest(l.set(2, 30))
  d = optSumForListTest(l.get(2))
  l.removeAt(0)
  e = l.length()
  f = optSumForListTest(l.get(0))
  l.remove(30)
  g = l.length()
  l.reverse()
  h = optSumForListTest(l.get(0))
  l.clear()
  i = l.length()
  a + b + c + oldVal + d + e + f + g + h + i

fun removeAllNodesOneAtATime() -> Int =
  l = List(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
  i = 0
  while i < 10
    l.removeAt(0)
    i = i + 1
  afterLoopLength = l.length()
  isEmpty = optSumForListTest(l.get(0))
  l.add(42)
  afterReAdd = optSumForListTest(l.get(0))
  afterLoopLength + isEmpty + afterReAdd

fun exerciseListInit() -> Int =
  l := List[Int]()
  before := l.length()
  l.add(7)
  before + l.length() + optSumForListTest(l.get(0))

fun exerciseListInitWithValues() -> Int =
  l := List(10, 20, 30)
  l.length() + optSumForListTest(l.get(0)) + optSumForListTest(l.get(2))

test "list add set remove at remove clear reverse all work correctly"
  # add(1,2,3,4,5): a=length=5, b=get(0)=1, c=get(4)=5
  # set(2,30): oldVal=3, d=get(2)=30 -> list [1,2,30,4,5]
  # removeAt(0): e=length=4, f=get(0)=2 -> list [2,30,4,5]
  # remove(30): g=length=3 -> list [2,4,5]
  # reverse(): h=get(0)=5 -> list [5,4,2]
  # clear(): i=length=0
  # 5+1+5+3+30+4+2+3+5+0 = 58
  assert exerciseList() == 58

test "removing every node in a loop leaves an empty correctly functioning list"
  # afterLoopLength=0, isEmpty(get(0) on empty list)=-1000, afterReAdd=42
  # 0 + -1000 + 42 = -958
  assert removeAllNodesOneAtATime() == -958

test "List[Int]() with no args runs init() and produces a working empty list"
  assert exerciseListInit() == 8

test "List(values) with positional args runs init(values) directly, infers T from args"
  assert exerciseListInitWithValues() == 43

test "list map and reduce work correctly"
  l = List(1, 2, 3, 4)
  doubled = l.map(|v| v * 2)
  total = doubled.reduce(0, |acc, v| acc + v)
  assert total == 20

test "list equals compares elementwise"
  a = List(1, 2, 3)
  b = List(1, 2, 3)
  c = List(1, 2)
  assert a.equals(b)
  assert !a.equals(c)

test "list sort orders elements correctly"
  l = List(3, 1, 4, 1, 5)
  sorted = l.sort(|a, b| a < b)
  assert optSumForListTest(sorted.get(0)) == 1
  assert optSumForListTest(sorted.get(4)) == 5

test "list join renders elements correctly with separator"
  l = List(1, 2, 3)
  assert l.join(",") == "1,2,3"

test "list toStr formats a plain list and a nested list of lists"
  l := List(1, 2, 3)
  assert l.toStr() == "[1,2,3]"
  chunks := l.chunk(2)
  assert chunks.toStr() == "[[1,2],[3]]"

test "chunk splits into sublists of at most size elements each"
  l := List(1, 2, 3, 4, 5)
  chunks := l.chunk(2)
  assert chunks.length() == 3
  assert chunks.get(0).unwrap().join(",") == "1,2"
  assert chunks.get(1).unwrap().join(",") == "3,4"
  assert chunks.get(2).unwrap().join(",") == "5"
  assert l.chunk(0).length() == 0

test "partition splits into elements matching a predicate, then the rest"
  l := List(1, 2, 3, 4, 5, 6)
  parts := l.partition(|v| v % 2 == 0)
  assert parts.length() == 2
  assert parts.get(0).unwrap().join(",") == "2,4,6"
  assert parts.get(1).unwrap().join(",") == "1,3,5"