local function read_file(path)
  local file = io.open(path, "r")
  if file == nil then
    error("Unable to open file: " .. path)
  end

  local contents = file:read("*a")
  file:close()
  return contents
end

local function read_file_if_exists(path)
  local file = io.open(path, "r")
  if file == nil then
    return nil
  end

  local contents = file:read("*a")
  file:close()
  return contents
end

local function escape_json(value)
  value = value:gsub("\\", "\\\\")
  value = value:gsub('"', '\\"')
  value = value:gsub("\n", "\\n")
  value = value:gsub("\r", "\\r")
  value = value:gsub("\t", "\\t")
  return value
end

local cjson_safe = nil
local cjson = nil

local function build_graphql_request_body(query)
  if cjson_safe ~= nil and cjson_safe.encode ~= nil then
    local encoded = cjson_safe.encode({ query = query })
    if encoded ~= nil then
      return encoded
    end
  end

  if cjson ~= nil and cjson.encode ~= nil then
    local ok_encode, encoded = pcall(cjson.encode, { query = query })
    if ok_encode and encoded ~= nil then
      return encoded
    end
  end

  local escaped_query = escape_json(query)
  return "{\"query\":\"" .. escaped_query .. "\"}"
end

local function build_persisted_document_request_body(document_id)
  if cjson_safe ~= nil and cjson_safe.encode ~= nil then
    local encoded = cjson_safe.encode({ documentId = document_id })
    if encoded ~= nil then
      return encoded
    end
  end

  if cjson ~= nil and cjson.encode ~= nil then
    local ok_encode, encoded = pcall(cjson.encode, { documentId = document_id })
    if ok_encode and encoded ~= nil then
      return encoded
    end
  end

  local escaped_document_id = escape_json(document_id)
  return "{\"documentId\":\"" .. escaped_document_id .. "\"}"
end

local function hash_string(value)
  local hash = 5381
  local max_u32 = 4294967296

  for i = 1, #value do
    hash = ((hash * 33) + value:byte(i)) % max_u32
  end

  return string.format("%08x", hash)
end
local ok_safe, cjson_safe_module = pcall(require, "cjson.safe")
if ok_safe then
  cjson_safe = cjson_safe_module
end

local ok_cjson, cjson_module = pcall(require, "cjson")
if ok_cjson then
  cjson = cjson_module
end

-- wrk runs response() in each worker's own Lua environment, while done()
-- runs in a separate environment. Keep counters global so done() can read
-- them from each worker via thread:get().
threads = {}
status_failures = 0
graphql_error_responses = 0
response_structure_failures = 0
checked_response_structure = false
sample_status_failure = nil
sample_graphql_error = nil
sample_structure_error = nil
local find = string.find

local expected_response_file = os.getenv("BENCH_EXPECTED_RESPONSE_FILE")
local expected_response

if expected_response_file ~= nil and expected_response_file ~= "" then
  expected_response = read_file(expected_response_file)
else
  expected_response = read_file_if_exists("expected_response.json") or read_file("bench/expected_response.json")
end

local expected_response_hash = hash_string(expected_response)

local function check_response_structure(body)
  if body == nil then
    return false
  end

  local response_hash = hash_string(body)
  return response_hash == expected_response_hash
end

local operation_file = os.getenv("BENCH_OPERATION_FILE")
local persisted_mode = os.getenv("BENCH_PERSISTED_MODE") == "true"
local request_body

if persisted_mode then
  local document_id = os.getenv("BENCH_DOCUMENT_ID")
  if document_id == nil or document_id == "" then
    document_id = "bench_test_query"
  end
  request_body = build_persisted_document_request_body(document_id)
else
  local query
  if operation_file ~= nil and operation_file ~= "" then
    query = read_file(operation_file)
  else
    query = read_file_if_exists("operation.graphql") or read_file("bench/operation.graphql")
  end
  request_body = build_graphql_request_body(query)
end

wrk.method = "POST"
wrk.headers["Content-Type"] = "application/json"
wrk.body = request_body

setup = function(thread)
  -- Store thread handles in the done() environment so validation results can
  -- be aggregated after all workers finish.
  threads[#threads + 1] = thread
end

response = function(status, headers, body)
  if status ~= 200 then
    status_failures = status_failures + 1
    if sample_status_failure == nil then
      sample_status_failure = body
    end
    return
  end

  if body and find(body, '"errors"', 1, true) then
    graphql_error_responses = graphql_error_responses + 1
    if sample_graphql_error == nil then
      sample_graphql_error = body
    end
    return
  end

  if checked_response_structure then
    return
  end

  checked_response_structure = true

  if not check_response_structure(body) then
    response_structure_failures = response_structure_failures + 1
    if sample_structure_error == nil then
      sample_structure_error = body
    end
  end
end

local function get_thread_number(thread, name)
  return thread:get(name) or 0
end

local function get_thread_sample(thread, name)
  local value = thread:get(name)
  if value ~= nil and value ~= "" then
    return value
  end

  return nil
end

done = function(summary, latency, requests)
  local total_status_failures = 0
  local total_graphql_error_responses = 0
  local total_response_structure_failures = 0
  local status_failure_sample = nil
  local graphql_error_sample = nil
  local structure_error_sample = nil

  -- done() cannot see mutations performed by response() directly. Pull the
  -- final global counter values from every worker environment instead.
  for _, thread in ipairs(threads) do
    total_status_failures = total_status_failures + get_thread_number(thread, "status_failures")
    total_graphql_error_responses = total_graphql_error_responses + get_thread_number(thread, "graphql_error_responses")
    total_response_structure_failures = total_response_structure_failures + get_thread_number(thread, "response_structure_failures")

    if status_failure_sample == nil then
      status_failure_sample = get_thread_sample(thread, "sample_status_failure")
    end
    if graphql_error_sample == nil then
      graphql_error_sample = get_thread_sample(thread, "sample_graphql_error")
    end
    if structure_error_sample == nil then
      structure_error_sample = get_thread_sample(thread, "sample_structure_error")
    end
  end

  io.write("VALIDATION_STATUS_FAILURES=" .. total_status_failures .. "\n")
  io.write("VALIDATION_GRAPHQL_ERRORS=" .. total_graphql_error_responses .. "\n")
  io.write("VALIDATION_RESPONSE_STRUCTURE_FAILURES=" .. total_response_structure_failures .. "\n")

  if status_failure_sample ~= nil then
    io.write("VALIDATION_STATUS_FAILURE_SAMPLE=" .. status_failure_sample .. "\n")
  end
  if graphql_error_sample ~= nil then
    io.write("VALIDATION_GRAPHQL_ERROR_SAMPLE=" .. graphql_error_sample .. "\n")
  end
  if structure_error_sample ~= nil then
    io.write("VALIDATION_RESPONSE_STRUCTURE_FAILURE_SAMPLE=" .. structure_error_sample .. "\n")
  end
end
