import { _mostCommon } from "agency-lang/stdlib-lib/builtins.js"

/** @module
Run a block several times and combine the results. `sample` collects
every result, `consensus` takes the majority vote, and `retry` /
`retryWithFeedback` re-run a block until one result passes a test.

```ts
import { consensus } from "std::strategy"

node main() {
  // Run the same prompt 5 times and keep the majority answer.
  let answer = consensus(5) as {
    return llm("Is 17 prime? Answer yes or no.")
  }
  print(answer)
}
```
*/

// Private helper: return the most common element in an array (majority vote).
// `mostCommon` used to be a public auto-imported stdlib function. We removed
// it as too niche to auto-import, and consensus() is now its only consumer.
def mostCommon(items: any[]): any {
  return _mostCommon(items)
}

export def sample(n: number, block: () -> any): any[] {
  """
  Run a block n times in parallel. Returns an array of all results.

  @param n - Number of times to run
  @param block - The block to execute
  """
  return fork(range(n)) as _ {
    return block()
  }
}

export def consensus(n: number, block: () -> any): any {
  """
  Run a block n times in parallel and return the most common result (majority vote).

  @param n - Number of times to run
  @param block - The block to execute
  """
  let results = sample(n, block)
  return mostCommon(results)
}

export def retry(n: number, test: (any) -> boolean, block: () -> any): any {
  """
  Run a block up to n times. Returns the first result that passes the test function. Returns null if all attempts fail.

  @param n - Maximum number of attempts
  @param test - The function that returns true when the result is acceptable
  @param block - The block to execute
  """
  for (i in range(n)) {
    let result = block()
    let passed = test(result)
    if (passed) {
      return result
    }
  }
  return null
}

export def retryWithFeedback(
  n: number,
  test: (any) -> boolean,
  block: (any, number) -> any,
): any {
  """
  Run a block up to n times. Each attempt receives the previous result and the attempt number (starting from 1). Returns the first result that passes the test, or the last result if all fail.

  @param n - Maximum number of attempts
  @param test - The function that returns true when the result is acceptable
  @param block - The block receiving (previousResult, attemptNumber)
  """
  let prev = null
  for (i in range(n)) {
    let result = block(prev, i + 1)
    let passed = test(result)
    if (passed) {
      return result
    }
    prev = result
  }
  return prev
}

/** What a `revise` check decides about one attempt: whether to accept it,
  and if not, the critique the next attempt should address. `feedback` is
  ignored when `accepted` is true. */
export type Critique = {
  accepted: boolean;
  feedback: string
}

export def revise(
  maxAttempts: number,
  check: (any) -> Critique,
  block: (string) -> any,
): any {
  """
  Run a block until its result is accepted, up to maxAttempts times. Each
  rejected attempt passes its critique to the next one, which receives it as
  its argument (an empty string on the first attempt). Returns the first
  accepted result, or the last result when every attempt is rejected.

  @param maxAttempts - Maximum number of attempts
  @param check - Judges a result and returns the critique to address next
  @param block - The work, receiving the previous critique
  """
  let result = null
  let feedback = ""
  let attempt = 0
  while (attempt < maxAttempts) {
    attempt = attempt + 1
    result = block(feedback)
    const critique = check(result)
    if (critique.accepted) {
      return result
    }
    feedback = critique.feedback
  }
  // Every attempt was rejected. The last one is still the best available
  // answer, and it is more useful to the caller than nothing.
  return result
}

export def firstValid(
  variants: any[],
  test: (any) -> boolean,
  block: (any) -> any,
): any {
  """
  Run a block for each variant in parallel, then return the first result that passes the test. Returns null if none pass.

  @param variants - Array of variants to try
  @param test - The function that returns true for valid results
  @param block - The block receiving each variant
  """
  let results = fork(variants) as v {
    return block(v)
  }
  for (r in results) {
    let passed = test(r)
    if (passed) {
      return r
    }
  }
  return null
}