///|
/// A keyed memoization map backed by one Memo per key.
///
/// `MemoMap[K, V]` provides a minimal parameterized-query pattern:
/// each key gets its own lazily-created `Memo[V]`, and subsequent reads
/// of that key reuse the cached memoized value.
///
/// Reading a key requires `K : Hash + Eq` and `V : Eq`, matching the
/// underlying `HashMap` lookup and `Memo` value semantics.
///
/// # Example
///
/// ```moonbit nocheck
/// let rt = Runtime()
/// let source = Signal(rt, 10)
/// let by_id = MemoMap(rt, (id : Int) => source.get() + id)
///
/// inspect(by_id.get(1), content="11")
/// inspect(by_id.get(1), content="11")  // cache hit for key=1
/// inspect(by_id.get(2), content="12")  // separate memo for key=2
/// ```
pub(all) struct MemoMap[K, V] {
  priv rt : Runtime
  priv compute : (K) -> V raise Failure
  priv label : String?
  priv entries : @hashmap.HashMap[K, Memo[V]]
}

///|
/// Constructor-syntax alias of `MemoMap::new`. Enables `MemoMap(rt, compute)`
/// / `MemoMap(rt, compute, label="...")` call sites.
pub fn[K : Hash + Eq, V] MemoMap::MemoMap(
  rt : Runtime,
  compute : (K) -> V raise Failure,
  label? : String,
) -> MemoMap[K, V] {
  MemoMap::new(rt, compute, label?)
}

///|
/// Creates a new empty MemoMap.
///
/// Entries are created lazily on first `get`/`get_result` per key.
pub fn[K : Hash + Eq, V] MemoMap::new(
  rt : Runtime,
  compute : (K) -> V raise Failure,
  label? : String,
) -> MemoMap[K, V] {
  { rt, compute, label, entries: @hashmap.HashMap([]) }
}

///|
/// Returns the value for `key`, computing lazily on first access.
///
/// Permissive read for test code and other top-level call sites.
/// If called while a tracking frame is active, records the per-key dependency.
pub fn[K : Hash + Eq, V : Eq] MemoMap::get(self : MemoMap[K, V], key : K) -> V {
  let value = self.get_or_create_memo(key).read_permissive()
  if self.rt.has_pending_events() {
    self.rt.drain_pending_events_if_idle()
  }
  value
}

///|
/// Returns the value for `key`, aborting if called outside a tracked context.
///
/// Stricter-contract peer of `get()`. Both methods record a dependency on
/// the per-key memo when invoked from inside another memo's compute
/// function; the difference is what happens at top level — `get()` reads
/// silently without recording a dep, whereas `get_tracked()` aborts via
/// `Memo::get`'s tracking guard. Prefer this when the call site is
/// expected to be tracking (catches "meant to call from inside a memo
/// but I'm not" bugs).
///
/// Caveat: `MemoMap::clear()` unconditionally disposes every per-key
/// wrapper, so a tracked dependent recorded via this method will hit
/// the disposed-dep guard on its next verify if `clear()` runs between
/// reads. `gc()` is safe — it keeps wrapper memos alive as long as a
/// tracked dependent is reachable from a root. The same caveat applies
/// to `get()`.
pub fn[K : Hash + Eq, V : Eq] MemoMap::get_tracked(
  self : MemoMap[K, V],
  key : K,
) -> V {
  self.get_or_create_memo(key).get()
}

///|
/// Returns the value for `key` as a Result, matching `Memo::get_result`.
///
/// Permissive read for callers that want cycle errors as values.
pub fn[K : Hash + Eq, V : Eq] MemoMap::get_result(
  self : MemoMap[K, V],
  key : K,
) -> Result[V, CycleError] {
  self.read_result(key)
}

///|
/// Package-private strict honest read. Requires an active tracked context;
/// reports mechanism failures (cycle / disposed per-key memo) as `ReadError`.
fn[K : Hash + Eq, V : Eq] MemoMap::get_strict_honest(
  self : MemoMap[K, V],
  key : K,
) -> Result[V, ReadError] {
  guard self.rt.core.tracking.stack.length() > 0 else {
    abort(
      "MemoMap::get() called outside tracked context. Use DerivedMap::read(key) or DerivedMap::read_or_abort(key) to read from outside the graph.",
    )
  }
  self.read_honest(key)
}

