local util = require "codeql.util"
local queryserver = require "codeql.queryserver"
local config = require "codeql.config"
local Path = require "plenary.path"
local vim = vim

local M = {}

function M.list_databases()
  local pok, pickers = pcall(require, "telescope.pickers")
  local fok, finders = pcall(require, "telescope.finders")
  local cok, conf = pcall(require, "telescope.config")
  local aok, actions = pcall(require, "telescope.actions")
  local sok, action_state = pcall(require, "telescope.actions.state")
  local uok, utils = pcall(require, "telescope.utils")
  if not (pok and fok and cok and aok and sok) then
    util.err_message "Telescope is not installed"
    return
  end
  local cmd = config.values.find_databases_cmd
  if not cmd then
    util.err_message "No find_databases_cmd set in config"
    return
  end
  if not type(cmd) == "table" then
    util.err_message "find_databases_cmd should be a list of strings"
    return
  end
  local opts = config.values.telescope_opts
  local entry_maker = config.values.database_list_entry_maker
  local previewer = config.values.database_list_previewer

  local results = utils.get_os_command_output(cmd)
  pickers
    .new(opts, {
      prompt_title = "CodeQL Databases",
      finder = finders.new_table {
        results = results,
        entry_maker = entry_maker,
      },
      previewer = previewer(opts),
      sorter = conf.values.generic_sorter(opts),
      attach_mappings = function(prompt_bufnr)
        actions.select_default:replace(function()
          actions.close(prompt_bufnr)
          local selection = action_state.get_selected_entry()
          M.set_database(selection.value)
        end)
        return true
      end,
    })
    :find()
end

function M.set_database(dbpath)
  local conf = config.values
  conf.ram_opts = util.resolve_ram()
  dbpath = vim.fn.fnamemodify(vim.trim(dbpath), ":p")
  local database
  if not dbpath then
    util.err_message("Incorrect database: " .. dbpath)
  elseif Path:new(dbpath):is_file() and vim.endswith(dbpath, ".zip") then
    -- extract the zip file
    local db_dir = string.format("%s/codeql_dbs/%s", vim.fn.stdpath "data", vim.fn.fnamemodify(dbpath, ":t:r"))
    -- make sure db_dir exists
    vim.fn.mkdir(vim.fn.fnamemodify(db_dir, ":h"), "p", 0777)
    -- extract the zip file
    vim.fn.system(string.format("unzip -q %s -d %s", dbpath, db_dir))
    local db_name = vim.trim(vim.fn.system(string.format('unzip -Z1 %s | head -n 1 | cut -d "/" -f 1', dbpath)))
    database = string.format("%s/%s/", db_dir, db_name)
  elseif util.is_dir(dbpath) then
    database = dbpath
    if not vim.endswith(database, "/") then
      database = database .. "/"
    end
  else
    util.err_message("Incorrect database: " .. dbpath)
  end
  if database then
    util.message("Database set to " .. database)
    local metadata = util.database_info(database)
    if not metadata then
      util.err_message("Could not load the database " .. vim.inspect(database))
      return
    end
    metadata.path = database

    if not util.is_blank(config.database) then
      queryserver.unregister_database(function()
        queryserver.register_database(metadata)
      end)
    else
      queryserver.register_database(metadata)
    end
  end
end

function M.unset_database(dbpath)
  if not util.is_blank(config.database) then
    queryserver.unregister_database(function() end)
  end
end

local function is_predicate_node(node)
  return node:type() == "charpred" or node:type() == "memberPredicate" or node:type() == "classlessPredicate"
end

local function is_predicate_identifier_node(parent, node)
  if parent then
    return (parent:type() == "charpred" and node:type() == "className")
      or (parent:type() == "classlessPredicate" and node:type() == "predicateName")
      or (parent:type() == "memberPredicate" and node:type() == "predicateName")
  end
end

