--- The otter.keeep takes care of all the otters
--- attached to the main buffers.
--- Each main buffer is associated with a raft of
--- otter buffers and their respective languages
--- and code chunks
local keeper = {}

local fn = require("otter.tools.functions")
local api = vim.api
local ts = vim.treesitter

---@class Raft
---@field languages string[]
---@field buffers table<string, integer>
---@field paths table<string, string>
---@field preambles table<string, string[]>
---@field postambles table<string, string[]>
---@field ignore_pattern table<string, string>
---@field otter_nr_to_lang table<integer, string>
---@field tsquery string?
---@field query vim.treesitter.Query
---@field parser vim.treesitter.LanguageTree
---@field otterls OtterLSInfo
---@field diagnostics_namespaces integer[]
---@field diagnostics_group integer?
---@field last_changetick integer?
---@field code_chunks table<string, CodeChunk[]>

---@class OtterLSInfo
---@field client_id integer?

---@class Rafts
---@field [number] Raft
---One raft per main buffer
---stored in the rafts table
keeper.rafts = {}

--- table of languages that can be injected
--- generated from the lanaguages
--- for which we have extensions
local injectable_languages = {}
for key, _ in pairs(OtterConfig.extensions) do
  table.insert(injectable_languages, key)
end

local function filter_lang(s)
  return s:gsub("^[%s%p]+", "")
end

---determine the language of the current node
---@param main_nr integer bufnr of the main buffer
---@param name string name of the capture
---@param node table node of the current capture
---@param metadata table metadata of the current capture
---@param current_language string? current language
---@return string?
local function determine_language(main_nr, name, node, metadata, current_language)
  local injection_language = metadata["injection.language"]
  if injection_language ~= nil then
    -- chunks using the metadata to specify the injected language
    -- e.g. html script tags
    if injection_language ~= "comment" then
      -- don't use comment as language,
      -- comments with language inside are handled in injection.combined
      return filter_lang(injection_language)
    end
  elseif metadata["injection.combined"] == true then
    -- chunks where the injected language is specified in the text of a comment
    local lang_capture = metadata[2]["text"]
    if lang_capture ~= nil then
      -- NOTE: this could be more elegant
      -- remove leading whitespace and comment characters erroneously captured as part of the langauge
      return filter_lang(lang_capture)
    end
  elseif name == "_lang" or name == "injection.language" then
    -- chunks where the name of the injected language is dynamic
    -- e.g. markdown code chunks
    return ts.get_node_text(node, main_nr, metadata)
  else
    return current_language
  end
end

---trims the leading whitespace from text
---@param text string
---@param bufnr integer host buffer number
---@param starting_ln integer
---@return string, integer
local function trim_leading_witespace(text, bufnr, starting_ln)
  if not OtterConfig.handle_leading_whitespace then
    return text, 0
  end

  -- Assume the first line is least indented
  -- the first line in the capture doesn't have its leading indent, so we grab from the buffer
  local split = vim.split(text, "\n", { trimempty = false })
  if #split == 0 then
    return text, 0
  end
  local first_line = vim.api.nvim_buf_get_lines(bufnr, starting_ln, starting_ln + 1, false)
  local leading = first_line[1]:match("^%s+")
  if not leading then
    return text, 0
  end
  for i, line in ipairs(split) do
    split[i] = line:gsub("^" .. leading, "")
  end
  return table.concat(split, "\n"), #leading
end

---@class CodeChunk
---@field range { from: [integer, integer], to: [integer, integer] }
---@field lang string
---@field node TSNode
---@field text string[]
---@field leading_offset number

