import Foundation

enum ToolCallParserError: LocalizedError, Sendable, Equatable {
  case multipleCalls
  case malformedCall
  case contentOutsideCall
  case invalidPayload
  case unexpectedFields([String])
  case missingName
  case missingArguments

  var errorDescription: String? {
    switch self {
    case .multipleCalls:
      "The model requested more than one tool in a single turn."
    case .malformedCall:
      "The model returned an incomplete tool-call envelope."
    case .contentOutsideCall:
      "The model mixed a tool call with other response content."
    case .invalidPayload:
      "The tool call was not valid JSON."
    case .unexpectedFields(let fields):
      "The tool call included unexpected fields: \(fields.joined(separator: ", "))."
    case .missingName:
      "The tool call did not include a valid function name."
    case .missingArguments:
      "The tool call arguments must be a JSON object."
    }
  }
}

struct ToolCallParser: Sendable {
  private static let openingTag = "<tool_call>"
  private static let closingTag = "</tool_call>"

  static func parse(_ text: String) throws -> AgentToolCall? {
    let openingRanges = text.ranges(of: openingTag)
    let closingRanges = text.ranges(of: closingTag)
    if openingRanges.isEmpty && closingRanges.isEmpty {
      return nil
    }
    guard openingRanges.count == 1, closingRanges.count == 1 else {
      throw ToolCallParserError.multipleCalls
    }
    guard
      let opening = openingRanges.first,
      let closing = closingRanges.first,
      opening.upperBound <= closing.lowerBound
    else {
      throw ToolCallParserError.malformedCall
    }

    let prefix = text[..<opening.lowerBound]
    let suffix = text[closing.upperBound...]
    guard
      prefix.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty,
      suffix.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
    else {
      throw ToolCallParserError.contentOutsideCall
    }

    let payloadText = String(text[opening.upperBound..<closing.lowerBound])
      .trimmingCharacters(in: .whitespacesAndNewlines)
    guard let data = payloadText.data(using: .utf8) else {
      throw ToolCallParserError.invalidPayload
    }
    let value: JSONValue
    do {
      value = try JSONDecoder().decode(JSONValue.self, from: data)
    } catch {
      throw ToolCallParserError.invalidPayload
    }
    guard case .object(let payload) = value else {
      throw ToolCallParserError.invalidPayload
    }
    let unexpected = Set(payload.keys)
      .subtracting(["name", "arguments"])
      .sorted()
    guard unexpected.isEmpty else {
      throw ToolCallParserError.unexpectedFields(unexpected)
    }
    guard
      let rawName = payload["name"]?.stringValue?
        .trimmingCharacters(in: .whitespacesAndNewlines),
      !rawName.isEmpty
    else {
      throw ToolCallParserError.missingName
    }
    guard case .object(let arguments) = payload["arguments"] else {
      throw ToolCallParserError.missingArguments
    }
    return AgentToolCall(name: rawName, arguments: arguments)
  }
}

private extension String {
  func ranges(of substring: String) -> [Range<String.Index>] {
    var result: [Range<String.Index>] = []
    var searchRange = startIndex..<endIndex
    while let range = range(of: substring, range: searchRange) {
      result.append(range)
      searchRange = range.upperBound..<endIndex
    }
    return result
  }
}