local util = require "codeql.util"
local panel = require "codeql.panel"
local codeql_ast_ns = vim.api.nvim_create_namespace "codeql_ast"
local vim = vim

local M = {}

M._entries = setmetatable({}, {
  __index = function(tbl, key)
    local entry = rawget(tbl, key)

    if not entry then
      entry = {}
      rawset(tbl, key, entry)
    end

    return entry
  end,
})

local function is_valid_graph(props)
  for _, prop in ipairs(props.tuples) do
    if prop[1] == "semmle.graphKind" and prop[2] == "tree" then
      return true
    end
  end
  return false
end

local function is_buf_visible(bufnr)
  local windows = vim.fn.win_findbuf(bufnr)
  return #windows > 0
end

local function close_buf_windows(bufnr)
  if not bufnr then
    return
  end

  util.for_each_buf_window(bufnr, function(window)
    vim.api.nvim_win_close(window, true)
  end)
end

local function close_buf(bufnr)
  if not bufnr then
    return
  end

  close_buf_windows(bufnr)

  if vim.api.nvim_buf_is_loaded(bufnr) then
    vim.cmd(string.format("bw! %d", bufnr))
  end
end

local function clear_entry(bufnr)
  local entry = M._entries[bufnr]

  close_buf(entry.display_bufnr)
  close_buf(entry.query_bufnr)
  M._entries[bufnr] = nil
end

local function setup_buf(for_buf)
  if M._entries[for_buf].display_bufnr then
    return M._entries[for_buf].display_bufnr
  end

  local buf = vim.api.nvim_create_buf(false, false)

  vim.api.nvim_buf_set_option(buf, "buftype", "nofile")
  vim.api.nvim_buf_set_option(buf, "swapfile", false)
  vim.api.nvim_buf_set_option(buf, "buflisted", false)
  vim.api.nvim_buf_set_option(buf, "filetype", "codeql_ast")

  vim.cmd(string.format("augroup CodeQLAST_%d", buf))
  vim.cmd "au!"
  vim.cmd(string.format([[autocmd CursorMoved <buffer=%d> lua require'codeql.ast'.highlight_node(%d)]], buf, for_buf))
  vim.cmd(string.format([[autocmd BufLeave <buffer=%d> lua require'codeql.ast'.clear_highlights(%d)]], buf, for_buf))
  vim.cmd(string.format([[autocmd BufWinEnter <buffer=%d> lua require'codeql.ast'.update(%d)]], buf, for_buf))
  vim.cmd "augroup END"

  vim.api.nvim_buf_set_keymap(
    buf,
    "n",
    "R",
    string.format(':lua require "codeql.ast".update(%d)<CR>', for_buf),
    { silent = true }
  )
  vim.api.nvim_buf_set_keymap(
    buf,
    "n",
    "<CR>",
    string.format(':lua require "codeql.ast".goto_node(%d)<CR>', for_buf),
    { silent = true }
  )
  vim.api.nvim_buf_attach(buf, false, {
    on_detach = function()
      clear_entry(for_buf)
    end,
  })

  return buf
end

