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

/** @module
@summary Sample alternatives, retry work, and revise results using feedback.

Use `sample` to collect parallel attempts and `consensus` to choose the
most common result. `retry` accepts the first result that passes a test.
`revise` gives each new attempt the prior attempt's critique.

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

node main() {
  const answer = retry(3, test: \text -> text.length > 0) as {
    return llm("Suggest a short title for a gardening guide")
  }
  print(answer)
}
```

Function budget decorators live in `std::decorators`.
*/

// 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 several times in parallel and collect the results in attempt order.
  """
  return fork(range(n)) as _ {
    return block()
  }
}

export def consensus(n: number, block: () -> any): any {
  """
  Run a block several times in parallel and return its most common result.
  """
  let results = sample(n, block)
  return mostCommon(results)
}

export def retry(n: number, test: (any) -> boolean, block: () -> any): any {
  """
  Retry a block until its result passes a test.
  Return the first passing result, or null when no attempt passes.
  """
  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 {
  """
  Retry a block using the prior result to improve each attempt.
  Return the first passing result, or the last result if no attempt passes.

  @param block - Receives null on the first attempt, then the previous result, followed by a one-based attempt number.
  """
  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
}

/** The first attempt receives an empty critique. Later attempts receive the
feedback from the last rejected result. When a budget guard stops the work,
the last completed attempt can reach the guard as a partial result.
If no attempt completed, an attempt's saved draft can be used instead. */
export def revise(
  maxAttempts: number,
  check: (any) -> Critique,
  block: (string) -> any,
): any {
  """
  Revise work using critiques until a check accepts it or attempts run out.
  Return the accepted result, or the last result when all attempts are rejected.

  @param block - Receives an empty string first, then the prior critique's feedback.
  """
  // `result` is the last completed attempt. `latest` is the attempt in
  // progress: when a guard trip stops the block, the block's saved draft,
  // if any, is bound into it.
  let result = null
  let latest = null
  let completed = false
  let feedback = ""
  let attempt = 0
  while (attempt < maxAttempts) {
    attempt = attempt + 1
    latest = block(feedback)
    result = latest
    completed = true
    const critique = check(result)
    if (critique.accepted) {
      return result
    }
    feedback = critique.feedback
  }
  return result

  // A guard trip stops revise. Hand back the best answer so far: the last
  // completed attempt, or else the draft the unfinished attempt saved. With
  // neither, return nothing, so the guard fails as it would without revise.
  // See docs/dev/runtime/saveDraft.md.
  finalize {
    if (completed) {
      return result
    }
    if (latest != null) {
      return latest
    }
  }
}

export def firstValid(
  variants: any[],
  test: (any) -> boolean,
  block: (any) -> any,
): any {
  """
  Run a block for each variant in parallel and return the first passing result
  in input order, or null when none passes.
  """
  let results = fork(variants) as v {
    return block(v)
  }
  for (r in results) {
    let passed = test(r)
    if (passed) {
      return r
    }
  }
  return null
}