local api = vim.api
local fn = vim.fn
local uv = vim.uv

local curl = require("plenary.curl")

local Utils = require("avante.utils")
local Config = require("avante.config")
local Path = require("avante.path")
local Providers = require("avante.providers")
local LLMToolHelpers = require("avante.llm_tools.helpers")
local LLMTools = require("avante.llm_tools")
local HistoryMessage = require("avante.history_message")

---@class avante.LLM
local M = {}

M.CANCEL_PATTERN = "AvanteLLMEscape"

------------------------------Prompt and type------------------------------

local group = api.nvim_create_augroup("avante_llm", { clear = true })

---@param content AvanteLLMMessageContent
---@param cb fun(title: string | nil): nil
function M.summarize_chat_thread_title(content, cb)
  local system_prompt =
    [[Summarize the content as a title for the chat thread. The title should be a concise and informative summary of the conversation, capturing the main points and key takeaways. It should be no longer than 100 words and should be written in a clear and engaging style. The title should be suitable for use as the title of a chat thread on a messaging platform or other communication medium. /no_think]]
  local response_content = ""
  local provider = Providers.get_memory_summary_provider()
  M.curl({
    provider = provider,
    prompt_opts = {
      system_prompt = system_prompt,
      messages = {
        { role = "user", content = content },
      },
    },
    handler_opts = {
      on_start = function(_) end,
      on_chunk = function(chunk)
        if not chunk then return end
        response_content = response_content .. chunk
      end,
      on_stop = function(stop_opts)
        if stop_opts.error ~= nil then
          Utils.error(string.format("summarize chat thread title failed: %s", vim.inspect(stop_opts.error)))
          return
        end
        if stop_opts.reason == "complete" then
          response_content = Utils.trim_think_content(response_content)
          response_content = Utils.trim(response_content, { prefix = "\n", suffix = "\n" })
          response_content = Utils.trim(response_content, { prefix = '"', suffix = '"' })
          local title = response_content
          cb(title)
        end
      end,
    },
  })
end

---@param prev_memory string | nil
---@param history_messages avante.HistoryMessage[]
---@param cb fun(memory: avante.ChatMemory | nil): nil
function M.summarize_memory(prev_memory, history_messages, cb)
  local system_prompt =
    [[You are an expert coding assistant. Your goal is to generate a concise, structured summary of the conversation below that captures all essential information needed to continue development after context replacement. Include tasks performed, code areas modified or reviewed, key decisions or assumptions, test results or errors, and outstanding tasks or next steps.]]
  if #history_messages == 0 then
    cb(nil)
    return
  end
  local latest_timestamp = nil
  local latest_message_uuid = nil
  for idx = #history_messages, 1, -1 do
    local message = history_messages[idx]
    if not message.is_dummy then
      latest_timestamp = message.timestamp
      latest_message_uuid = message.uuid
      break
    end
  end
  if not latest_timestamp or not latest_message_uuid then
    cb(nil)
    return
  end
  local conversation_items = vim
    .iter(history_messages)
    :map(function(msg) return msg.message.role .. ": " .. Utils.message_to_text(msg, history_messages) end)
    :totable()
  local conversation_text = table.concat(conversation_items, "\n")
  local user_prompt = "Here is the conversation so far:\n"
    .. conversation_text
    .. "\n\nPlease summarize this conversation, covering:\n1. Tasks performed and outcomes\n2. Code files, modules, or functions modified or examined\n3. Important decisions or assumptions made\n4. Errors encountered and test or build results\n5. Remaining tasks, open questions, or next steps\nProvide the summary in a clear, concise format."
  if prev_memory then user_prompt = user_prompt .. "\n\nThe previous summary is:\n\n" .. prev_memory end
  local messages = {
    {
      role = "user",
      content = user_prompt,
    },
  }
  local response_content = ""
  local provider = Providers.get_memory_summary_provider()
  M.curl({
    provider = provider,
    prompt_opts = {
      system_prompt = system_prompt,
      messages = messages,
    },
    handler_opts = {
      on_start = function(_) end,
      on_chunk = function(chunk)
        if not chunk then return end
        response_content = response_content .. chunk
      end,
      on_stop = function(stop_opts)
        if stop_opts.error ~= nil then
          Utils.error(string.format("summarize memory failed: %s", vim.inspect(stop_opts.error)))
          return
        end
        if stop_opts.reason == "complete" then
          response_content = Utils.trim_think_content(response_content)
          local memory = {
            content = response_content,
            last_summarized_timestamp = latest_timestamp,
            last_message_uuid = latest_message_uuid,
          }
          cb(memory)
        else
          cb(nil)
        end
      end,
    },
  })
