local Utils = require("avante.utils")
local Clipboard = require("avante.clipboard")
local P = require("avante.providers")
local HistoryMessage = require("avante.history.message")
local JsonParser = require("avante.libs.jsonparser")
local Config = require("avante.config")

---@class AvanteProviderFunctor
local M = {}

M.api_key_name = "ANTHROPIC_API_KEY"
M.support_prompt_caching = true

M.role_map = {
  user = "user",
  assistant = "assistant",
}

---@param headers table<string, string>
---@return integer|nil
function M:get_rate_limit_sleep_time(headers)
  local remaining_tokens = tonumber(headers["anthropic-ratelimit-tokens-remaining"])
  if remaining_tokens == nil then return end
  if remaining_tokens > 10000 then return end
  local reset_dt_str = headers["anthropic-ratelimit-tokens-reset"]
  if remaining_tokens ~= 0 then reset_dt_str = reset_dt_str or headers["anthropic-ratelimit-requests-reset"] end
  local reset_dt, err = Utils.parse_iso8601_date(reset_dt_str)
  if err then
    Utils.warn(err)
    return
  end
  local now = Utils.utc_now()
  return Utils.datetime_diff(tostring(now), tostring(reset_dt))
end

---@param tool AvanteLLMTool
---@return AvanteClaudeTool
function M:transform_tool(tool)
  local input_schema_properties, required = Utils.llm_tool_param_fields_to_json_schema(tool.param.fields)
  return {
    name = tool.name,
    description = tool.get_description and tool.get_description() or tool.description,
    input_schema = {
      type = "object",
      properties = input_schema_properties,
      required = required,
    },
  }
end

function M:is_disable_stream() return false end