---Extract code chunks from the specified buffer.
---Updates M.rafts[main_nr].code_chunks
---@param main_nr integer The main buffer number
---@param lang string? language to extract. All languages if nil.
---@param exclude_eval_false boolean? Exclude code chunks with eval: false
---@param range_start_row integer? Row to start from, inclusive, 1-indexed.
---@param range_end_row integer? Row to end at, inclusive, 1-indexed.
---@return table<string, CodeChunk[]>
keeper.extract_code_chunks = function(main_nr, lang, exclude_eval_false, range_start_row, range_end_row)
  local query = keeper.rafts[main_nr].query
  local parser = keeper.rafts[main_nr].parser
  local tree = parser:parse()
  assert(tree, "[otter] Treesitter failed to parse buffer " .. main_nr)
  local root = tree[1]:root()

  ---@type table<string, CodeChunk[]>
  local code_chunks = {}
  local lang_capture = nil
  for _, match, metadata in query:iter_matches(root, main_nr, 0, -1, { all = true }) do
    for id, nodes in pairs(match) do
      local name = query.captures[id]

      for _, node in ipairs(nodes) do
        local text
        lang_capture = determine_language(main_nr, name, node, metadata, lang_capture)
        if
          lang_capture
          and (name == "content" or name == "injection.content")
          and (lang == nil or lang_capture == lang)
        then
          -- the actual code content
          text = ts.get_node_text(node, main_nr, { metadata = metadata[id] })
          -- remove surrounding quotes
          -- (workaround for treesitter offsets not properly processed)
          text, _ = fn.strip_wrapping_quotes(text)
          if exclude_eval_false and string.find(text, "| *eval: *false") then
            text = ""
          end

          -- Use the range from the `offset!` directive if available
          local start_row, start_col, end_row, end_col
          if metadata[id] and metadata[id].range then
            start_row, start_col, end_row, end_col = unpack(metadata[id].range)
          else
            start_row, start_col, end_row, end_col = node:range()
          end

          if
            range_start_row ~= nil
            and range_end_row ~= nil
            and ((start_row >= range_end_row and range_end_row > 0) or end_row < range_start_row)
          then
            goto continue
          end
          local leading_offset
          text, leading_offset = trim_leading_witespace(text, main_nr, start_row)
          local result = {
            range = { from = { start_row, start_col }, to = { end_row, end_col } },
            lang = lang_capture,
            node = node,
            text = fn.lines(text),
            leading_offset = leading_offset,
          }
          if code_chunks[lang_capture] == nil then
            code_chunks[lang_capture] = {}
          end
          table.insert(code_chunks[lang_capture], result)
          -- reset current language
          lang_capture = nil
        elseif fn.contains(injectable_languages, name) then
          -- chunks where the name of the language is the name of the capture
          if lang == nil or name == lang then
            text = ts.get_node_text(node, main_nr, { metadata = metadata[id] })
            text, _ = fn.strip_wrapping_quotes(text)

            ---@type integer
            ---@diagnostic disable-next-line: assign-type-mismatch
            local start_row, start_col, end_row, end_col = node:range()
            local leading_offset
            text, leading_offset = trim_leading_witespace(text, main_nr, start_row)
            local result = {
              range = { from = { start_row, start_col }, to = { end_row, end_col } },
              lang = name,
              node = node,
              text = fn.lines(text),
              leading_offset = leading_offset,
            }
            if code_chunks[name] == nil then
              code_chunks[name] = {}
            end
            table.insert(code_chunks[name], result)
          end
        end
        ::continue::
      end
    end
  end
  return code_chunks
end

--- Get the language context of a position
--- @param main_nr integer? bufnr of the parent buffer. Default is 0
--- @param position table? position (row, col). Default is the current cursor position (1,0)-based
--- @return string? language nil if no language context is found
--- @return integer? start_row
--- @return integer? start_col
--- @return integer? end_row
--- @return integer? end_col
keeper.get_current_language_context = function(main_nr, position)
  main_nr = main_nr or api.nvim_get_current_buf()
  position = position or api.nvim_win_get_cursor(0)
  if keeper.rafts[main_nr] == nil then
    return nil
  end
  local row, col = unpack(position)
  row = row - 1
  col = col

  local query = keeper.rafts[main_nr].query
  local parser = keeper.rafts[main_nr].parser
  local trees = parser:parse()
  assert(trees, "[otter] Treesitter failed to parse buffer " .. main_nr)
  local tree = trees[1]
  local root = tree:root()
  local lang_capture = nil
  for _, match, metadata in query:iter_matches(root, main_nr, 0, -1, { all = true }) do
    for id, nodes in pairs(match) do
      local name = query.captures[id]

      for _, node in ipairs(nodes) do
        lang_capture = determine_language(main_nr, name, node, metadata, lang_capture)
        local start_row, start_col, end_row, end_col = node:range()
        end_row = end_row - 1

        local language = nil
        if lang_capture and (name == "content" or name == "injection.content") then
          -- chunks where the name of the injected language is dynamic
          -- e.g. markdown code chunks
          if ts.is_in_node_range(node, row, col) then
            language = lang_capture
          end
          -- chunks where the name of the language is the name of the capture
        elseif fn.contains(injectable_languages, name) then
          if ts.is_in_node_range(node, row, col) then
            language = name
          end
        end

        if language then
          if OtterConfig.handle_leading_whitespace then
            local buf = keeper.rafts[main_nr].buffers[language]
            if buf then
              local lines = vim.api.nvim_buf_get_lines(buf, end_row - 1, end_row, false)
              if lines[1] then
                end_col = #lines[1]
              end
            end
          end
          return language, start_row, start_col, end_row, end_col
        end
      end
    end
  end
  return nil