end

---@param user_input string
---@param cb fun(error: string | nil): nil
function M.generate_todos(user_input, cb)
  local system_prompt =
    [[You are an expert coding assistant. Please generate a todo list to complete the task based on the user input and pass the todo list to the add_todos tool.]]
  local messages = {
    { role = "user", content = user_input },
  }

  local provider = Providers[Config.provider]
  local tools = {
    require("avante.llm_tools.add_todos"),
  }

  local history_messages = {}
  cb = Utils.call_once(cb)

  M.curl({
    provider = provider,
    prompt_opts = {
      system_prompt = system_prompt,
      messages = messages,
      tools = tools,
    },
    handler_opts = {
      on_start = function() end,
      on_chunk = function() end,
      on_messages_add = function(msgs)
        msgs = vim.islist(msgs) and msgs or { msgs }
        for _, msg in ipairs(msgs) do
          if not msg.uuid then msg.uuid = Utils.uuid() end
          local idx = nil
          for i, m in ipairs(history_messages) do
            if m.uuid == msg.uuid then
              idx = i
              break
            end
          end
          if idx ~= nil then
            history_messages[idx] = msg
          else
            table.insert(history_messages, msg)
          end
        end
      end,
      on_stop = function(stop_opts)
        if stop_opts.error ~= nil then
          Utils.error(string.format("generate todos failed: %s", vim.inspect(stop_opts.error)))
          return
        end
        if stop_opts.reason == "tool_use" then
          local uncalled_tool_uses = Utils.get_uncalled_tool_uses(history_messages)
          for _, partial_tool_use in ipairs(uncalled_tool_uses) do
            if partial_tool_use.state == "generated" and partial_tool_use.name == "add_todos" then
              LLMTools.process_tool_use(tools, partial_tool_use, function() end, function() cb() end, {})
              cb()
            end
          end
        else
          cb()
        end
      end,
    },
  })
end