///|
/// Package-private honest permissive read — the single source of truth for the
/// read channel. A read of a per-key memo that gc disposed (while the map still
/// caches the entry) returns `Err(Disposed(_))`; a cycle returns
/// `Err(Cycle(_))`; a compute `raise Failure` is a defect and still aborts.
fn[K : Hash + Eq, V : Eq] MemoMap::read_honest(
  self : MemoMap[K, V],
  key : K,
) -> Result[V, ReadError] {
  let memo = self.get_or_create_memo(key)
  guard !memo.is_disposed() else { return Err(ReadError::disposed(memo.id())) }
  let result = memo.get_result_inner() catch {
    e => {
      self.rt.drain_pending_events_direct()
      abort("MemoMap compute raised: " + e.to_string())
    }
  }
  if self.rt.has_pending_events() {
    self.rt.drain_pending_events_if_idle()
  }
  match result {
    Ok(v) => Ok(v)
    Err(e) => Err(ReadError::cycle(e))
  }
}

///|
/// Package-private strict Result read. Requires an active tracked context.
/// Compatibility projection of `get_strict_honest` onto `CycleError`.
fn[K : Hash + Eq, V : Eq] MemoMap::get_strict_result(
  self : MemoMap[K, V],
  key : K,
) -> Result[V, CycleError] {
  guard self.rt.core.tracking.stack.length() > 0 else {
    abort(
      "MemoMap::get() called outside tracked context. Use DerivedMap::read(key) or DerivedMap::read_or_abort(key) to read from outside the graph.",
    )
  }
  self.read_result(key)
}

///|
/// Package-private permissive Result read. Records a dependency when tracked.
/// Compatibility projection of `read_honest` onto `CycleError`: a disposed
/// per-key memo aborts (legacy contract), a cycle is `Err(CycleError)`.
fn[K : Hash + Eq, V : Eq] MemoMap::read_result(
  self : MemoMap[K, V],
  key : K,
) -> Result[V, CycleError] {
  match self.read_honest(key) {
    Ok(v) => Ok(v)
    Err(ReadError::Cycle(e)) => Err(e)
    Err(ReadError::Disposed(_)) =>
      abort("MemoMap::get called on a disposed memo")
  }
}

///|
/// Returns the value for `key`, or `fallback` if a cycle is detected.
pub fn[K : Hash + Eq, V : Eq] MemoMap::get_or(
  self : MemoMap[K, V],
  key : K,
  fallback : V,
) -> V {
  match self.get_result(key) {
    Ok(v) => v
    Err(_) => fallback
  }
}

///|
/// Returns the value for `key`, or computes a fallback from the cycle error.
pub fn[K : Hash + Eq, V : Eq] MemoMap::get_or_else(
  self : MemoMap[K, V],
  key : K,
  fallback : (CycleError) -> V,
) -> V {
  match self.get_result(key) {
    Ok(v) => v
    Err(e) => fallback(e)
  }
}

///|
/// Returns true if a memo entry for `key` has already been created.
pub fn[K : Hash + Eq, V] MemoMap::contains(
  self : MemoMap[K, V],
  key : K,
) -> Bool {
  self.entries.contains(key)
}

///|
/// Returns the number of memo entries that have been created so far.
pub fn[K, V] MemoMap::length(self : MemoMap[K, V]) -> Int {
  self.entries.length()
}

///|
/// Returns the memo for `key`, creating and caching it on first access.
fn[K : Hash + Eq, V : Eq] MemoMap::get_or_create_memo(
  self : MemoMap[K, V],
  key : K,
) -> Memo[V] {
  match self.entries.get(key) {
    Some(memo) => memo
    None => {
      // Each key gets a dedicated memo; key is captured by value in closure.
      let memo = match self.label {
        Some(label) => Memo::new(self.rt, () => (self.compute)(key), label~)
        None => Memo::new(self.rt, () => (self.compute)(key))
      }
      self.entries.set(key, memo)
      memo
    }
  }
}

///|
/// Removes entries whose memos have been disposed (e.g. by gc()).
///
/// Two-pass: collects keys to remove, then removes them (cannot mutate
/// HashMap during iteration). Returns the number of entries swept.
///
/// Call after `Runtime::gc()` to clean up stale entries.
pub fn[K : Hash + Eq, V] MemoMap::sweep(self : MemoMap[K, V]) -> Int {
  let to_remove : Array[K] = []
  for key, memo in self.entries {
    if memo.is_disposed() {
      to_remove.push(key)
    }
  }
  for key in to_remove {
    self.entries.remove(key)
  }
  to_remove.length()
}

///|
/// Disposes every wrapper memo and empties the entry map.
///
/// Unlike `sweep()`, which only removes entries whose memos are already
/// disposed, `clear()` actively disposes each wrapper memo. Use this after
/// a structural event (e.g. pipeline rebuild) that invalidates all cached
/// keys — `sweep()` is insufficient because MemoMap's wrappers are
/// Runtime-owned and are not reached by chain-scope disposal.
pub fn[K, V] MemoMap::clear(self : MemoMap[K, V]) -> Unit {
  for _, memo in self.entries {
    memo.dispose()
  }
  self.entries.clear()
}