end

---find the leading_offset of the given line number, and buffer number. Returns 0 if the line number
---isn't in a chunk.
---@param line_nr number
---@param main_nr number
keeper.get_leading_offset = function(line_nr, main_nr)
  if not OtterConfig.handle_leading_whitespace then
    return 0
  end

  local lang_chunks = keeper.rafts[main_nr].code_chunks
  for _, chunks in pairs(lang_chunks) do
    for _, chunk in ipairs(chunks) do
      if line_nr >= chunk.range.from[1] and line_nr <= chunk.range.to[1] then
        return chunk.leading_offset
      end
    end
  end
  return 0
end

---adjusts IN PLACE the position to include the start and end
---@param obj table
---@param main_nr number
---@param invert boolean?
---@param exclude_end boolean?
---@param known_offset number?
keeper.modify_position = function(obj, main_nr, invert, exclude_end, known_offset)
  if not OtterConfig.handle_leading_whitespace or known_offset == 0 then
    return
  end

  local sign = invert and -1 or 1
  local offset = known_offset

  -- there are apparently a lot of ranges that different language servers can use
  local ranges = { "range", "targetSelectionRange", "targetRange", "originSelectionRange" }
  for _, range in ipairs(ranges) do
    if obj[range] then
      local start = obj[range].start
      local end_ = obj[range]["end"]
      offset = offset or keeper.get_leading_offset(start.line, main_nr) * sign
      obj[range].start.character = start.character + offset
      if not exclude_end then
        obj[range]["end"].character = end_.character + offset
      end
    end
  end

  if obj.position then
    local pos = obj.position
    offset = offset or keeper.get_leading_offset(pos.line, main_nr) * sign
    obj.position.character = pos.character + offset
  end

  if obj.documentChanges then
    for _, change in ipairs(obj.documentChanges) do
      if change.edits then
        for _, edit in ipairs(change.edits) do
          keeper.modify_position(edit, main_nr, invert, exclude_end, offset)
        end
      end
    end
  end

  if obj.changes then
    for _, change in pairs(obj.changes) do
      for _, edit in ipairs(change) do
        keeper.modify_position(edit, main_nr, invert, exclude_end, offset)
      end
    end
  end

  if obj.newText then
    offset = offset or keeper.get_leading_offset(obj.range.start, main_nr) * sign
    local str = ""
    for _ = 1, offset, 1 do
      str = str .. " "
    end
    -- Put indents in front of newline, but ignore newlines that are followed by newlines
    obj.newText = string.gsub(obj.newText, "(\n)([^\n])", "%1" .. str .. "%2")
    obj.newText = string.gsub(obj.newText, "\n$", "\n" .. str) -- match a potential newline at the end
  end
end

---@param main_nr integer bufnr of the parent buffer
---@return boolean has_raft true if the buffer has otters attached
keeper.has_raft = function(main_nr)
  return keeper.rafts[main_nr] ~= nil
end

---@alias SyncResult
---| '"success"' -- The sync was successful
---| '"no_raft"' -- The buffer has no raft
---| '"textlock_active"' -- The buffer is currently locked
---| '"error"' -- Some other error occurred