local function print_node(bufnr, node, lines, level)
  -- keep track of what nodes are printed in each line
  M._entries[bufnr].tree2buf[#lines + 1] = node
  node.line = #lines

  -- print node
  local line = string.format(
    "%s%s [%d, %d] - [%d, %d]",
    string.rep(" ", 2 * level),
    node.label ~= "" and node.label or "???",
    node.location.startLine,
    node.location.startColumn,
    node.location.endLine,
    node.location.endColumn
  )
  line = string.gsub(line, "\n", "")
  table.insert(lines, line)

  -- print each child
  if #node.children > 0 then
    for _, child in ipairs(node.children) do
      print_node(bufnr, child, lines, level + 1)
    end
  end
end

local function print_tree(bufnr, roots)
  M._entries[bufnr].tree2buf = {}
  local lines = {}
  for _, root in ipairs(roots) do
    print_node(bufnr, root, lines, 0)
  end
  return lines
end

local function sort_node(node)
  table.sort(node.children, function(left, right)
    return left.order < right.order
  end)
  if #node.children > 0 then
    for _, child in ipairs(node.children) do
      sort_node(child)
    end
  end
end

local function sort_tree(roots)
  table.sort(roots, function(left, right)
    return left.order < right.order
  end)
  for _, root in ipairs(roots) do
    sort_node(root)
  end
end

local function is_cursor_in_node(cursor, node)
  local start_row = node.location.startLine
  local start_col = node.location.startColumn
  local end_row = node.location.endLine
  local end_col = node.location.endColumn
  local cursor_row = cursor[1]
  local cursor_col = cursor[2] + 1

  -- print(format('(%d,%d) in [%d,%d] - [%d-%d]?',
  --   cursor_row, cursor_col,
  --   start_row, start_col,
  --   end_row, end_col))

  if cursor_row < start_row or cursor_row > end_row then
    return false
  elseif cursor_row > start_row and cursor_row < end_row then
    return true
  elseif start_row == end_row then
    return cursor_col >= start_col and cursor_col <= end_col
  elseif end_row > start_row then
    if cursor_row == start_row then
      return cursor_col >= start_col
    elseif cursor_row == end_row then
      return cursor_col <= end_col
    end
  end
  util.err_message "Error checking cursor position"
end

local function get_node_at_cursor(cursor, node, matching_nodes)
  if is_cursor_in_node(cursor, node) then
    table.insert(matching_nodes, node)
  end
  if #node.children > 0 then
    for _, child in ipairs(node.children) do
      get_node_at_cursor(cursor, child, matching_nodes)
    end
  end
end

-- exported functions

function M.clear_highlights(bufnr, namespace)
  if not bufnr then
    return
  end
  namespace = namespace or codeql_ast_ns
  vim.api.nvim_buf_clear_namespace(bufnr, namespace, 0, -1)
end

function M.clear_ast_highlights(bufnr)
  M.clear_highlights(M._entries[bufnr].display_bufnr)
end

function M.highlight_ast_nodes(bufnr, nodes)
  local entry = M._entries[bufnr]
  local results = entry.results
  local display_buf = entry.display_bufnr
  local lines = {}

  if not results or not display_buf then
    return
  end

  for _, node in ipairs(nodes) do
    local lnum = node.line
    table.insert(lines, lnum)
    local ast_lines = vim.api.nvim_buf_get_lines(display_buf, lnum, lnum + 1, false)
    if ast_lines[1] then
      local _, _, ws, _ = string.find(ast_lines[1], "(%s*)(.+)")
      vim.api.nvim_buf_add_highlight(display_buf, codeql_ast_ns, "CodeqlAstFocus", lnum, #ws, -1)
    end
  end

  return lines
end

function M.highlight_node(bufnr)
  M.clear_highlights(bufnr)

  local node = M._entries[bufnr].tree2buf[vim.fn.line "."]

  if not node then
    return
  end

  local loc = node.location
  local start_row = loc.startLine
  local start_col = loc.startColumn
  local end_row = loc.endLine
  local end_col = loc.endColumn
  vim.highlight.range(
    bufnr,
    codeql_ast_ns,
    "CodeqlAstFocus",
    { start_row - 1, start_col - 1 },
    { end_row - 1, end_col }
  )

  util.for_each_buf_window(bufnr, function(window)
    pcall(vim.api.nvim_win_set_cursor, window, { start_row, start_col })
  end)
end

function M.goto_node(bufnr)
  local line = vim.fn.line "."
  local node = M._entries[bufnr].tree2buf[line]
  local loc = node.location

  local bufwin = vim.fn.win_findbuf(bufnr)[1]
  if bufwin then
    vim.api.nvim_set_current_win(bufwin)
    M.clear_highlights(bufnr)
    vim.api.nvim_win_set_cursor(bufwin, { loc.startLine, loc.startColumn })
  end
end

function M.open(bufnr)
  bufnr = bufnr or vim.api.nvim_get_current_buf()
  local display_buf = setup_buf(bufnr)
  local current_window = vim.api.nvim_get_current_win()

  M._entries[bufnr].display_bufnr = display_buf
  vim.cmd "vsplit"
  vim.cmd(string.format("buffer %d", display_buf))

  vim.api.nvim_win_set_option(0, "spell", false)
  vim.api.nvim_win_set_option(0, "number", false)
  vim.api.nvim_win_set_option(0, "relativenumber", false)
  vim.api.nvim_win_set_option(0, "cursorline", false)

  vim.api.nvim_set_current_win(current_window)

  return display_buf
end

function M.update(bufnr)
  bufnr = bufnr or vim.api.nvim_get_current_buf()
  local display_buf = M._entries[bufnr].display_bufnr

  -- Don't bother updating if the playground isn't shown
  if not display_buf or not is_buf_visible(display_buf) then
    return
  end

  local results = M._entries[bufnr].results

  vim.api.nvim_buf_set_lines(display_buf, 0, -1, false, results.lines)
  -- if print_virt_hl then
  --   printer.print_hl_groups(bufnr, display_buf)
  -- end
end

function M.toggle(bufnr)
  bufnr = bufnr or vim.api.nvim_get_current_buf()
  local display_buf = M._entries[bufnr].display_bufnr

  if display_buf and is_buf_visible(display_buf) then
    close_buf_windows(M._entries[bufnr].query_bufnr)
    close_buf_windows(display_buf)
  else
    M.open(bufnr)
  end
end

function M.highlight_ast_node_from_buffer(bufnr)
  M.clear_ast_highlights(bufnr)

  local display_buf = M._entries[bufnr].display_bufnr
  if not display_buf then
    return
  end

  local bufwin = vim.fn.win_findbuf(bufnr)[1]
  if not bufwin then
    bufwin = vim.api.nvim_get_current_win()
  end
  if bufwin then
    local cursor = vim.api.nvim_win_get_cursor(bufwin)
    local roots = M._entries[bufnr].roots

    local nodes_at_cursor = {}
    for _, root in ipairs(roots) do
      get_node_at_cursor(cursor, root, nodes_at_cursor)

      -- if we already found a root with nodes at cursor, skip the rest
      if #nodes_at_cursor > 0 then
        break
      end
    end

    if #nodes_at_cursor == 0 then
      return
    end

    local lnums = M.highlight_ast_nodes(bufnr, nodes_at_cursor)
    if lnums[#lnums] then
      util.for_each_buf_window(display_buf, function(window)
        pcall(vim.api.nvim_win_set_cursor, window, { lnums[#lnums], 0 })
      end)
    end
  end
end

M._highlight_ast_node_debounced = util.debounce(M.highlight_ast_node_from_buffer, function()
  return 25
end)

function M.attach(bufnr)
  vim.cmd(string.format("augroup CodeQLAST %d", bufnr))
  vim.cmd "au!"
  vim.cmd(
    string.format(
      [[autocmd CursorMoved <buffer=%d> lua require'codeql.ast'._highlight_ast_node_debounced(%d)]],
      bufnr,
      bufnr
    )
  )
  vim.cmd(
    string.format([[autocmd BufLeave <buffer=%d> lua require'codeql.ast'.clear_ast_highlights(%d)]], bufnr, bufnr)
  )
  vim.cmd "augroup END"
end

function M.detach(bufnr)
  clear_entry(bufnr)
  vim.cmd(string.format("autocmd! CodeQLAST_%d CursorMoved", bufnr))
  vim.cmd(string.format("autocmd! CodeQLAST_%d BufLeave", bufnr))
end

function M.build_ast(jsonPath, bufnr)
  if not util.is_file(jsonPath) then
    return
  end
  local results = util.read_json_file(jsonPath)

  local nodeTuples = results.nodes
  local edgeTuples = results.edges
  local graphProperties = results.graphProperties

  if not is_valid_graph(graphProperties) then
    util.err_message "Invalid AST tree"
    return
  end

  vim.api.nvim_buf_clear_namespace(bufnr, panel.ns, 0, -1)

  local idToItem = {}
  local parentToChildren = {}
  local childToParent = {}
  local astOrder = {}
  local roots = {}
  local edgeLabels = {}

  -- Build up the parent-child relationships
  for _, tuple in ipairs(edgeTuples.tuples) do
    local source, target, tupleType, value = unpack(tuple)
    -- local target = tuple[2]
    -- local tupleType = tuple[3]
    -- local orderValue = tuple[4]
    local sourceId = source.id
    local targetId = target.id

    if tupleType == "semmle.order" then
      astOrder[targetId] = tonumber(value)
    elseif tupleType == "semmle.label" then
      childToParent[targetId] = sourceId
      local children = parentToChildren[sourceId] or {}
      table.insert(children, targetId)
      parentToChildren[sourceId] = children

      -- ignore values that indicate a numeric order.
      if not tonumber(value) then
        edgeLabels[targetId] = value
      end
    else
      -- ignore other tupleTypes since they are not needed by the ast viewer
    end
  end

  -- populate parents and children
  for _, tuple in ipairs(nodeTuples.tuples) do
    local entity, tupleType, value = unpack(tuple)
    -- local tupleType = tuple[2]
    -- local orderValue = tuple[3]
    local id = entity.id

    if tupleType == "semmle.order" then
      astOrder[id] = tonumber(value)
    elseif tupleType == "semmle.label" then
      -- If an edge label exists, include it and separate from the node label using ':'
      local nodeLabel = value and value or entity.label
      local edgeLabel = edgeLabels[id]
      local label
      if edgeLabel then
        label = string.format("%s: %s", edgeLabel, nodeLabel)
      else
        label = nodeLabel
      end

      local item = {
        id = id,
        label = label,
        location = entity.url,
        children = {},
        order = math.huge,
      }

      idToItem[id] = item
      local parent = idToItem[childToParent[id] or -1]
      if parent then
        local astItem = item
        astItem.parent = parent
        table.insert(parent.children, astItem)
      end

      local children = parentToChildren[id] or {}
      for _, childId in ipairs(children) do
        local child = idToItem[childId]
        if child then
          child.parent = item
          table.insert(item.children, child)
        end
      end
    else
      -- ignore other tupleTypes since they are not needed by the ast viewer
    end
  end

  -- find the roots and add the order
  for _, item in pairs(idToItem) do
    item.order = astOrder[item.id] and astOrder[item.id] or math.huge
    if not item.parent then
      table.insert(roots, item)
    end
  end

  sort_tree(roots)

  --util.message("found " .. #roots .. " roots")

  M.attach(bufnr)
  local text_lines = print_tree(bufnr, roots)
  M._entries[bufnr].roots = roots
  M._entries[bufnr].results = { lines = text_lines }
  M.open(bufnr)
end

return M
