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

/** @module
  @summary Coordinate concurrent work with locks and worker pools.
  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"

  let counter = 0

  node main() {
    fork(range(5), shared: true) as _ {
      withLock("counter") as {
        counter = counter + 1
      }
    }
    print(counter)
  }
  ```
*/

/**
Use a stable, namespaced name such as "std::tty" for a shared resource.
Parent and child code launched with `std::agency.run()` share the lock.
The lock releases when the block returns, throws, or unwinds for an interrupt.

Reacquiring a held lock from the same owner throws. A fork inside a locked
block can deadlock if it tries to acquire that lock: the outer block waits
for the fork while holding the lock. Use `timeoutMs` to bound waits when
several locks could form a cycle.
*/
export def withLock(
  name: string,
  timeoutMs: number | null = null,
  warnAfterMs: number | null = null,
  block: () -> any = null,
): any {
  """
  Run a block while holding a named lock shared by branches of one run.

  @param timeoutMs - Maximum wait to acquire the lock, or null to wait indefinitely.
  @param warnAfterMs - Wait before a diagnostic warning, or null for the 30-second default.
  """
  return _withLock(name, timeoutMs, warnAfterMs, block)
}

/**
Results follow the input order. Items are assigned to workers round-robin,
so a worker processes its own items in sequence.
*/
export def concurrencyPool(
  size: number,
  items: any[],
  block: (item: any) -> any,
): any[] {
  """
  Process items concurrently with at most `size` workers.

  @param size - Positive whole number of workers.
  """
  // 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
}