---@return AvanteClaudeMessage[]
function M:parse_messages(opts)
  ---@type AvanteClaudeMessage[]
  local messages = {}

  local provider_conf, _ = P.parse_config(self)

  ---@type {idx: integer, length: integer}[]
  local messages_with_length = {}
  for idx, message in ipairs(opts.messages) do
    table.insert(messages_with_length, { idx = idx, length = Utils.tokens.calculate_tokens(message.content) })
  end

  table.sort(messages_with_length, function(a, b) return a.length > b.length end)

  local has_tool_use = false
  for _, message in ipairs(opts.messages) do
    local content_items = message.content
    local message_content = {}
    if type(content_items) == "string" then
      if message.role == "assistant" then content_items = content_items:gsub("%s+$", "") end
      if content_items ~= "" then
        table.insert(message_content, {
          type = "text",
          text = content_items,
        })
      end
    elseif type(content_items) == "table" then
      ---@cast content_items AvanteLLMMessageContentItem[]
      for _, item in ipairs(content_items) do
        if type(item) == "string" then
          if message.role == "assistant" then item = item:gsub("%s+$", "") end
          table.insert(message_content, { type = "text", text = item })
        elseif type(item) == "table" and item.type == "text" then
          table.insert(message_content, { type = "text", text = item.text })
        elseif type(item) == "table" and item.type == "image" then
          table.insert(message_content, { type = "image", source = item.source })
        elseif not provider_conf.disable_tools and type(item) == "table" and item.type == "tool_use" then
          has_tool_use = true
          table.insert(message_content, { type = "tool_use", name = item.name, id = item.id, input = item.input })
        elseif
          not provider_conf.disable_tools
          and type(item) == "table"
          and item.type == "tool_result"
          and has_tool_use
        then
          table.insert(
            message_content,
            { type = "tool_result", tool_use_id = item.tool_use_id, content = item.content, is_error = item.is_error }
          )
        elseif type(item) == "table" and item.type == "thinking" then
          table.insert(message_content, { type = "thinking", thinking = item.thinking, signature = item.signature })
        elseif type(item) == "table" and item.type == "redacted_thinking" then
          table.insert(message_content, { type = "redacted_thinking", data = item.data })
        end
      end
    end
    if #message_content > 0 then
      table.insert(messages, {
        role = self.role_map[message.role],
        content = message_content,
      })
    end
  end

  if Clipboard.support_paste_image() and opts.image_paths and #opts.image_paths > 0 then
    local message_content = messages[#messages].content
    for _, image_path in ipairs(opts.image_paths) do
      table.insert(message_content, {
        type = "image",
        source = {
          type = "base64",
          media_type = "image/png",
          data = Clipboard.get_base64_content(image_path),
        },
      })
    end
    messages[#messages].content = message_content
  end

  return messages
end

---@param usage avante.AnthropicTokenUsage | nil
---@return avante.LLMTokenUsage | nil
function M.transform_anthropic_usage(usage)
  if not usage then return nil end
  ---@type avante.LLMTokenUsage
  local res = {
    prompt_tokens = usage.cache_creation_input_tokens and (usage.input_tokens + usage.cache_creation_input_tokens)
      or usage.input_tokens,
    completion_tokens = usage.cache_read_input_tokens and (usage.output_tokens + usage.cache_read_input_tokens)
      or usage.output_tokens,
  }

  return res
end

function M:parse_response(ctx, data_stream, event_state, opts)
  if event_state == nil then
    if data_stream:match('"message_start"') then
      event_state = "message_start"
    elseif data_stream:match('"message_delta"') then
      event_state = "message_delta"
    elseif data_stream:match('"message_stop"') then
      event_state = "message_stop"
    elseif data_stream:match('"content_block_start"') then
      event_state = "content_block_start"
    elseif data_stream:match('"content_block_delta"') then
      event_state = "content_block_delta"
    elseif data_stream:match('"content_block_stop"') then
      event_state = "content_block_stop"
    end
  end
  if ctx.content_blocks == nil then ctx.content_blocks = {} end

  ---@param content AvanteLLMMessageContentItem
  ---@param uuid? string
  ---@return avante.HistoryMessage
  local function new_assistant_message(content, uuid)
    assert(
      event_state == "content_block_start"
        or event_state == "content_block_delta"
        or event_state == "content_block_stop",
      "called with unexpected event_state: " .. event_state
    )
    return HistoryMessage:new("assistant", content, {
      state = event_state == "content_block_stop" and "generated" or "generating",
      turn_id = ctx.turn_id,
      uuid = uuid,
    })
  end

  if event_state == "message_start" then
    local ok, jsn = pcall(vim.json.decode, data_stream)
    if not ok then return end
    ctx.usage = jsn.message.usage
  elseif event_state == "content_block_start" then
    local ok, jsn = pcall(vim.json.decode, data_stream)
    if not ok then return end
    local content_block = jsn.content_block
    content_block.stoppped = false
    ctx.content_blocks[jsn.index + 1] = content_block
    if content_block.type == "text" then
      local msg = new_assistant_message(content_block.text)
      content_block.uuid = msg.uuid
      if opts.on_messages_add then opts.on_messages_add({ msg }) end
    elseif content_block.type == "thinking" then
      if opts.on_chunk then opts.on_chunk("<think>\n") end
      if opts.on_messages_add then
        local msg = new_assistant_message({
          type = "thinking",
          thinking = content_block.thinking,
          signature = content_block.signature,
        })
        content_block.uuid = msg.uuid
        opts.on_messages_add({ msg })
      end
    elseif content_block.type == "tool_use" then
      if opts.on_messages_add then
        local incomplete_json = JsonParser.parse(content_block.input_json)
        local msg = new_assistant_message({
          type = "tool_use",
          name = content_block.name,
          id = content_block.id,
          input = incomplete_json or { dummy = "" },
        })
        content_block.uuid = msg.uuid
        opts.on_messages_add({ msg })
        -- opts.on_stop({ reason = "tool_use", streaming_tool_use = true })
      end
    end
  elseif event_state == "content_block_delta" then
    local ok, jsn = pcall(vim.json.decode, data_stream)
    if not ok then return end
    local content_block = ctx.content_blocks[jsn.index + 1]
    if jsn.delta.type == "input_json_delta" then
      if not content_block.input_json then content_block.input_json = "" end
      content_block.input_json = content_block.input_json .. jsn.delta.partial_json
      return
    elseif jsn.delta.type == "thinking_delta" then
      content_block.thinking = content_block.thinking .. jsn.delta.thinking
      if opts.on_chunk then opts.on_chunk(jsn.delta.thinking) end
      if opts.on_messages_add then
        local msg = new_assistant_message({
          type = "thinking",
          thinking = content_block.thinking,
          signature = content_block.signature,
        }, content_block.uuid)
        opts.on_messages_add({ msg })
      end
    elseif jsn.delta.type == "text_delta" then
      content_block.text = content_block.text .. jsn.delta.text
      if opts.on_chunk then opts.on_chunk(jsn.delta.text) end
      if opts.on_messages_add then
        local msg = new_assistant_message(content_block.text, content_block.uuid)
        opts.on_messages_add({ msg })
      end
    elseif jsn.delta.type == "signature_delta" then
      if ctx.content_blocks[jsn.index + 1].signature == nil then ctx.content_blocks[jsn.index + 1].signature = "" end
      ctx.content_blocks[jsn.index + 1].signature = ctx.content_blocks[jsn.index + 1].signature .. jsn.delta.signature
    end
  elseif event_state == "content_block_stop" then
    local ok, jsn = pcall(vim.json.decode, data_stream)
    if not ok then return end
    local content_block = ctx.content_blocks[jsn.index + 1]
    content_block.stoppped = true
    if content_block.type == "text" then
      if opts.on_messages_add then
        local msg = new_assistant_message(content_block.text, content_block.uuid)
        opts.on_messages_add({ msg })
      end
    elseif content_block.type == "thinking" then
      if opts.on_chunk then
        if content_block.thinking and content_block.thinking ~= vim.NIL and content_block.thinking:sub(-1) ~= "\n" then
          opts.on_chunk("\n</think>\n\n")
        else
          opts.on_chunk("</think>\n\n")
        end
      end
      if opts.on_messages_add then
        local msg = new_assistant_message({
          type = "thinking",
          thinking = content_block.thinking,
          signature = content_block.signature,
        }, content_block.uuid)
        opts.on_messages_add({ msg })
      end
    elseif content_block.type == "tool_use" then
      if opts.on_messages_add then
        local ok_, complete_json = pcall(vim.json.decode, content_block.input_json)
        if not ok_ then complete_json = nil end
        local msg = new_assistant_message({
          type = "tool_use",
          name = content_block.name,
          id = content_block.id,
          input = complete_json or { dummy = "" },
        }, content_block.uuid)
        opts.on_messages_add({ msg })
      end
    end
  elseif event_state == "message_delta" then
    local ok, jsn = pcall(vim.json.decode, data_stream)
    if not ok then return end
    if jsn.usage and ctx.usage then ctx.usage.output_tokens = ctx.usage.output_tokens + jsn.usage.output_tokens end
    if jsn.delta.stop_reason == "end_turn" then
      opts.on_stop({ reason = "complete", usage = M.transform_anthropic_usage(ctx.usage) })
    elseif jsn.delta.stop_reason == "max_tokens" then
      opts.on_stop({ reason = "max_tokens", usage = M.transform_anthropic_usage(ctx.usage) })
    elseif jsn.delta.stop_reason == "tool_use" then
      opts.on_stop({
        reason = "tool_use",
        usage = M.transform_anthropic_usage(ctx.usage),
      })
    end
    return
  elseif event_state == "error" then
    opts.on_stop({ reason = "error", error = vim.json.decode(data_stream) })
  end
end

---@param prompt_opts AvantePromptOptions
---@return AvanteCurlOutput|nil
function M:parse_curl_args(prompt_opts)
  local provider_conf, request_body = P.parse_config(self)
  local disable_tools = provider_conf.disable_tools or false

  local headers = {
    ["Content-Type"] = "application/json",
    ["anthropic-version"] = "2023-06-01",
    ["anthropic-beta"] = "prompt-caching-2024-07-31",
  }

  if P.env.require_api_key(provider_conf) then
    local api_key = self.parse_api_key()
    if not api_key then
      Utils.error("Claude: API key is not set. Please set " .. M.api_key_name)
      return nil
    end
    headers["x-api-key"] = api_key
  end

  local messages = self:parse_messages(prompt_opts)

  local tools = {}
  if not disable_tools and prompt_opts.tools then
    for _, tool in ipairs(prompt_opts.tools) do
      if Config.mode == "agentic" then
        if tool.name == "create_file" then goto continue end
        if tool.name == "view" then goto continue end
        if tool.name == "str_replace" then goto continue end
        if tool.name == "create" then goto continue end
        if tool.name == "insert" then goto continue end
        if tool.name == "undo_edit" then goto continue end
        if tool.name == "replace_in_file" then goto continue end
      end
      table.insert(tools, self:transform_tool(tool))
      ::continue::
    end
  end

  if prompt_opts.tools and #prompt_opts.tools > 0 and Config.mode == "agentic" then
    if provider_conf.model:match("claude%-sonnet%-4%-5") then
      table.insert(tools, {
        type = "text_editor_20250728",
        name = "str_replace_based_edit_tool",
      })
    elseif provider_conf.model:match("claude%-sonnet%-4") then
      table.insert(tools, {
        type = "text_editor_20250429",
        name = "str_replace_based_edit_tool",
      })
    elseif provider_conf.model:match("claude%-3%-7%-sonnet") then
      table.insert(tools, {
        type = "text_editor_20250124",
        name = "str_replace_editor",
      })
    elseif provider_conf.model:match("claude%-3%-5%-sonnet") then
      table.insert(tools, {
        type = "text_editor_20250124",
        name = "str_replace_editor",
      })
    end
  end

  if self.support_prompt_caching then
    if #messages > 0 then
      local found = false
      for i = #messages, 1, -1 do
        local message = messages[i]
        message = vim.deepcopy(message)
        ---@cast message AvanteClaudeMessage
        local content = message.content
        ---@cast content AvanteClaudeMessageContentTextItem[]
        for j = #content, 1, -1 do
          local item = content[j]
          if item.type == "text" then
            item.cache_control = { type = "ephemeral" }
            found = true
            break
          end
        end
        if found then
          messages[i] = message
          break
        end
      end
    end
    if #tools > 0 then
      local last_tool = vim.deepcopy(tools[#tools])
      last_tool.cache_control = { type = "ephemeral" }
      tools[#tools] = last_tool
    end
  end

  return {
    url = Utils.url_join(provider_conf.endpoint, "/v1/messages"),
    proxy = provider_conf.proxy,
    insecure = provider_conf.allow_insecure,
    headers = Utils.tbl_override(headers, self.extra_headers),
    body = vim.tbl_deep_extend("force", {
      model = provider_conf.model,
      system = {
        {
          type = "text",
          text = prompt_opts.system_prompt,
          cache_control = self.support_prompt_caching and { type = "ephemeral" } or nil,
        },
      },
      messages = messages,
      tools = tools,
      stream = true,
    }, request_body),
  }
end

function M.on_error(result)
  if result.status == 429 then return end
  if not result.body then
    return Utils.error("API request failed with status " .. result.status, { once = true, title = "Avante" })
  end

  local ok, body = pcall(vim.json.decode, result.body)
  if not (ok and body and body.error) then
    return Utils.error("Failed to parse error response", { once = true, title = "Avante" })
  end

  local error_msg = body.error.message
  local error_type = body.error.type

  if error_type == "insufficient_quota" then
    error_msg = "You don't have any credits or have exceeded your quota. Please check your plan and billing details."
  elseif error_type == "invalid_request_error" and error_msg:match("temperature") then
    error_msg = "Invalid temperature value. Please ensure it's between 0 and 1."
  end

  Utils.error(error_msg, { once = true, title = "Avante" })
end

return M