import { hostedSearchTools, searchTools } from "std::agents/lib/search"
import { AGENT_SPILL_OVER, llmOptions, withContext } from "std::agents/lib/shared"
import { webTools, spillTools, SAVE_DRAFT_HINT, ASK_USER_HINT } from "std::agents/lib/toolkits"
import { guarded } from "std::decorators"
import { agentSkill } from "std::skills"
import { evalValue } from "std::statelog"
import {
  ThreadMessage,
  assistantMessage,
  ensureSystemMessage,
  threadIsNew,
  userMessage,
 } from "std::thread"

import { communicationTools } from "./lib/toolkits.agency"

/** @module
@summary Answer research questions using web and Wikipedia sources.

Use `researcherAgent` for a cited answer from retrieved sources. It asks
for an answer again when the first reply is too short and reports a failure
if that follow-up is blank. No separate reviewer checks the answer.

As an LLM tool, this handoff continues the caller's conversation. From code,
use a `thread` block to isolate the research. `keepDraft: false` makes a
budget stop a failure even when a draft is available. The draft is then in
the failure's data as `draft`.
*/

static const systemPrompt = """You are a research analyst. Answer the task using your source tools.

- Ground every factual claim in a source and cite it. A citation is the full URL of the page, written out (https://...), so the reader can open it. Do not cite a tool call or a search result number.
- If your tools cannot retrieve the needed information, say so plainly. NEVER fabricate sources or facts.
- If no source describes the thing the task asks about (a feature, a product, an event), say that no source describes it. Do not describe it from memory.
- Do your own looking up: never ask the user for something a source could tell you. Answer with what you found, and say what you could not find.
- If a page will not load, try another source; do not stop on one failure.

## Communicating with the user
- Tell the user what you are doing with the `whatIAmDoing` tool when you start and when you change approach. Each call costs a tool round; spend the rounds on sources.${ASK_USER_HINT}${SAVE_DRAFT_HINT}
"""

static const skills = agentSkill("researcher")

export def buildTools(): any[] {
  """
  Return the researcher's source lookup and communication tools.
  """
  return [
    ...webTools(),
    ...spillTools(),
    read.partial(useAgentCwd: true),
    ...searchTools(),
    skills,
    ...communicationTools(),
  ]
}

// Retry short first replies, but accept any nonblank follow-up.
static const MIN_ANSWER_CHARS = 20

static const NO_ANSWER_FOLLOWUP = "You returned no answer. Write your answer now from what you have read, citing sources."

def answerIsBlank(answer: any): boolean {
  if (answer == null) {
    return true
  }
  const text: string = "${answer}"
  return text.trim().length == 0
}

def answerIsEmpty(answer: any): boolean {
  if (answer == null) {
    return true
  }
  const text: string = "${answer}"
  return text.trim().length < MIN_ANSWER_CHARS
}

def research(
  task: string,
  context: string,
  history: ThreadMessage[],
  model: string,
  provider: string,
  extraTools: any[],
): Result<string> {
  const researchTools = buildTools()
  const fresh = threadIsNew()
  // A handoff runs on the caller's thread, which may already hold the
  // persona; the push is skipped when it does.
  ensureSystemMessage(systemPrompt)
  if (fresh) {
    // The conversation so far, so the question is read in its context. A
    // shared session already holds its history and is not seeded again.
    for (message in history) {
      if (message.role == "assistant") {
        assistantMessage(message.content)
      } else {
        userMessage(message.content)
      }
    }
  }
  const question = withContext(task: task, context: context)
  let answer: any = llm(
    question,
    llmOptions(
      model: model,
      provider: provider,
      tools: [...researchTools, ...extraTools, saveDraft],
      spillToolResultsOver: AGENT_SPILL_OVER,
      hostedTools: hostedSearchTools(model, provider),
      // A research run is many searches and fetches; the default ten
      // rounds end it before the sources are in.
      maxToolCallRounds: 30
    ),
  )
  if (answerIsEmpty(answer)) {
    // Keep the sources in context and remove tools so the follow-up writes an answer.
    answer = llm(NO_ANSWER_FOLLOWUP, llmOptions(model: model, provider: provider))
    if (answerIsBlank(answer)) {
      return failure("the model returned no answer twice")
    }
  }
  saveDraft(answer)
  return success(answer)
}

@guarded(cost: maxCost, time: maxTime, keepDraft: keepDraft)
export handoff def researcherAgent(
  task: string,
  context: string = "",
  maxCost: number = $50.00,
  maxTime: number = 30m,
  model: string = "",
  provider: string = "",
  extraTools: any[] = [],
  history: ThreadMessage[] = [],
  keepDraft: boolean = true,
): Result<string> {
  """
  Answer a research question from web and Wikipedia sources and return a
  cited answer.

  @param task - The research question
  @param context - Extra material folded into the prompt, or ""
  @param history - Prior messages, oldest first, used only when the active thread is empty.
  @param maxCost - Cost budget in USD.
  @param maxTime - Time budget in milliseconds.
  @param model - Model override, or "" for the ambient model
  @param provider - Provider for the model override, ignored when model is empty.
  @param extraTools - Extra tools to offer the LLM, appended to the built-in set
  @param keepDraft - False returns a budget failure even when a draft was saved.
  """

  // Return position: when the budget runs out, the draft the research
  // pass saved reaches the guard and comes back as the result.
  return research(
    task: task,
    context: context,
    history: history,
    model: model,
    provider: provider,
    extraTools: extraTools,
  )
}

// Approve research fetches and searches for evals, and reject other interrupts.
def researchOnlyHandler(intr) {
  return match(intr.effect) {
    "std::http::fetch" => approve()
    "std::http::fetchJSON" => approve()
    "std::http::fetchMarkdown" => approve()
    "std::search" => approve()
    "std::tavilySearch" => approve()
    "std::wikipedia::search" => approve()
    "std::wikipedia::summary" => approve()
    "std::wikipedia::article" => approve()
    _ => reject()
  }
}

/** What an eval hands the researcher: the question, and the conversation
  before it when the test has one. The shape the `evals/researcher` suite
  uses as its input. */
export type ResearchEvalInput = {
  task: string;
  history?: ThreadMessage[]
}

/** Run this module through an eval suite with `agency eval run stdlib/agents/researcher.agency:evalMain --suite <dir>`. */
node evalMain(input: ResearchEvalInput): string {
  evalValue(input)
  let history: ThreadMessage[] = []
  if (input.history != null) {
    history = input.history
  }
  let answer: Result<string> = failure("the researcher never ran")
  handle {
    answer = researcherAgent(
      task: input.task,
      history: history,
      maxCost: $3.00,
      maxTime: 6m,
    )
  } with researchOnlyHandler
  if (answer is failure(err)) {
    return "research agent failed: ${err}"
  }
  return answer.value
}