---@param opts AvanteGeneratePromptsOptions
---@return AvantePromptOptions
function M.generate_prompts(opts)
  local provider = opts.provider or Providers[Config.provider]
  local mode = opts.mode or Config.mode

  -- Check if the instructions contains an image path
  local image_paths = {}
  if opts.prompt_opts and opts.prompt_opts.image_paths then
    image_paths = vim.list_extend(image_paths, opts.prompt_opts.image_paths)
  end

  local project_root = Utils.root.get()
  Path.prompts.initialize(Path.prompts.get_templates_dir(project_root), project_root)

  local system_info = Utils.get_system_info()

  local selected_files = opts.selected_files or {}

  if opts.selected_filepaths then
    for _, filepath in ipairs(opts.selected_filepaths) do
      local lines, error = Utils.read_file_from_buf_or_disk(filepath)
      lines = lines or {}
      local filetype = Utils.get_filetype(filepath)
      if error ~= nil then
        Utils.error("error reading file: " .. error)
      else
        local content = table.concat(lines, "\n")
        table.insert(selected_files, { path = filepath, content = content, file_type = filetype })
      end
    end
  end

  local viewed_files = {}

  if opts.history_messages then
    for _, message in ipairs(opts.history_messages) do
      local content = message.message.content
      if type(content) ~= "table" then goto continue end
      for _, item in ipairs(content) do
        if type(item) ~= "table" then goto continue1 end
        if item.type ~= "tool_use" then goto continue1 end
        local tool_name = item.name
        if tool_name ~= "view" then goto continue1 end
        local path = item.input.path
        if path then
          local uniform_path = Utils.uniform_path(path)
          viewed_files[uniform_path] = item.id
        end
        ::continue1::
      end
      ::continue::
    end
  end

  selected_files = vim.iter(selected_files):filter(function(file) return viewed_files[file.path] == nil end):totable()

  local template_opts = {
    ask = opts.ask, -- TODO: add mode without ask instruction
    code_lang = opts.code_lang,
    selected_files = selected_files,
    selected_code = opts.selected_code,
    recently_viewed_files = opts.recently_viewed_files,
    project_context = opts.project_context,
    diagnostics = opts.diagnostics,
    system_info = system_info,
    model_name = provider.model or "unknown",
    memory = opts.memory,
  }

  if opts.get_todos then
    local todos = opts.get_todos()
    if todos and #todos > 0 then template_opts.todos = vim.json.encode(todos) end
  end

  local system_prompt
  if opts.prompt_opts and opts.prompt_opts.system_prompt then
    system_prompt = opts.prompt_opts.system_prompt
  else
    system_prompt = Path.prompts.render_mode(mode, template_opts)
  end

  if Config.system_prompt ~= nil then
    local custom_system_prompt = Config.system_prompt
    if type(custom_system_prompt) == "function" then custom_system_prompt = custom_system_prompt() end
    if custom_system_prompt ~= nil and custom_system_prompt ~= "" and custom_system_prompt ~= "null" then
      system_prompt = system_prompt .. "\n\n" .. custom_system_prompt
    end
  end

  ---@type AvanteLLMMessage[]
  local context_messages = {}
  if opts.prompt_opts and opts.prompt_opts.messages then
    context_messages = vim.list_extend(context_messages, opts.prompt_opts.messages)
  end

  if opts.project_context ~= nil and opts.project_context ~= "" and opts.project_context ~= "null" then
    local project_context = Path.prompts.render_file("_project.avanterules", template_opts)
    if project_context ~= "" then
      table.insert(context_messages, { role = "user", content = project_context, visible = false, is_context = true })
    end
  end

  if opts.diagnostics ~= nil and opts.diagnostics ~= "" and opts.diagnostics ~= "null" then
    local diagnostics = Path.prompts.render_file("_diagnostics.avanterules", template_opts)
    if diagnostics ~= "" then
      table.insert(context_messages, { role = "user", content = diagnostics, visible = false, is_context = true })
    end
  end

  if #selected_files > 0 or opts.selected_code ~= nil then
    local code_context = Path.prompts.render_file("_context.avanterules", template_opts)
    if code_context ~= "" then
      table.insert(context_messages, { role = "user", content = code_context, visible = false, is_context = true })
    end
  end

  if opts.memory ~= nil and opts.memory ~= "" and opts.memory ~= "null" then
    local memory = Path.prompts.render_file("_memory.avanterules", template_opts)
    if memory ~= "" then
      table.insert(context_messages, { role = "user", content = memory, visible = false, is_context = true })
    end
  end

  local pending_compaction_history_messages = {}
  if opts.prompt_opts and opts.prompt_opts.pending_compaction_history_messages then
    pending_compaction_history_messages =
      vim.list_extend(pending_compaction_history_messages, opts.prompt_opts.pending_compaction_history_messages)
  end

  local context_window = provider.context_window

  if context_window and context_window > 0 then
    Utils.debug("Context window", context_window)
    if opts.get_tokens_usage then
      local tokens_usage = opts.get_tokens_usage()
      if tokens_usage then
        local target_tokens = context_window * 0.9
        local tokens_count = tokens_usage.prompt_tokens + tokens_usage.completion_tokens
        Utils.debug("Tokens count", tokens_count)
        if tokens_count > target_tokens then pending_compaction_history_messages = opts.history_messages end
      end
    end
  end

  ---@type AvanteLLMMessage[]
  local messages = vim.deepcopy(context_messages)
  for _, msg in ipairs(opts.history_messages or {}) do
    local message = msg.message
    if msg.is_user_submission then
      message = vim.deepcopy(message)
      local content = message.content
      if type(content) == "string" then
        message.content = "<task>" .. content .. "</task>"
      elseif type(content) == "table" then
        for idx, item in ipairs(content) do
          if type(item) == "string" then
            item = "<task>" .. item .. "</task>"
            content[idx] = item
          elseif type(item) == "table" and item.type == "text" then
            item.content = "<task>" .. item.content .. "</task>"
            content[idx] = item
          end
        end
      end
    end
    table.insert(messages, message)
  end

  messages = vim
    .iter(messages)
    :filter(function(msg) return type(msg.content) ~= "string" or msg.content ~= "" end)
    :totable()

  if opts.instructions ~= nil and opts.instructions ~= "" then
    messages = vim.list_extend(messages, { { role = "user", content = opts.instructions } })
  end

  opts.session_ctx = opts.session_ctx or {}
  opts.session_ctx.system_prompt = system_prompt
  opts.session_ctx.messages = messages

  local tools = {}
  if opts.tools then tools = vim.list_extend(tools, opts.tools) end
  if opts.prompt_opts and opts.prompt_opts.tools then tools = vim.list_extend(tools, opts.prompt_opts.tools) end

  ---@type AvantePromptOptions
  return {
    system_prompt = system_prompt,
    messages = messages,
    image_paths = image_paths,
    tools = tools,
    pending_compaction_history_messages = pending_compaction_history_messages,
  }