---@return TSNode?, TSNode?
function M.get_node_at_cursor()
  local winnr = vim.api.nvim_get_current_win()
  local bufnr = vim.api.nvim_get_current_buf()
  local cursor = vim.api.nvim_win_get_cursor(winnr)
  local cursor_range = { cursor[1] - 1, cursor[2] }
  local ok, root_lang_tree = pcall(vim.treesitter.get_parser, bufnr, vim.treesitter.language.get_lang(vim.bo.filetype))
  if not ok then
    return
  end
  local root = root_lang_tree:trees()[1]:root()
  return root:named_descendant_for_range(cursor_range[1], cursor_range[2], cursor_range[1], cursor_range[2]), root
end

function M.get_eval_position()
  local bufnr = vim.api.nvim_get_current_buf()
  if vim.treesitter.highlighter.active[bufnr] == nil then
    util.err_message "No treesitter parser"
    return
  end
  local node, root = M.get_node_at_cursor()
  local orig_node = node
  if not node then
    util.err_message "Error getting node at cursor. Make sure treesitter CodeQL parser is installed"
    return
  end
  local parent = node:parent()
  -- ascend the AST until we get to a predicate node
  while parent and parent ~= root and not is_predicate_node(node) and not is_predicate_identifier_node(parent, node) do
    node = parent
    parent = node:parent()
  end
  if parent == root then
    -- We got to the root node, evaluate the whole query
    return
  elseif is_predicate_identifier_node(parent, node) then
    local srow, scol, erow, ecol = node:range()
    local midname = math.floor((scol + ecol) / 2)
    return { srow + 1, midname, erow + 1, midname }
  elseif is_predicate_node(node) then
    -- descend the predicate node till we find the name node
    for child in node:iter_children() do
      if is_predicate_identifier_node(node, child) then
        --util.message(string.format("Evaluating '%s' predicate", vim.treesitter.get_node_text(child, bufnr)[1]))
        local srow, scol, erow, ecol = child:range()
        local midname = math.floor((scol + ecol) / 2)
        return { srow + 1, midname, erow + 1, midname }
      end
    end
    vim.notify("No predicate identifier node found", 2)
    return
  else
    vim.notify("No predicate node found " .. vim.inspect(orig_node:type()), 2)
  end
end

function M.smart_quick_evaluate()
  local position = M.get_eval_position()
  if position then
    M.query(true, position)
  else
    M.query(false, util.get_current_position())
  end
end

function M.quick_evaluate()
  M.query(true, util.get_current_position())
end

function M.run_query()
  M.query(false, util.get_current_position())
end

function M.query(quick_eval, position)
  local db = config.database
  if not db or not db.path then
    util.err_message "Missing database. Use `:QL db set` command"
    return
  end

  local dbPath = db.path
  local queryPath = vim.fn.expand "%:p"

  local libPaths = util.resolve_library_path(queryPath)
  if not libPaths then
    vim.notify(string.format("Cannot resolve QL library paths for %s", queryPath), 2)
    return
  end

  local opts = {
    quick_eval = quick_eval,
    bufnr = vim.api.nvim_get_current_buf(),
    query = queryPath,
    dbPath = dbPath,
    startLine = position[1],
    startColumn = position[2],
    endLine = position[3],
    endColumn = position[4],
    metadata = util.query_info(queryPath),
    libraryPath = libPaths.libraryPath,
    dbschemePath = libPaths.dbscheme,
  }

  queryserver.run_query(opts)
end

local templated_queries = {
  c = { qlpack = "codeql/cpp-all" },
  cpp = { qlpack = "codeql/cpp-all" },
  java = { qlpack = "codeql/java-all" },
  cs = { qlpack = "codeql/csharp-all" },
  javascript = { qlpack = "codeql/javascript-all" },
  python = { qlpack = "codeql/python-all" },
  ql = { qlpack = "codeql/ql", path = "ide-contextual-queries/" },
  ruby = { qlpack = "codeql/ruby-all", path = "ide-contextual-queries/" },
  go = { qlpack = "codeql/go-all" },
  yaml = { qlpack = "codeql/actions-all", path = "ql/lib/ide-contextual-queries/" },
  yml = { qlpack = "codeql/actions-all", path = "ql/lib/ide-contextual-queries/" },
  swift = { qlpack = "codeql/swift-all", path = "ide-contextual-queries/" },
}

