import Foundation

enum DolphinPromptMessage: Sendable, Equatable {
  case user(String)
  case assistant(String)
  case assistantToolCall(AgentToolCall)
  case toolResponse(String)
}

struct DolphinPromptRenderer: Sendable {
  static func render(
    systemPrompt: String,
    messages: [DolphinPromptMessage],
    tools: [AgentToolDefinition],
    addGenerationPrompt: Bool = true
  ) -> String {
    var prompt = renderSystemPrompt(
      sanitize(systemPrompt),
      tools: tools
    )

    for message in messages {
      switch message {
      case .user(let content):
        prompt += chatMessage(role: "user", content: sanitize(content))
      case .assistant(let content):
        prompt += chatMessage(role: "assistant", content: sanitize(content))
      case .assistantToolCall(let call):
        prompt += "<|im_start|>assistant\n"
        prompt += "<tool_call>\n"
        prompt += sanitize(toolCallJSON(call))
        prompt += "\n</tool_call><|im_end|>\n"
      case .toolResponse(let content):
        prompt += "<|im_start|>user\n"
        prompt += "<tool_response>\n"
        prompt += sanitize(content)
        prompt += "\n</tool_response><|im_end|>\n"
      }
    }

    if addGenerationPrompt {
      prompt += "<|im_start|>assistant\n"
    }
    return prompt
  }

  static func toolCallJSON(_ call: AgentToolCall) -> String {
    JSONValue.object([
      "arguments": .object(call.arguments),
      "name": .string(call.name),
    ]).canonicalJSON
  }

  private static func renderSystemPrompt(
    _ systemPrompt: String,
    tools: [AgentToolDefinition]
  ) -> String {
    let governedPrompt = systemPrompt + """


      # Runtime protocol

      - Use no tool for greetings, introductions, thanks, or ordinary conversation.
      - Use memory_save only for an explicit request to remember or store information.
      - Use memory_search only when the user asks to recall or search saved information.
      - After a tool result, answer naturally and concisely.
      - Keep internal reasoning private. Return only the next tool call or the final answer; never emit analysis, chain-of-thought, or <think> blocks.
      - Never reproduce or imitate internal receipts, call IDs, argument JSON, observations, or Activity logs in the answer.
      - If no tools are listed, answer without a tool call.
      """
    guard !tools.isEmpty else {
      return chatMessage(role: "system", content: governedPrompt)
    }

    var value = "<|im_start|>system\n\(governedPrompt)"
    value += "\n\n# Tools\n\n"
    value += "You may call one function to assist with the user query.\n\n"
    value += "You are provided with function signatures within <tools></tools> XML tags:\n<tools>"
    for tool in tools.sorted(by: { $0.id < $1.id }) {
      value += "\n\(toolDefinitionJSON(tool))"
    }
    value += "\n</tools>\n\n"
    value += "For a function call, return one JSON object with function name and arguments within <tool_call></tool_call> XML tags:\n"
    value += "<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n"
    return value
  }

  private static func toolDefinitionJSON(_ tool: AgentToolDefinition) -> String {
    JSONValue.object([
      "function": .object([
        "description": .string(tool.summary),
        "name": .string(tool.id),
        "parameters": tool.parameters,
      ]),
      "type": .string("function"),
    ]).canonicalJSON
  }

  private static func chatMessage(role: String, content: String) -> String {
    "<|im_start|>\(role)\n\(content)<|im_end|>\n"
  }

  private static func sanitize(_ value: String) -> String {
    var result = value
    while let start = result.range(of: "<|"),
      let end = result.range(
        of: "|>",
        range: start.upperBound..<result.endIndex
      )
    {
      result.replaceSubrange(start.lowerBound..<end.upperBound, with: "[special_token]")
    }
    return result
      .replacingOccurrences(of: "<tool_call>", with: "[tool_call]")
      .replacingOccurrences(of: "</tool_call>", with: "[/tool_call]")
      .replacingOccurrences(of: "<tool_response>", with: "[tool_response]")
      .replacingOccurrences(of: "</tool_response>", with: "[/tool_response]")
      .replacingOccurrences(of: "<tools>", with: "[tools]")
      .replacingOccurrences(of: "</tools>", with: "[/tools]")
  }
}