end

---@param opts AvanteGeneratePromptsOptions
---@return integer
function M.calculate_tokens(opts)
  local prompt_opts = M.generate_prompts(opts)
  local tokens = Utils.tokens.calculate_tokens(prompt_opts.system_prompt)
  for _, message in ipairs(prompt_opts.messages) do
    tokens = tokens + Utils.tokens.calculate_tokens(message.content)
  end
  return tokens
end

local parse_headers = function(headers_file)
  local headers = {}
  local file = io.open(headers_file, "r")
  if file then
    for line in file:lines() do
      line = line:gsub("\r$", "")
      local key, value = line:match("^%s*(.-)%s*:%s*(.*)$")
      if key and value then headers[key] = value end
    end
    file:close()
  end
  return headers
end

---@param opts avante.CurlOpts
function M.curl(opts)
  local provider = opts.provider
  local prompt_opts = opts.prompt_opts
  local handler_opts = opts.handler_opts

  local orig_on_stop = handler_opts.on_stop
  local stopped = false
  ---@param stop_opts AvanteLLMStopCallbackOptions
  handler_opts.on_stop = function(stop_opts)
    if stop_opts and not stop_opts.streaming_tool_use then
      if stopped then return end
      stopped = true
    end
    if orig_on_stop then return orig_on_stop(stop_opts) end
  end

  ---@type AvanteCurlOutput
  local spec = provider:parse_curl_args(prompt_opts)

  ---@type string
  local current_event_state = nil
  local turn_ctx = {}
  turn_ctx.turn_id = Utils.uuid()

  local response_body = ""
  ---@param line string
  local function parse_stream_data(line)
    local event = line:match("^event:%s*(.+)$")
    if event then
      current_event_state = event
      return
    end
    local data_match = line:match("^data:%s*(.+)$")
    if data_match then
      response_body = ""
      provider:parse_response(turn_ctx, data_match, current_event_state, handler_opts)
    else
      response_body = response_body .. line
      local ok, jsn = pcall(vim.json.decode, response_body)
      if ok then
        if jsn.error then
          handler_opts.on_stop({ reason = "error", error = jsn.error })
        else
          provider:parse_response(turn_ctx, response_body, current_event_state, handler_opts)
        end
        response_body = ""
      end
    end
  end

  local function parse_response_without_stream(data)
    provider:parse_response_without_stream(data, current_event_state, handler_opts)
  end

  local completed = false

  local active_job

  local temp_file = fn.tempname()
  local curl_body_file = temp_file .. "-request-body.json"
  local resp_body_file = temp_file .. "-response-body.txt"
  local headers_file = temp_file .. "-response-headers.txt"
  local json_content = vim.json.encode(spec.body)
  fn.writefile(vim.split(json_content, "\n"), curl_body_file)

  Utils.debug("curl request body file:", curl_body_file)
  Utils.debug("curl response body file:", resp_body_file)
  Utils.debug("curl headers file:", headers_file)

  local function cleanup()
    if Config.debug then return end
    vim.schedule(function()
      fn.delete(curl_body_file)
      pcall(fn.delete, resp_body_file)
      fn.delete(headers_file)
    end)
  end

  local headers_reported = false

  active_job = curl.post(spec.url, {
    headers = spec.headers,
    proxy = spec.proxy,
    insecure = spec.insecure,
    body = curl_body_file,
    raw = spec.rawArgs,
    dump = { "-D", headers_file },
    stream = function(err, data, _)
      if not headers_reported and opts.on_response_headers then
        headers_reported = true
        opts.on_response_headers(parse_headers(headers_file))
      end
      if err then
        completed = true
        handler_opts.on_stop({ reason = "error", error = err })
        return
      end
      if not data then return end
      if Config.debug then
        if type(data) == "string" then
          local file = io.open(resp_body_file, "a")
          if file then
            file:write(data .. "\n")
            file:close()
          end
        end
      end
      vim.schedule(function()
        if provider.parse_stream_data ~= nil then
          provider:parse_stream_data(turn_ctx, data, handler_opts)
        else
          parse_stream_data(data)
        end
      end)
    end,
    on_error = function(err)
      if err.exit == 23 then
        local xdg_runtime_dir = os.getenv("XDG_RUNTIME_DIR")
        if not xdg_runtime_dir or fn.isdirectory(xdg_runtime_dir) == 0 then
          Utils.error(
            "$XDG_RUNTIME_DIR="
              .. xdg_runtime_dir
              .. " is set but does not exist. curl could not write output. Please make sure it exists, or unset.",
            { title = "Avante" }
          )
        elseif not uv.fs_access(xdg_runtime_dir, "w") then
          Utils.error(
            "$XDG_RUNTIME_DIR="
              .. xdg_runtime_dir
              .. " exists but is not writable. curl could not write output. Please make sure it is writable, or unset.",
            { title = "Avante" }
          )
        end
      end

      active_job = nil
      if not completed then
        completed = true
        cleanup()
        handler_opts.on_stop({ reason = "error", error = err })
      end
    end,
    callback = function(result)
      active_job = nil
      cleanup()
      local headers_map = vim.iter(result.headers):fold({}, function(acc, value)
        local pieces = vim.split(value, ":")
        local key = pieces[1]
        local remain = vim.list_slice(pieces, 2)
        if not remain then return acc end
        local val = Utils.trim_spaces(table.concat(remain, ":"))
        acc[key] = val
        return acc
      end)
      if result.status >= 400 then
        if provider.on_error then
          provider.on_error(result)
        else
          Utils.error("API request failed with status " .. result.status, { once = true, title = "Avante" })
        end
        local retry_after = 10
        if headers_map["retry-after"] then retry_after = tonumber(headers_map["retry-after"]) or 10 end
        if result.status == 429 then
          handler_opts.on_stop({ reason = "rate_limit", retry_after = retry_after })
          return
        end
        vim.schedule(function()
          if not completed then
            completed = true
            handler_opts.on_stop({
              reason = "error",
              error = "API request failed with status " .. result.status .. ". Body: " .. vim.inspect(result.body),
            })
          end
        end)
      end

      -- If stream is not enabled, then handle the response here
      if provider:is_disable_stream() and result.status == 200 then
        vim.schedule(function()
          completed = true
          parse_response_without_stream(result.body)
        end)
      end

      if result.status == 200 and spec.url:match("https://openrouter.ai") then
        local content_type = headers_map["content-type"]
        if content_type and content_type:match("text/html") then
          handler_opts.on_stop({
            reason = "error",
            error = "Your openrouter endpoint setting is incorrect, please set it to https://openrouter.ai/api/v1",
          })
        end
      end
    end,
  })

  api.nvim_create_autocmd("User", {
    group = group,
    pattern = M.CANCEL_PATTERN,
    once = true,
    callback = function()
      -- Error: cannot resume dead coroutine
      if active_job then
        -- Mark as completed first to prevent error handler from running
        completed = true

        -- Attempt to shutdown the active job, but ignore any errors
        xpcall(function() active_job:shutdown() end, function(err)
          Utils.debug("Ignored error during job shutdown: " .. vim.inspect(err))
          return err
        end)

        Utils.debug("LLM request cancelled")
        active_job = nil

        -- Clean up and notify of cancellation
        cleanup()
        vim.schedule(function() handler_opts.on_stop({ reason = "cancelled" }) end)
      end
    end,
  })

  return active_job