function M.run_print_cfg()
  local bufnr = vim.api.nvim_get_current_buf()
  local bufname = vim.fn.bufname(bufnr)

  -- not a ql:/ buffer
  if not vim.startswith(bufname, "ql:/") then
    util.err_message "Not a ql:/ buffer"
    return
  end

  local fname = vim.split(bufname, "://")[2]
  M.run_templated_query("printCfg", fname)
end

function M.run_print_ast()
  local bufnr = vim.api.nvim_get_current_buf()
  local bufname = vim.fn.bufname(bufnr)

  -- not a ql:/ buffer
  if not vim.startswith(bufname, "ql:/") then
    util.err_message "Not a ql:/ buffer"
    return
  end

  local fname = vim.split(bufname, "://")[2]
  M.run_templated_query("printAst", fname)
end

function M.run_templated_query(query_name, fname)
  local bufnr = vim.api.nvim_get_current_buf()
  local dbPath = config.database.path
  local ft = vim.fn.fnamemodify(fname, ":e")
  if not templated_queries[ft] then
    util.err_message(string.format("%s does not support %s file type", query_name, ft))
    return
  end
  local qlpack = templated_queries[ft].qlpack
  local path_modifier = templated_queries[ft].path or ""
  local qlpacks = util.resolve_qlpacks()
  local path
  if qlpacks and qlpacks[qlpack] then
    path = qlpacks[qlpack][1]
  else
    path = vim.fn.getcwd()
  end
  local position = vim.api.nvim_win_get_cursor(0)
  local queryPath = string.format("%s/%s%s.ql", path, path_modifier, query_name)
  if util.is_file(queryPath) then
    local libPaths = util.resolve_library_path(queryPath)
    if not libPaths then
      vim.notify("Cannot resolve QL library paths for: " .. query_name, 2)
      return
    end
    local opts = {
      quick_eval = false,
      bufnr = bufnr,
      query = queryPath,
      dbPath = dbPath,
      metadata = util.query_info(queryPath),
      libraryPath = libPaths.libraryPath,
      dbschemePath = libPaths.dbscheme,
      templateValues = {
        selectedSourceFile = "/" .. fname,
        selectedSourceLine = tostring(position[1]),
        selectedSourceColumn = tostring(position[2] + 1),
      },
    }
    require("codeql.queryserver").run_query(opts)
  end
end

