local util = require "codeql.util"
local loader = require "codeql.loader"
local config = require "codeql.config"
local rpc = require "vim.lsp.rpc"
local protocol = require "vim.lsp.protocol"
local vim = vim

local client_index = 0
local progress_id = -1
local last_rpc_msg_id = -1
local lsp_client_id
local lsp_progress_handler
local progress_stage

local function next_client_id()
  client_index = client_index + 1
  return client_index
end

local function next_progress_id()
  progress_id = progress_id + 1
  return progress_id
end

local M = {}

M.client = nil

function M.start_client(opts)
  local client_id = next_client_id()

  local callbacks = opts.callbacks or {}
  local name = opts.name or tostring(client_id)
  local log_prefix = string.format("QueryServer[%s]", name)
  local handlers = {}

  local function resolve_callback(method)
    return callbacks[method] -- or default_callbacks[method]
  end

  function handlers.notification(method, params)
    local callback = resolve_callback(method)
    if callback then
      -- Method name is provided here for convenience.
      callback(method, params, client_id)
    end
  end

  function handlers.server_request(method, params)
    local callback = resolve_callback(method)
    if callback then
      return callback(method, params, client_id)
    end
    return nil, rpc.rpc_response_error(protocol.ErrorCodes.MethodNotFound)
  end

  function handlers.on_error(code, err)
    util.err_message(string.format("%s: Error %s: %s", log_prefix, rpc.client_errors[code], vim.inspect(err)))
    if opts.on_error then
      local status, usererr = pcall(opts.on_error, code, err)
      if not status then
        util.err_message(string.format("%s user on_error failed: ", log_prefix, tostring(usererr)))
      end
    end
  end

  function handlers.on_exit(code, signal)
    if opts.on_exit then
      pcall(opts.on_exit, code, signal)
    end
  end

  -- Start the RPC client.
  local client = rpc.start(opts.cmd, handlers, {
    env = opts.cmd_env,
  })
  client.id = client_id
  return client
end

function M.get_lsp_client_id()
  if lsp_client_id then
    return lsp_client_id
  else
    local lsp_clients = vim.lsp.get_active_clients()
    for _, lsp_client in ipairs(lsp_clients) do
      if lsp_client.name == "codeqlls" then
        lsp_client_id = lsp_client.id
        return lsp_client_id
      end
    end
  end
end

function M.get_lsp_handler()
  if lsp_progress_handler then
    return lsp_progress_handler
  end
  lsp_progress_handler = vim.lsp.handlers["$/progress"]
  return lsp_progress_handler
end

function M.start_server()
  if M.client then
    return M.client
  end

  util.message "Starting CodeQL Query Server"
  local cmd = {
    "codeql",
    "execute",
    "query-server2",
    "--debug",
    "--tuple-counting",
    "--threads=0",
    "--evaluator-log-level",
    "5",
    "-v",
    "--log-to-stderr",
    --"--additional-packs", "/Users/pwntester/src/github.com/github/securitylab-codeql",
  }
  local conf = config.values
  vim.list_extend(cmd, conf.ram_opts)

  local opts = {
    cmd = cmd,
    offset_encoding = { "utf-8", "utf-16" },
    callbacks = {

      -- progress update
      ["ql/progressUpdated"] = function(_, params, _)
        local client_id = M.get_lsp_client_id()
        local lsp_handler = M.get_lsp_handler()
        local message = params.message
        local progress = (100 * params.step) / params.maxStep
        if progress == 0 and not progress_stage then
          progress_stage = "begin"
        elseif progress == 0 then
          return
        elseif params.step >= params.maxStep - 1 then
          progress_stage = "end"
        end

        local handler
        if client_id and lsp_handler and type(lsp_handler) == "function" then
          -- use LSP progress handler
          handler = function(result)
            -- https://github.com/neovim/neovim/blob/b8ad1bfe8bc23ed5ffbfe43df5fda3501f1d2802/runtime/lua/vim/lsp/handlers.lua#L24
            local ctx = {
              client_id = M.get_lsp_client_id(),
            }
            --vim.lsp.handlers["$/progress"](nil, {value = { message = "foo"}, token = "foo"}, { client_id = vim.lsp.get_active_clients()[1].id })
            --print(vim.inspect(result), vim.inspect(ctx))
            lsp_handler(nil, result, ctx)
          end
        else
          -- Use vim.notify
          handler = function(result)
            if result.message then
              util.message("Query execution progress: " .. result.value.message)
            end
          end
        end

        if progress_stage == "begin" then
          handler {
            token = "CodeQLToken",
            value = {
              kind = progress_stage,
              title = "CodeQL",
              percentage = 0,
            },
          }
          progress_stage = "report"
        elseif progress_stage == "report" then
          handler {
            token = "CodeQLToken",
            value = {
              kind = progress_stage,
              message = message,
              percentage = progress,
            },
          }
        elseif progress_stage == "end" then
          handler {
            token = "CodeQLToken",
            value = {
              kind = progress_stage,
              percentage = 100,
            },
          }
          progress_stage = nil
        end
      end,
    },
  }
  return M.start_client(opts)
end