end

---@param opts AvanteLLMStreamOptions
function M._stream(opts)
  -- Reset the cancellation flag at the start of a new request
  if LLMToolHelpers then LLMToolHelpers.is_cancelled = false end

  local provider = opts.provider or Providers[Config.provider]
  opts.session_ctx = opts.session_ctx or {}

  if not opts.session_ctx.on_messages_add then opts.session_ctx.on_messages_add = opts.on_messages_add end
  if not opts.session_ctx.on_state_change then opts.session_ctx.on_state_change = opts.on_state_change end
  if not opts.session_ctx.on_start then opts.session_ctx.on_start = opts.on_start end
  if not opts.session_ctx.on_chunk then opts.session_ctx.on_chunk = opts.on_chunk end
  if not opts.session_ctx.on_stop then opts.session_ctx.on_stop = opts.on_stop end
  if not opts.session_ctx.on_tool_log then opts.session_ctx.on_tool_log = opts.on_tool_log end
  if not opts.session_ctx.get_history_messages then
    opts.session_ctx.get_history_messages = opts.get_history_messages
  end

  ---@cast provider AvanteProviderFunctor

  local prompt_opts = M.generate_prompts(opts)

  if
    prompt_opts.pending_compaction_history_messages
    and #prompt_opts.pending_compaction_history_messages > 0
    and opts.on_memory_summarize
  then
    opts.on_memory_summarize(prompt_opts.pending_compaction_history_messages)
    return
  end

  local resp_headers = {}

  ---@type AvanteHandlerOptions
  local handler_opts = {
    on_messages_add = opts.on_messages_add,
    on_state_change = opts.on_state_change,
    update_tokens_usage = opts.update_tokens_usage,
    on_start = opts.on_start,
    on_chunk = opts.on_chunk,
    on_stop = function(stop_opts)
      if stop_opts.usage and opts.update_tokens_usage then opts.update_tokens_usage(stop_opts.usage) end

      ---@param partial_tool_use_list AvantePartialLLMToolUse[]
      ---@param tool_use_index integer
      ---@param tool_results AvanteLLMToolResult[]
      local function handle_next_tool_use(partial_tool_use_list, tool_use_index, tool_results, streaming_tool_use)
        if tool_use_index > #partial_tool_use_list then
          ---@type avante.HistoryMessage[]
          local messages = {}
          for _, tool_result in ipairs(tool_results) do
            messages[#messages + 1] = HistoryMessage:new({
              role = "user",
              content = {
                {
                  type = "tool_result",
                  tool_use_id = tool_result.tool_use_id,
                  content = tool_result.content,
                  is_error = tool_result.is_error,
                  is_user_declined = tool_result.is_user_declined,
                },
              },
            })
          end
          if opts.on_messages_add then opts.on_messages_add(messages) end
          local the_last_tool_use = partial_tool_use_list[#partial_tool_use_list]
          if the_last_tool_use and the_last_tool_use.name == "attempt_completion" then
            opts.on_stop({ reason = "complete" })
            return
          end
          local new_opts = vim.tbl_deep_extend("force", opts, {
            history_messages = opts.get_history_messages(),
          })
          if provider.get_rate_limit_sleep_time then
            local sleep_time = provider:get_rate_limit_sleep_time(resp_headers)
            if sleep_time and sleep_time > 0 then
              Utils.info("Rate limit reached. Sleeping for " .. sleep_time .. " seconds ...")
              vim.defer_fn(function() M._stream(new_opts) end, sleep_time * 1000)
              return
            end
          end
          M._stream(new_opts)
          return
        end
        local partial_tool_use = partial_tool_use_list[tool_use_index]
        ---@param result string | nil
        ---@param error string | nil
        local function handle_tool_result(result, error)
          -- Special handling for cancellation signal from tools
          if error == LLMToolHelpers.CANCEL_TOKEN then
            Utils.debug("Tool execution was cancelled by user")
            if opts.on_chunk then opts.on_chunk("\n*[Request cancelled by user during tool execution.]*\n") end
            if opts.on_messages_add then
              local message = HistoryMessage:new({
                role = "assistant",
                content = "\n*[Request cancelled by user during tool execution.]*\n",
              }, {
                just_for_display = true,
              })
              opts.on_messages_add({ message })
            end
            return opts.on_stop({ reason = "cancelled" })
          end

          local is_user_declined = error and error:match("^User declined")
          local tool_result = {
            tool_use_id = partial_tool_use.id,
            content = error ~= nil and error or result,
            is_error = error ~= nil, -- Keep this as error to prevent processing as success
            is_user_declined = is_user_declined ~= nil,
          }
          table.insert(tool_results, tool_result)
          return handle_next_tool_use(partial_tool_use_list, tool_use_index + 1, tool_results)
        end
        local is_edit_tool_use = Utils.is_edit_func_call_tool_use(partial_tool_use)
        local support_streaming = false
        local llm_tool = vim.iter(prompt_opts.tools):find(function(tool) return tool.name == partial_tool_use.name end)
        if llm_tool then support_streaming = llm_tool.support_streaming == true end
        if partial_tool_use.state == "generating" and not is_edit_tool_use and not support_streaming then return end
        if type(partial_tool_use.input) == "table" then partial_tool_use.input.tool_use_id = partial_tool_use.id end
        if partial_tool_use.state == "generating" then
          if type(partial_tool_use.input) == "table" then
            partial_tool_use.input.streaming = true
            LLMTools.process_tool_use(
              prompt_opts.tools,
              partial_tool_use,
              function() end,
              function() end,
              opts.session_ctx
            )
          end
          return
        else
          if streaming_tool_use then return end
        end
        -- Either on_complete handles the tool result asynchronously or we receive the result and error synchronously when either is not nil
        local result, error = LLMTools.process_tool_use(
          prompt_opts.tools,
          partial_tool_use,
          opts.on_tool_log,
          handle_tool_result,
          opts.session_ctx
        )
        if result ~= nil or error ~= nil then return handle_tool_result(result, error) end
      end
      if stop_opts.reason == "cancelled" then
        if opts.on_chunk then opts.on_chunk("\n*[Request cancelled by user.]*\n") end
        if opts.on_messages_add then
          local message = HistoryMessage:new({
            role = "assistant",
            content = "\n*[Request cancelled by user.]*\n",
          }, {
            just_for_display = true,
          })
          opts.on_messages_add({ message })
        end
        return opts.on_stop({ reason = "cancelled" })
      end
      local history_messages = opts.get_history_messages and opts.get_history_messages({ all = true }) or {}
      local uncalled_tool_uses = Utils.get_uncalled_tool_uses(history_messages)
      if stop_opts.reason == "complete" and Config.mode == "agentic" then
        local completed_attempt_completion_tool_use = nil
        for idx = #history_messages, 1, -1 do
          local message = history_messages[idx]
          if message.is_user_submission then break end
          if not Utils.is_tool_use_message(message) then goto continue end
          if message.message.content[1].name ~= "attempt_completion" then break end
          completed_attempt_completion_tool_use = message
          if message then break end
          ::continue::
        end
        local unfinished_todos = {}
        if opts.get_todos then
          local todos = opts.get_todos()
          unfinished_todos = vim.tbl_filter(
            function(todo) return todo.status ~= "done" or todo.status ~= "cancelled" end,
            todos
          )
        end
        local user_reminder_count = opts.session_ctx.user_reminder_count or 0
        if
          not completed_attempt_completion_tool_use
          and opts.on_messages_add
          and (user_reminder_count < 3 or #unfinished_todos > 0)
        then
          opts.session_ctx.user_reminder_count = user_reminder_count + 1
          Utils.debug("user reminder count", user_reminder_count)
          local message
          if #unfinished_todos > 0 then
            message = HistoryMessage:new({
              role = "user",
              content = "<user-reminder>You should use tool calls to answer the question, for example, use update_todo_status if the task step is done or cancelled.</user-reminder>",
            }, {
              visible = false,
            })
          else
            message = HistoryMessage:new({
              role = "user",
              content = "<user-reminder>You should use tool calls to answer the question, for example, use attempt_completion if the job is done.</user-reminder>",
            }, {
              visible = false,
            })
          end
          opts.on_messages_add({ message })
          local new_opts = vim.tbl_deep_extend("force", opts, {
            history_messages = opts.get_history_messages(),
          })
          if provider.get_rate_limit_sleep_time then
            local sleep_time = provider:get_rate_limit_sleep_time(resp_headers)
            if sleep_time and sleep_time > 0 then
              Utils.info("Rate limit reached. Sleeping for " .. sleep_time .. " seconds ...")
              vim.defer_fn(function() M._stream(new_opts) end, sleep_time * 1000)
              return
            end
          end
          M._stream(new_opts)
          return
        end
      end
      if stop_opts.reason == "tool_use" then
        opts.session_ctx.user_reminder_count = 0
        return handle_next_tool_use(uncalled_tool_uses, 1, {}, stop_opts.streaming_tool_use)
      end
      if stop_opts.reason == "rate_limit" then
        local msg_content = "*[Rate limit reached. Retrying in " .. stop_opts.retry_after .. " seconds ...]*"
        if opts.on_chunk then opts.on_chunk("\n" .. msg_content .. "\n") end
        local message
        if opts.on_messages_add then
          message = HistoryMessage:new({
            role = "assistant",
            content = "\n\n" .. msg_content,
          }, {
            just_for_display = true,
          })
          opts.on_messages_add({ message })
        end
        local timer = vim.loop.new_timer()
        if timer then
          local retry_after = stop_opts.retry_after
          local function countdown()
            timer:start(
              1000,
              0,
              vim.schedule_wrap(function()
                if retry_after > 0 then retry_after = retry_after - 1 end
                local msg_content_ = "*[Rate limit reached. Retrying in " .. retry_after .. " seconds ...]*"
                if opts.on_chunk then opts.on_chunk([[\033[1A\033[K]] .. "\n" .. msg_content_ .. "\n") end
                if opts.on_messages_add and message then
                  message.message.content = "\n\n" .. msg_content_
                  opts.on_messages_add({ message })
                end
                countdown()
              end)
            )
          end
          countdown()
        end
        Utils.info("Rate limit reached. Retrying in " .. stop_opts.retry_after .. " seconds", { title = "Avante" })
        vim.defer_fn(function()
          if timer then timer:stop() end
          M._stream(opts)
        end, stop_opts.retry_after * 1000)
        return
      end
      return opts.on_stop(stop_opts)
    end,
  }

  return M.curl({
    provider = provider,
    prompt_opts = prompt_opts,
    handler_opts = handler_opts,
    on_response_headers = function(headers) resp_headers = headers end,
  })
end

local function _merge_response(first_response, second_response, opts)
  local prompt = "\n" .. Config.dual_boost.prompt
  prompt = prompt
    :gsub("{{[%s]*provider1_output[%s]*}}", function() return first_response end)
    :gsub("{{[%s]*provider2_output[%s]*}}", function() return second_response end)

  prompt = prompt .. "\n"

  if opts.instructions == nil then opts.instructions = "" end

  -- append this reference prompt to the prompt_opts messages at last
  opts.instructions = opts.instructions .. prompt

  M._stream(opts)
end

local function _collector_process_responses(collector, opts)
  if not collector[1] or not collector[2] then
    Utils.error("One or both responses failed to complete")
    return
  end
  _merge_response(collector[1], collector[2], opts)
end

local function _collector_add_response(collector, index, response, opts)
  collector[index] = response
  collector.count = collector.count + 1

  if collector.count == 2 then
    collector.timer:stop()
    _collector_process_responses(collector, opts)
  end
end

function M._dual_boost_stream(opts, Provider1, Provider2)
  Utils.debug("Starting Dual Boost Stream")

  local collector = {
    count = 0,
    responses = {},
    timer = uv.new_timer(),
    timeout_ms = Config.dual_boost.timeout,
  }

  -- Setup timeout
  collector.timer:start(
    collector.timeout_ms,
    0,
    vim.schedule_wrap(function()
      if collector.count < 2 then
        Utils.warn("Dual boost stream timeout reached")
        collector.timer:stop()
        -- Process whatever responses we have
        _collector_process_responses(collector, opts)
      end
    end)
  )

  -- Create options for both streams
  local function create_stream_opts(index)
    local response = ""
    return vim.tbl_extend("force", opts, {
      on_chunk = function(chunk)
        if chunk then response = response .. chunk end
      end,
      on_stop = function(stop_opts)
        if stop_opts.error then
          Utils.error(string.format("Stream %d failed: %s", index, stop_opts.error))
          return
        end
        Utils.debug(string.format("Response %d completed", index))
        _collector_add_response(collector, index, response, opts)
      end,
    })
  end

  -- Start both streams
  local success, err = xpcall(function()
    local opts1 = create_stream_opts(1)
    opts1.provider = Provider1
    M._stream(opts1)
    local opts2 = create_stream_opts(2)
    opts2.provider = Provider2
    M._stream(opts2)
  end, function(err) return err end)
  if not success then Utils.error("Failed to start dual_boost streams: " .. tostring(err)) end
end

---@param opts AvanteLLMStreamOptions
function M.stream(opts)
  local is_completed = false
  if opts.on_tool_log ~= nil then
    local original_on_tool_log = opts.on_tool_log
    opts.on_tool_log = vim.schedule_wrap(function(...)
      if not original_on_tool_log then return end
      return original_on_tool_log(...)
    end)
  end
  if opts.on_chunk ~= nil then
    local original_on_chunk = opts.on_chunk
    opts.on_chunk = vim.schedule_wrap(function(chunk)
      if is_completed then return end
      if original_on_chunk then return original_on_chunk(chunk) end
    end)
  end
  if opts.on_stop ~= nil then
    local original_on_stop = opts.on_stop
    opts.on_stop = vim.schedule_wrap(function(stop_opts)
      if is_completed then return end
      if stop_opts.reason == "complete" or stop_opts.reason == "error" or stop_opts.reason == "cancelled" then
        is_completed = true
      end
      return original_on_stop(stop_opts)
    end)
  end

  local valid_dual_boost_modes = {
    legacy = true,
  }

  opts.mode = opts.mode or Config.mode

  if Config.dual_boost.enabled and valid_dual_boost_modes[opts.mode] then
    M._dual_boost_stream(
      opts,
      Providers[Config.dual_boost.first_provider],
      Providers[Config.dual_boost.second_provider]
    )
  else
    M._stream(opts)
  end
end

function M.cancel_inflight_request()
  if LLMToolHelpers.is_cancelled ~= nil then LLMToolHelpers.is_cancelled = true end
  if LLMToolHelpers.confirm_popup ~= nil then
    LLMToolHelpers.confirm_popup:cancel()
    LLMToolHelpers.confirm_popup = nil
  end

  api.nvim_exec_autocmds("User", { pattern = M.CANCEL_PATTERN })
end

return M