import { Feedback } from "std::agents/lib/feedback"
import {
  llmOptions,
  withContext,
  docsOnlyHandler,
 } from "std::agents/lib/shared"
import {
  writingReviewAgent,
  WritingReviewEvalInput,
 } from "std::agents/writing/review"
import { guarded } from "std::decorators"
import { evalValue, setAgentName } from "std::statelog"

/** @module
@summary Review prose and return a rewrite with the findings applied.

`writingRewriteAgent` reviews the text, then rewrites it when the review
has findings. `passes` bounds the number of review-and-rewrite rounds.
A review without findings ends the loop, and a failed review fails the call.

```ts
import { writingRewriteAgent } from "std::agents/writing/rewrite"

node main() {
  print(writingRewriteAgent(text: "The report was written by Alice."))
}
```
*/

static const rewritePrompt = """You rewrite prose to apply a reviewer's findings. You get the text, the task it was written for, and the findings. Return the rewritten text and nothing else: no preamble, no notes, no list of what you changed.

Apply every finding. When a finding says a passage should be cut, cut it. When a finding gives a rewrite, use that rewrite unless it conflicts with another finding. When a finding names a problem without a rewrite, fix the problem in the plainest way.

Keep every fact, number, and identifier (names of files, functions, flags, and symbols) exactly as the text has them. Add nothing the text and the task do not say. Keep the text's format: a Markdown document stays Markdown, a code comment stays a comment, a reply stays a reply.
"""

def applyFindings(
  text: string,
  task: string,
  guidelines: string,
  context: string,
  findings: Feedback[],
  model: string,
  provider: string,
): string {
  const findingLines = map(findings, \item -> "- ${item.feedback}").join("\n")
  let guidelinesSection = ""
  if (guidelines != "") {
    guidelinesSection = """The writing guidelines the text must follow:
<guidelines>
${guidelines}
</guidelines>

"""
  }
  const prompt = """
${rewritePrompt}

${guidelinesSection}The text was written for this task:
<task>
${withContext(task: task, context: context)}
</task>

The reviewer's findings:
<findings>
${findingLines}
</findings>

<text>
${text}
</text>
"""
  return llm(prompt, llmOptions(model: model, provider: provider))
}

/** Start with one pass and increase `passes` when another review of the
rewritten text is useful. On the eval suite, two passes scored 0.795
against 0.769 for one pass. The rewrite call receives the review findings
and the original task and guidelines. */
@guarded(cost: maxCost, time: maxTime)
export def writingRewriteAgent(
  text: string,
  task: string = "",
  guidelines: string = "",
  context: string = "",
  passes: number = 1,
  maxCost: number = $5.00,
  maxTime: number = 10m,
  model: string = "",
  provider: string = "",
): Result<string> {
  """
  Review prose and return a rewrite with the findings applied.
  Return unchanged text when the review has no findings.

  @param task - The intended reader and what the text must communicate.
  @param guidelines - Writing guidelines that override conflicting built-in review rules.
  @param passes - Maximum review-and-rewrite rounds, stopping early when a review has no findings.
  @param maxCost - Cost budget in USD for all rounds together.
  @param maxTime - Time budget in milliseconds for all rounds together.
  @param model - Model override, or "" for the ambient model.
  @param provider - Provider for the model override, ignored when model is empty.
  """
  let current = text
  let remaining = passes
  while (remaining > 0) {
    remaining = remaining - 1
    const reviewed = writingReviewAgent(
      text: current,
      task: task,
      guidelines: guidelines,
      context: context,
      maxCost: maxCost,
      maxTime: maxTime,
      model: model,
      provider: provider,
      failOpen: false,
    )
    // A review that could not run is not a clean pass: report it, rather
    // than returning the text as if nothing needed fixing.
    if (reviewed is failure(err)) {
      return failure("the writing reviewer failed: ${err}")
    }
    if (reviewed is success(findings)) {
      if (findings.length == 0) {
        remaining = 0
      } else {
        current = applyFindings(
          text: current,
          task: task,
          guidelines: guidelines,
          context: context,
          findings: findings,
          model: model,
          provider: provider,
        )
      }
    }
  }
  return current
}

// Takes the same input as the reviewer's `evalMain` in
// std::agents/writing/review. Effects are decided here as in the reviewer:
// the seeded input file is read `with approve`, and the agent runs under
// `docsOnlyHandler`, so a budget trip comes back as text rather than an
// interrupt escaping the entry point.
/** Run this module through an eval suite with `agency eval run stdlib/agents/writing/rewrite.agency:evalMain --suite <dir>`. */
node evalMain(input: WritingReviewEvalInput): string {
  setAgentName("writing-rewrite")
  evalValue(input)
  const source = read(input.sourceFile) with approve
  if (source is failure(readErr)) {
    return "could not read ${input.sourceFile}: ${readErr}"
  }
  let rewritten: Result<string> = failure("the rewriter never ran")
  handle {
    rewritten = writingRewriteAgent(text: source.value, task: input.assignment, passes: 1)
  } with docsOnlyHandler
  if (rewritten is failure(err)) {
    return "rewrite agent failed: ${err}"
  }
  return rewritten.value
}