function M.run_query(opts)
  local dbPath = config.database.path
  if not dbPath then
    --util.err_message "Cannot find dataset folder. Did you `:QL db set`?"
    return
  end

  if not M.client then
    M.client = M.start_server()
  end

  local bufnr = opts.bufnr
  local queryPath = opts.query
  local bqrsPath = string.format(vim.fn.tempname(), ".bqrs")

  local additionalPacks = util.get_additional_packs()

  local runQuery_params = {
    body = {
      db = dbPath,
      additionalPacks = {
        additionalPacks or "",
      },
      externalInputs = {},
      singletonExternalInputs = opts.templateValues or {},
      outputPath = bqrsPath,
      queryPath = queryPath,
      -- do we want Datalog Intermediary Language dumps?
      -- https://codeql.github.com/docs/codeql-overview/codeql-glossary/#dil
      -- dilPath = "",
      -- logPath = "",
      target = opts.quick_eval and {
        quickEval = {
          quickEvalPos = {
            fileName = queryPath,
            line = opts.startLine,
            column = opts.startColumn,
            endLine = opts.endLine,
            endColumn = opts.endColumn,
          },
        },
      } or {
        query = { xx = "" },
      },
    },
    progressId = next_progress_id(),
  }

  local runQuery_callback = function(err, resp)
    if err then
      util.err_message("ERROR: " .. vim.inspect(err))
    end
    if not resp then
      -- we may have got an RPC error, so print it
      -- this is possible if the language server crashed, etc
      return
    end
    if resp["resultType"] == 0 then
      if util.is_file(bqrsPath) then
        util.bqrs_info({
          bqrs_path = bqrsPath,
          bufnr = bufnr,
          db_path = dbPath,
          query_path = queryPath,
          query_kind = opts.metadata["kind"],
          query_id = opts.metadata["id"],
          save_bqrs = true,
        }, loader.process_results)
      end
    elseif resp["resultType"] == 1 then
      util.err_message("ERROR: Other: " .. resp["message"])
    elseif resp["resultType"] == 2 then
      util.err_message("ERROR: Compilation Error: " .. resp["message"])
    elseif resp["resultType"] == 3 then
      util.err_message("ERROR: OOM Error: " .. resp["message"])
    elseif resp["resultType"] == 4 then
      util.err_message("ERROR: Query cancelled: " .. resp["message"])
    elseif resp["resultType"] == 5 then
      util.err_message("ERROR: DB Scheme mismatch: " .. resp["message"])
    elseif resp["resultType"] == 6 then
      util.err_message("ERROR: DB Scheme no upgrade found: " .. resp["message"])
    else
      util.err_message "Query run failed. Database may be locked by a different Query Server"
    end
  end

  -- run query
  --util.message(string.format("Running query %s", queryPath))

  _, last_rpc_msg_id = M.client.request("evaluation/runQuery", runQuery_params, runQuery_callback)
end

function M.register_database(database)
  if not M.client then
    M.client = M.start_server()
  end
  config.database = database
  local lang = config.database.languages[1]
  -- https://github.com/github/vscode-codeql/blob/0e2c03f572226a85ef4691621ca07424c06a2162/extensions/ql-vscode/src/common/query-language.ts#L48-L57
  local langTodbScheme = {
    javascript = "semmlecode.javascript.dbscheme",
    cpp = "semmlecode.cpp.dbscheme",
    java = "semmlecode.dbscheme",
    python = "semmlecode.python.dbscheme",
    csharp = "semmlecode.csharp.dbscheme",
    go = "go.dbscheme",
    ruby = "ruby.dbscheme",
    ql = "ql.dbscheme",
    swift = "swift.dbscheme",
    yaml = "yaml.dbscheme",
  }
  local dbschemePath = config.database.datasetFolder .. "/" .. langTodbScheme[lang]
  if not vim.fn.filereadable(dbschemePath) then
    util.err_message "Cannot find dbscheme file"
    return
  else
    util.message(string.format("Using dbscheme file %s", dbschemePath))
  end
  local resp = util.database_upgrades(dbschemePath)
  if resp ~= vim.NIL and #resp.scripts > 0 then
    util.database_upgrade(config.database.path)
  end
  util.message(string.format("Registering database %s", config.database.path))
  local params = {
    body = {
      databases = {
        config.database.path,
      },
      progressId = next_progress_id(),
    },
  }
  M.client.request("evaluation/registerDatabases", params, function(err, result)
    if err then
      util.err_message(string.format("Error registering database %s", vim.inspect(err)))
    else
      util.message(string.format("Successfully registered %s", result.registeredDatabases[1]))
      -- TODO: add option to open the drawer automatically
      --require 'codeql.explorer'.draw()
    end
  end)
end

function M.unregister_database(cb)
  if util.is_blank(config.database) then
    vim.notify "No database registered"
    return
  end
  if not M.client then
    M.client = M.start_server()
  end
  util.message(string.format("Deregistering database %s", config.database.path))
  local params = {
    body = {
      databases = {
        config.database.path,
      },
      progressId = next_progress_id(),
    },
  }
  M.client.request("evaluation/deregisterDatabases", params, function(err, result)
    if err then
      util.err_message(string.format("Error registering database %s", vim.inspect(err)))
    elseif #result.registeredDatabases == 0 then
      util.message(string.format("Successfully deregistered %s", config.database.path))
      config.database = nil
    end
    -- call the callback
    if cb then
      cb()
    end
  end)
end

function M.cancel_query()
  if last_rpc_msg_id < 0 then
    return
  end
  if not M.client then
    M.client = M.start_server()
  end
  util.message(M.client.notify("$/cancelRequest", {
    id = last_rpc_msg_id,
  }))
end

function M.stop_server()
  if M.client then
    local handle = M.client.handle
    handle:kill()
    M.client = nil
  end
end

return M
