import { map } from "std::array"
import {
  Policy,
  InterruptEffect,
  InterruptDataKey,
  InterruptDataVal,
 } from "std::policy"
import { highlight } from "std::syntax"
import { _writePolicyFile } from "agency-lang/stdlib-lib/policy.js"
import { args } from "std::system"
import { systemMessage, userMessage, assistantMessage } from "std::thread"

callback("onToolCallStart") as data {
  print("Calling ${data.toolName}...")
}

callback("onToolCallEnd") as data {
  print("Finished calling ${data.toolName}.")
}

systemPrompt = read("system.md", __dirname + "/prompts") with approve

type Question = {
  type: "askQuestion";
  question: string
}

type Done = {
  type: "done"
}

type Action = Question | Done

type PolicyAction = "approve" | "reject" | "propagate"

let currentPolicy: Policy = {}

def addToPolicy(
  interruptEffect: InterruptEffect,
  matchArgName: InterruptDataKey,
  matchArgValue: InterruptDataVal,
  action: PolicyAction!,
): Result<Policy> {
  """
  Add a rule to the current policy for the given interrupt effect, matching on the specified argument and value, with the specified action.
  Use only "approve", "reject", or "propagate" as the action. DO NOT use "allow" or "deny".
  Returns the current policy after adding the rule.
  """
  const match = {
    [matchArgName]: matchArgValue
  }
  if (!currentPolicy[interruptEffect]) {
    currentPolicy[interruptEffect] = []
  }
  currentPolicy[interruptEffect].push({
    match: match,
    action: action
  })

  return success(currentPolicy)
}

def getPolicy(): Policy {
  """
  Get the current policy object.
  """
  return currentPolicy
}

def setPolicy(newPolicy: Policy): Policy {
  """
  Replace the current policy with a new policy object. Returns the new policy.

  @param newPolicy - The new policy object to set as the current policy.
  """
  currentPolicy = newPolicy
  return currentPolicy
}

def deleteFromPolicy(interruptEffect: InterruptEffect, index: number): boolean {
  """
  Delete the rule at the specified index for the given interrupt effect from the current policy.
  Returns true if the rule was successfully deleted, or false if the index was out of bounds.
  """
  const rules = currentPolicy[interruptEffect]
  if (rules != null && index < rules.length) {
    rules.splice(index, 1)
    return true
  }
  return false
}

def savePolicy(path: string) {
  """
  Validate and save the current policy to a JSON file.

  @param path - The file path where the policy should be saved.
  """
  return try _writePolicyFile(path, currentPolicy, [])
}

def get(array: any[], index: number) {
  if (index < array.length) {
    return success(array[index])
  }
  return failure("Index out of bounds: ${index} >= ${array.length}")
}

def highlightJSON(obj: Record<string, any>): string {
  return highlight(JSON.stringify(obj, null, 2), "json")
}

node main() {
  const cliArgs = args()
  if (cliArgs.length < 2) {
    print("Usage: agency policy gen <file>")
    return
  }

  const interruptEffects = try JSON.parse(cliArgs[0])
  const outputPath = cliArgs[1]
  const existingPolicy = try JSON.parse(get(cliArgs, 2))

  if (interruptEffects is failure(error)) {
    print("Failed to parse interrupt effects: ${error}")
    return
  }

  let contextMessage = "This agent can produce the following interrupts:\n${map(interruptEffects.value, \effect -> "- ${effect}\n").join("")}"

  if (existingPolicy is success(thePolicy)) {
    currentPolicy = thePolicy
    contextMessage = """
${contextMessage}

Current policy:
```json
${highlightJSON(currentPolicy)}
```

What would you like to change?
    """
  } else {
    contextMessage = """
${contextMessage}

What actions would you like to allow?
    """
  }

  print(contextMessage)

  let userInput = input("> ")
  // didn't catch `policy` var is undefined??
  // doesn't typecheck in string segments??
  const prompt = """
  Current policy: ${JSON.stringify(currentPolicy, null, 2)}

  Decide the next step. If the user is describing what they want, build or update the policy and use addToPolicy, setPolicy, or deleteFromPolicy. Show the user the policy and get confirmation that it looks good. When you are done, return  { type: "done" }. Otherwise, ask the user a question to clarify what they want, and return { type: "askQuestion", question: "..." }.
  """

  systemMessage(systemPrompt catch "")
  assistantMessage(contextMessage)
  userMessage(userInput)

  const tools = [addToPolicy, getPolicy, setPolicy, deleteFromPolicy]

  let action: Action = llm(prompt, {
    tools: tools
  })
  while (action.type != "done") {
    print(action.question)
    let userInput = input("> ")
    action: Action = llm(userInput, {
      tools: tools
    })
  }
  print("Done. Final policy:")
  printJSON(currentPolicy)
  const result = savePolicy(outputPath)
  if (result is failure(error)) {
    print("Failed to save policy: ${error}")
    return
  }
  print("Saved policy to ${outputPath}.")
}