import { _withLock } from "agency-lang/stdlib-lib/concurrency.js"

/** @module
  Concurrency primitives for coordinating work inside one Agency run. Use
  `withLock` to guard a shared resource so branches take turns instead of
  clobbering each other.

  ```ts
  import { withLock } from "std::concurrency"

  node main() {
    fork(range(5), shared: true) as _ {
      withLock("counter") as {
        // only one branch runs this block at a time
        updateSharedState()
      }
    }
  }
  ```
*/

/** Caller notes:
    - In subprocesses launched with `std::agency.run()`, the lock is coordinated by
      the parent process, so parent and child code share the same mutex.
    - Same-owner reentrancy on the same lock throws immediately.
    - Multi-lock deadlock cycles are not detected automatically; use `timeoutMs`
      when a bounded wait is required.
    - A fork branch may run inside a lock as long as it does not reacquire the same
      lock. Reacquiring the same lock from a fork spawned inside the lock can deadlock,
      because the outer scope waits for the fork while the branch waits for the outer
      lock to release. */
export def withLock(
  name: string,
  timeoutMs: number | null = null,
  warnAfterMs: number | null = null,
  block: () -> any = null,
): any {
  """
  Run a block while holding a named per-run mutex. Branches of the same run
  that use the same lock name execute this block one at a time. Branches
  using different names continue concurrently. The lock is released
  automatically when the block returns, throws, or unwinds for an interrupt.

  @param name - Lock name; use stable names like "std::tty" for shared resources
  @param timeoutMs - Maximum time to wait before failing without acquiring the lock; null waits indefinitely
  @param warnAfterMs - Wait time before printing a diagnostic warning; defaults to 30s
  @param block - The work to run while holding the lock
  """
  return _withLock(name, timeoutMs, warnAfterMs, block)
}

export def concurrencyPool(
  size: number,
  items: any[],
  block: (item: any) -> any,
): any[] {
  """
  Run a block in parallel across a bounded pool of workers.

  @param size - Number of concurrent workers to run
  @param items - The items to process
  @param block - The work to run in each worker
  """
  // fork branches only share module globals (with shared: true), never
  // function locals, so workers cannot pull from a shared queue or write
  // into a shared results array. Instead: deal the items round-robin into
  // one chunk per worker, fork over the chunks, and reassemble by index
  // from the branch RETURN values — the one channel fork reports through.
  // Pair each item with its index by hand: std map passes only the item to
  // its function, never an index.
  const indexed: any[] = []
  let index = 0
  for (item in items) {
    indexed.push([item, index])
    index = index + 1
  }
  const chunks: any[] = []
  let worker = 0
  while (worker < size) {
    chunks.push([])
    worker = worker + 1
  }
  let position = 0
  for (pair in indexed) {
    chunks[position % size].push(pair)
    position = position + 1
  }
  const chunkResults = fork(chunks) as chunk {
    const chunkOut: any[] = []
    for (pair in chunk) {
      chunkOut.push([pair[1], block(pair[0])])
    }
    return chunkOut
  }
  const results: any[] = map(items, \_ -> null)
  for (chunkOut in chunkResults) {
    for (entry in chunkOut) {
      results[entry[0]] = entry[1]
    }
  }
  return results
}