--- Synchronize the raft of otters attached to a buffer
---@param main_nr integer bufnr of the parent buffer
---@param language string|nil only sync one otter buffer matching a language
---@return SyncResult result
keeper.sync_raft = function(main_nr, language)
  if not keeper.has_raft(main_nr) then
    return "no_raft"
  end
  local all_code_chunks
  local changetick = api.nvim_buf_get_changedtick(main_nr)
  if keeper.rafts[main_nr].last_changetick == changetick then
    all_code_chunks = keeper.rafts[main_nr].code_chunks
    return "success"
  end

  ---@param callback function
  ---@return SyncResult
  ---
  --- Assumption: If textlock is currently active,
  --- we can't sync, but those are also the cases
  --- in which it is not necessary to sync
  --- and can be delayed until the textlock is released.
  --- The lsp request should still be valid
  --- NOTE: We may be able to get rid of this entirely in nvim v0.11
  local function do_with_maybe_texlock(callback)
    local texlock_err_msg = "E565: Not allowed to change text or change window"
    local success, result = pcall(callback)
    if success then
      return "success"
    end

    result = tostring(result)
    if result:match(texlock_err_msg) then
      vim.schedule(callback)
      return "textlock_active"
    else
      return "error"
    end
  end

  all_code_chunks = keeper.extract_code_chunks(main_nr)

  keeper.rafts[main_nr].last_changetick = changetick
  keeper.rafts[main_nr].code_chunks = all_code_chunks

  local result
  local langs
  if language == nil then
    langs = keeper.rafts[main_nr].languages
  else
    langs = { language }
  end
  for _, lang in ipairs(langs) do
    local otter_nr = keeper.rafts[main_nr].buffers[lang]
    if otter_nr ~= nil then
      local code_chunks = all_code_chunks[lang]
      if code_chunks ~= nil then
        local nmax = code_chunks[#code_chunks].range["to"][1] -- last code line

        -- create list with empty lines the length of the buffer
        local ls = fn.empty_lines(nmax)

        -- set preamble lines
        local preamble = keeper.rafts[main_nr].preambles[lang]
        for i, l in ipairs(preamble) do
          table.remove(ls, i)
          table.insert(ls, i, l)
        end

        -- collect language lines
        -- are allowed to overwrite the preamble
        -- apply ignore_pattern filtering on read
        local pattern = keeper.rafts[main_nr].ignore_pattern[lang]
        for _, t in ipairs(code_chunks) do
          local start_index = t.range["from"][1]
          for i, l in ipairs(t.text) do
            local index = start_index + i
            if pattern == nil or not string.match(l, pattern) then
              ls[index] = l
            end
          end
        end

        -- set postamble lines
        local postamble = keeper.rafts[main_nr].postambles[lang]
        for _, l in ipairs(postamble) do
          table.insert(ls, l)
        end

        -- set code lines
        result = do_with_maybe_texlock(function()
          api.nvim_buf_set_lines(otter_nr, 0, -1, false, ls)
        end)
      else
        -- no code chunks so we wipe the otter buffer
        result = do_with_maybe_texlock(function()
          api.nvim_buf_set_lines(otter_nr, 0, -1, false, {})
        end)
      end
    end
  end

  return result
end

--- Export the raft of otters as files.
--- Asks for filename for each language.
---@param force boolean
keeper.export_raft = function(force)
  local main_nr = api.nvim_get_current_buf()
  keeper.sync_raft(main_nr)
  for _, otter_nr in pairs(keeper.rafts[main_nr].buffers) do
    local path = api.nvim_buf_get_name(otter_nr)
    local lang = keeper.rafts[main_nr].otter_nr_to_lang[otter_nr]
    local extension = OtterConfig.extensions[lang] or lang
    path = fn.otterpath_to_plain_path(path) .. "." .. extension
    vim.notify("Exporting otter: " .. lang)
    local new_path = vim.fn.input({ prompt = "New path: ", default = path, completion = "file" })
    if new_path ~= "" then
      api.nvim_set_current_buf(otter_nr)
      vim.lsp.buf.format({ bufnr = otter_nr })
      vim.cmd.write({ new_path, bang = force })
    end
    api.nvim_set_current_buf(main_nr)
  end
end

---Export only one language to a pre-specified filename
---@param language string
---@param fname string
---@param force boolean
keeper.export_otter_as = function(language, fname, force)
  local main_nr = api.nvim_get_current_buf()
  keeper.sync_raft(main_nr)
  for _, otter_nr in pairs(keeper.rafts[main_nr].buffers) do
    local path = api.nvim_buf_get_name(otter_nr)
    local lang = keeper.rafts[main_nr].otter_nr_to_lang[otter_nr]
    if lang ~= language then
      return
    end
    path = path:match("(.*[/\\])") .. fname
    api.nvim_set_current_buf(otter_nr)
    vim.lsp.buf.format()
    vim.cmd.write({ path, bang = force })
    api.nvim_set_current_buf(main_nr)
  end
end

---Get lines of code chunks managed by otter in a range of the current buffer.
---@param exclude_eval_false boolean | nil Exclude code chunks with eval: false
---@param row_start integer Row to start from, inclusive, 1-indexed.
---@param row_end integer Row to end at, inclusive, 1-indexed.
---@return string[]|nil Lines of code
keeper.get_language_lines = function(exclude_eval_false, row_start, row_end)
  local main_nr = vim.api.nvim_get_current_buf()
  keeper.sync_raft(main_nr)
  local lang = keeper.get_current_language_context()
  if lang == nil then
    return
  end

  local chunks = keeper.extract_code_chunks(main_nr, lang, exclude_eval_false, row_start, row_end)[lang]
  if not chunks or next(chunks) == nil then
    return
  end
  local code = {}
  for _, c in ipairs(chunks) do
    table.insert(code, fn.concat(c.text))
  end
  return code
end

---Get lines of code chunks managed by otter around the cursor in the current buffer.
---@return string|nil Lines of code
keeper.get_language_lines_around_cursor = function()
  local main_nr = vim.api.nvim_get_current_buf()
  local row, col = unpack(api.nvim_win_get_cursor(0))
  row = row - 1
  col = col

  local query = keeper.rafts[main_nr].query
  local parser = keeper.rafts[main_nr].parser
  local trees = parser:parse()
  assert(trees, "[otter] Treesitter failed to parse buffer " .. main_nr)
  local tree = trees[1]
  local root = tree:root()

  for _, match, metadata in query:iter_matches(root, main_nr, 0, -1, { all = true }) do
    for id, nodes in pairs(match) do
      local name = query.captures[id]

      -- TODO: maybe can be removed with nvim v0.10
      if type(nodes) ~= "table" then
        nodes = { nodes }
      end

      for _, node in ipairs(nodes) do
        if name == "content" then
          if ts.is_in_node_range(node, row, col) then
            return ts.get_node_text(node, main_nr, metadata)
          end
          -- chunks where the name of the language is the name of the capture
        elseif fn.contains(injectable_languages, name) then
          if ts.is_in_node_range(node, row, col) then
            return ts.get_node_text(node, main_nr, metadata)
          end
        end
      end
    end
  end
end

---Get lines of code chunks managed by otter in the current buffer up to the cursor.
---@param exclude_eval_false boolean|nil Exclude code chunks with eval: false
keeper.get_language_lines_to_cursor = function(exclude_eval_false)
  local row, _ = unpack(vim.api.nvim_win_get_cursor(0))
  row = row + 1
  return keeper.get_language_lines(exclude_eval_false, 0, row)
end

---Get lines of code chunks managed by otter in the current buffer from the cursor to the end.
---@param exclude_eval_false boolean|nil Exclude code chunks with eval: false
keeper.get_language_lines_from_cursor = function(exclude_eval_false)
  local row, _ = unpack(vim.api.nvim_win_get_cursor(0))
  row = row + 1
  return keeper.get_language_lines(exclude_eval_false, row, -1)
end

---Get lines of code chunks managed by otter in the current buffer from the cursor to the end.
---@param exclude_eval_false boolean|nil Exclude code chunks with eval: false
keeper.get_language_lines_in_visual_selection = function(exclude_eval_false)
  local lang = keeper.get_current_language_context()
  if lang == nil then
    return
  end
  local row_start, _ = unpack(api.nvim_buf_get_mark(0, "<"))
  local row_end, _ = unpack(api.nvim_buf_get_mark(0, ">"))
  return keeper.get_language_lines(exclude_eval_false, row_start, row_end)
end

return keeper
