local FILE = vim.env.NVIM_BENCH_FILE
  or '/home/lillis/projects/grammars/tree-sitter-zig/foo.zig'
local PARSER = vim.env.NVIM_BENCH_PARSER
  or '/home/lillis/projects/grammars/tree-sitter-zig/zig.so'
local QUERY_FILE = vim.env.NVIM_BENCH_QUERY
  or '/home/lillis/projects/grammars/tree-sitter-zig/queries/highlights.scm'
local REPS = tonumber(vim.env.NVIM_BENCH_REPS) or 3

vim.treesitter.language.add('zig', { path = PARSER })

local fd = assert(io.open(FILE, 'r'))
local source = fd:read('*a')
fd:close()
local fd2 = assert(io.open(QUERY_FILE, 'r'))
local query_text = fd2:read('*a')
fd2:close()

local lines = vim.split(source, '\n', { plain = true })
local buf = vim.api.nvim_create_buf(false, true)
vim.api.nvim_buf_set_lines(buf, 0, -1, false, lines)
vim.bo[buf].filetype = 'zig'

local function ns_to_ms(ns) return ns / 1e6 end

local parser = vim.treesitter.get_parser(buf, 'zig')
parser:invalidate(true)
local t0 = vim.uv.hrtime()
parser:parse(true)
local parse_ms = ns_to_ms(vim.uv.hrtime() - t0)
local tree = parser:trees()[1]

local query = vim.treesitter.query.parse('zig', query_text)

local function stats(times)
  table.sort(times)
  local sum = 0
  for _, v in ipairs(times) do sum = sum + v end
  return {
    mean = sum / #times,
    median = times[math.floor(#times / 2) + 1],
    min = times[1],
    max = times[#times],
  }
end

local function bench(label, fn)
  local n = fn()  -- warm
  local times = {}
  for i = 1, REPS do
    local t = vim.uv.hrtime()
    fn()
    times[i] = ns_to_ms(vim.uv.hrtime() - t)
  end
  local s = stats(times)
  io.stdout:write(string.format(
    '%-22s n=%-7d mean=%9.2f ms  median=%9.2f ms  min=%9.2f ms  max=%9.2f ms\n',
    label, n, s.mean, s.median, s.min, s.max))
end

io.stdout:write(string.format('file=%s lines=%d size=%d parse_ms=%.3f reps=%d\n',
  FILE, #lines, #source, parse_ms, REPS))

bench('zig iter_captures', function()
  local n = 0
  for _ in query:iter_captures(tree:root(), buf, 0, -1) do n = n + 1 end
  return n
end)

bench('zig iter_matches', function()
  local n = 0
  for _ in query:iter_matches(tree:root(), buf, 0, -1) do n = n + 1 end
  return n
end)