function M.get_permalink()
  local bufnr = vim.api.nvim_get_current_buf()
  local bufname = vim.api.nvim_buf_get_name(bufnr)
  local line = vim.api.nvim_win_get_cursor(0)[1]
  local uri = string.match(bufname, "ql://(.*)")
  if not uri then
    util.err_message "Cannot copy permalink for this buffer"
    return
  end
  local chunks = vim.split(uri, "/")
  if #chunks < 4 then
    util.err_message "Cannot copy permalink for this buffer"
    return
  end
  local owner = chunks[1]
  local name = chunks[2]
  local revisionId = chunks[3]
  local path = table.concat(chunks, "/", 4, #chunks)
  return string.format("https://github.com/%s/%s/blob/%s/%s#L%d", owner, name, revisionId, path, line)
end

function M.copy_permalink()
  local permalink = M.get_permalink()
  vim.fn.setreg("+", permalink)
end

function M.open_in_browser()
  local permalink = M.get_permalink()
  -- replace github.com with github.dev
  permalink = string.gsub(permalink, "github.com", "github.dev")
  vim.fn.system(string.format("open %s", permalink))
end

function M.deprecated(cmd, replacement)
  if not replacement then
    replacement = ":QL"
  end
  vim.notify("'" .. cmd .. "'command is deprecated. Please use `" .. replacement .. "` instead.")
end

function M.setup(opts)
  if vim.fn.executable "codeql" then
    config.setup(opts or {})

    -- highlight groups
    vim.cmd [[highlight default link CodeqlAstFocus CursorLine]]
    vim.cmd [[highlight default link CodeqlRange Error]]

    -- deprecated commands
    vim.cmd [[command! -nargs=1 -complete=file SetDatabase lua require'codeql'.deprecated("SetDatabase");require'codeql'.set_database(<f-args>)]]
    vim.cmd [[command! UnsetDatabase lua require'codeql'.deprecated("UnsetDatabase");require'codeql.queryserver'.unregister_database()]]
    vim.cmd [[command! CancelQuery lua require'codeql'.deprecated("CancelQuery");require'codeql.queryserver'.cancel_query()]]
    vim.cmd [[command! RunQuery lua require'codeql'.deprecated("RunQuery");require'codeql'.run_query()]]
    vim.cmd [[command! QuickEvalPredicate require'codeql'.deprecated("QuickEvalPredicate");lua require'codeql'.smart_quick_evaluate()]]
    vim.cmd [[command! -range QuickEval lua require'codeql'.deprecated("QuickEval");require'codeql'.quick_evaluate()]]
    vim.cmd [[command! StopServer lua require'codeql'.deprecated("StopServer");require'codeql.queryserver'.stop_server()]]
    vim.cmd [[command! History lua require'codeql'.deprecated("History");require'codeql.history'.menu()]]
    vim.cmd [[command! PrintAST lua require'codeql'.deprecated("PrintAST");require'codeql'.run_print_ast()]]
    vim.cmd [[command! -nargs=1 -complete=file LoadSarif lua require'codeql'.deprecated("LoadSarif");require'codeql.loader'.load_sarif_results(<f-args>)]]
    vim.cmd [[command! ArchiveTree lua require'codeql'.deprecated("ArchiveTree");require'codeql.explorer'.draw()]]
    vim.cmd [[command! -nargs=1 LoadMRVAScan lua require'codeql'.deprecated("LoadMRVAScan");require'codeql.mrva.panel'.load(<f-args>)]]
    vim.cmd [[command! CopyPermalink lua require'codeql'.deprecated("CopyPermalink");require'codeql'.copy_permalink()]]
    -- new QL command
    vim.api.nvim_create_user_command("QL", function(copts)
      require("codeql").command(unpack(copts.fargs))
    end, { complete = require("codeql").command_complete, nargs = "*" })
    -- autocommands
    vim.cmd [[augroup codeql]]
    vim.cmd [[au!]]
    vim.cmd [[au BufEnter * if &ft ==# 'codeql_panel' | execute("lua require'codeql.panel'.apply_mappings()") | endif]]
    vim.cmd [[autocmd FileType ql lua require'codeql.util'.apply_mappings()]]

    if require("codeql.config").values.format_on_save then
      vim.cmd [[autocmd FileType ql autocmd BufWrite <buffer> lua vim.lsp.buf.format()]]
    end
    vim.cmd [[augroup END]]

    -- mappings
    vim.cmd [[nnoremap <Plug>(CodeQLGoToDefinition) <cmd>lua require'codeql.defs'.find_at_cursor('definitions')<CR>]]
    vim.cmd [[nnoremap <Plug>(CodeQLFindReferences) <cmd>lua require'codeql.defs'.find_at_cursor('references')<CR>]]
  end
end

local commands = {
  database = {
    set = {
      description = "Set the database to use for CodeQL queries",
      args = {
        {
          name = "database",
          description = "The path to the database to use",
          type = "string",
        },
      },
      handler = function(args)
        M.set_database(args)
      end,
    },
    unset = {
      description = "Unset the database to use for CodeQL queries",
      handler = function()
        M.unset_database()
      end,
    },
    browse = {
      description = "Browse the CodeQL database",
      handler = function()
        require("codeql.explorer").draw()
      end,
    },
    grep = {
      description = "Grep the CodeQL database",
      handler = function()
        require("codeql.grepper").grep_source()
      end,
    },
    list = {
      description = "List the CodeQL databases known to CodeQL.nvim",
      handler = function()
        M.list_databases()
      end,
    },
  },
  query = {
    run = {
      description = "Run the current query",
      handler = function()
        M.run_query()
      end,
    },
    cancel = {
      description = "Cancel the current query",
      handler = function()
        require("codeql.queryserver").cancel_query()
      end,
    },
    eval = {
      description = "Evaluate the current predicate or statements",
      handler = function()
        M.smart_quick_evaluate()
      end,
    },
  },
  server = {
    stop = {
      description = "Stop the CodeQL query server",
      handler = function()
        require("codeql.queryserver").stop_server()
      end,
    },
  },
  history = {
    list = {
      description = "View the CodeQL query history",
      handler = function()
        require("codeql.history").menu()
      end,
    },
  },
  view = {
    ast = {
      description = "Print the AST for the current query",
      handler = function()
        M.run_print_ast()
      end,
    },
    cfg = {
      description = "Print the CFG for the current scope",
      handler = function()
        M.run_print_cfg()
      end,
    },
  },
  ast = {
    print = {
      description = "Print the AST for the current query",
      handler = function()
        M.deprecated("QL ast print", "QL view ast")
      end,
    },
  },
  sarif = {
    load = {
      description = "Load a SARIF file",
      args = {
        {
          name = "sarif",
          description = "The path to the SARIF file to load",
          type = "string",
        },
      },
      handler = function(args)
        require("codeql.loader").load_sarif_results(args)
      end,
    },
    ["permalink"] = {
      description = "Copy a permalink to the current SARIF result",
      handler = function()
        M.copy_permalink()
      end,
    },
    ["browse"] = {
      description = "Copy a permalink to the current SARIF result",
      handler = function()
        M.open_in_browser()
      end,
    },
  },
  mrva = {
    list = {
      description = "List the MRVA scans",
      handler = function()
        require("codeql.mrva.panel").list_sessions()
      end,
    },
    load = {
      description = "Load an MRVA scan",
      args = {
        {
          name = "scan",
          description = "The path to the MRVA scan to load",
          type = "string",
        },
      },
      handler = function(args)
        require("codeql.mrva.panel").load(args)
      end,
    },
  },
}
-- define an alias for database
commands["db"] = commands["database"]

function M.command_complete(argLead, cmdLine)
  -- ArgLead		the leading portion of the argument currently being completed on
  -- CmdLine		the entire command line
  -- CursorPos	the cursor position in it (byte index)
  local command_keys = vim.tbl_keys(commands)
  local parts = vim.split(vim.trim(cmdLine), " ")

  local get_options = function(options)
    local valid_options = {}
    for _, option in pairs(options) do
      if string.sub(option, 1, #argLead) == argLead then
        table.insert(valid_options, option)
      end
    end
    return valid_options
  end

  if #parts == 1 then
    return command_keys
  elseif #parts == 2 and not vim.tbl_contains(command_keys, parts[2]) then
    return get_options(command_keys)
  elseif #parts == 2 and vim.tbl_contains(command_keys, parts[2]) or #parts == 3 then
    local obj = commands[parts[2]]
    if obj then
      return get_options(vim.tbl_keys(obj))
    end
  end
end

function M.command(object, action, ...)
  if not object or not action then
    util.err_message "Missing arguments"
    return
  end
  local command = commands[object] and commands[object][action]
  if command then
    command.handler(...)
  else
    util.err_message "Unknown command"
  end
end

return M
