local M = {}

local context = require('chatgpt_nvim.context')
local handler = require('chatgpt_nvim.handler')
local config = require('chatgpt_nvim.config')
local ui = require('chatgpt_nvim.ui')

local ok_yaml, lyaml = pcall(require, "lyaml")

local function copy_to_clipboard(text)
  vim.fn.setreg('+', text)
end

local function parse_response(raw)
  local conf = config.load()
  if not ok_yaml then
    vim.api.nvim_err_writeln("lyaml not available. Install with `luarocks install lyaml`.")
    return nil
  end
  local ok, data = pcall(lyaml.load, raw)
  if not ok or not data then
    vim.api.nvim_err_writeln("Failed to parse YAML response.")
    ui.debug_log("RAW response that failed parsing:\n" .. raw)
    return nil
  end
  ui.debug_log("Successfully parsed YAML response.")
  return data
end

local function is_subpath(root, path)
  local root_abs = vim.fn.fnamemodify(root, ":p")
  local target_abs = vim.fn.fnamemodify(path, ":p")
  return target_abs:sub(1, #root_abs) == root_abs
end

local function read_file(path)
  local fd = vim.loop.fs_open(path, "r", 438)
  if not fd then
    return nil
  end
  local stat = vim.loop.fs_fstat(fd)
  if not stat then
    vim.loop.fs_close(fd)
    return nil
  end
  local data = vim.loop.fs_read(fd, stat.size, 0)
  vim.loop.fs_close(fd)
  return data
end

local function is_directory(path)
  local stat = vim.loop.fs_stat(path)
  return stat and stat.type == "directory"
end

-- We'll create two token estimate functions: one basic, one improved
local function estimate_tokens_basic(text)
  local approx_chars_per_token = 4
  local length = #text
  return math.floor(length / approx_chars_per_token)
end

local function estimate_tokens_improved(text)
  -- Word-based approach, assume ~0.75 token/word
  local words = #vim.split(text, "%s+")
  local approximate_tokens = math.floor(words * 0.75)
  ui.debug_log("Using improved token estimate: " .. approximate_tokens .. " tokens")
  return approximate_tokens
end

local function get_estimate_fn()
  local conf = config.load()
  if conf.improved_debug then
    return estimate_tokens_improved
  else
    return estimate_tokens_basic
  end
end

---------------------------------------------------------------------------------
-- Step-by-Step Handling (replaces chunking)
---------------------------------------------------------------------------------
local function handle_step_by_step_if_needed(prompt, estimate_fn)
  local conf = config.load()
  local token_count = estimate_fn(prompt)
  -- If step-by-step is disabled or token count is within limit, return the original prompt
  if not conf.enable_step_by_step or token_count <= (conf.token_limit or 8000) then
    return { prompt }
  end

  -- If we exceed the token limit, create a single message prompting the user to split tasks
  local step_prompt = [[
    It appears this request might exceed the model's token limit if done all at once.
    Please break down the tasks into smaller steps and handle them one by one.
    At each step, we'll provide relevant files or context if needed.
    Thank you!
  ]]

  return { step_prompt }
end

-- Close an existing buffer by name (if it exists)
local function close_existing_buffer_by_name(pattern)
  for _, b in ipairs(vim.api.nvim_list_bufs()) do
    local name = vim.api.nvim_buf_get_name(b)
    if name:match(pattern) then
      vim.api.nvim_buf_delete(b, { force = true })
    end
  end
end

-- Show the user a preview buffer with the proposed changes (unchanged).
local function preview_changes(changes)
  close_existing_buffer_by_name("ChatGPT_Changes_Preview$")

  local bufnr = vim.api.nvim_create_buf(false, true)
  vim.api.nvim_buf_set_name(bufnr, "ChatGPT_Changes_Preview")
  vim.api.nvim_buf_set_option(bufnr, "filetype", "diff")
  vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, {
    "# Preview of Changes:",
    "# (Close this window to apply changes or use :q to cancel)",
    ""
  })
  for _, fileinfo in ipairs(changes) do
    local indicator = (fileinfo.delete == true) and "Delete file" or "Write file"
    vim.api.nvim_buf_set_lines(bufnr, -1, -1, false, {
      string.format("=== %s: %s ===", indicator, fileinfo.path or "<no path>")
    })
    if fileinfo.content then
      local lines = vim.split(fileinfo.content, "\n")
      for _, line in ipairs(lines) do
        vim.api.nvim_buf_set_lines(bufnr, -1, -1, false, { line })
      end
    end
    vim.api.nvim_buf_set_lines(bufnr, -1, -1, false, { "" })
  end

  vim.cmd("vsplit")
  vim.cmd("buffer " .. bufnr)
end

-- Minimal partial acceptance from previous example
local function partial_accept(changes)
  close_existing_buffer_by_name("ChatGPT_Partial_Accept$")

  local bufnr = vim.api.nvim_create_buf(false, true)
  vim.api.nvim_buf_set_name(bufnr, "ChatGPT_Partial_Accept")
  vim.api.nvim_buf_set_option(bufnr, "filetype", "diff")

  local lines = {
    "# Remove or comment out (prepend '#') any changes you do NOT want, then :wq, :x, or :bd to finalize partial acceptance",
    ""
  }
  for _, fileinfo in ipairs(changes) do
    local action = (fileinfo.delete == true) and "[DELETE]" or "[WRITE]"
    table.insert(lines, string.format("%s %s", action, fileinfo.path or "<no path>"))
    if fileinfo.content then
      local content_lines = vim.split(fileinfo.content, "\n")
      for _, cl in ipairs(content_lines) do
        table.insert(lines, "  " .. cl)
      end
    end
    table.insert(lines, "")
  end

  vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, lines)

  local final_changes = {}
  local function on_write()
    local edited_lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
    local keep_current = false
    local current_fileinfo = { path = nil, content = nil, delete = false }
    local content_accum = {}

    for _, line in ipairs(edited_lines) do
      if line:match("^#") or line == "" then
        goto continue
      end
      local del_match = line:match("^%[DELETE%] (.+)")
      local write_match = line:match("^%[WRITE%] (.+)")
      if del_match then
        if keep_current and (current_fileinfo.path ~= nil) then
          if #content_accum > 0 then
            current_fileinfo.content = table.concat(content_accum, "\n")
          end
          table.insert(final_changes, current_fileinfo)
        end
        keep_current = true
        current_fileinfo = { path = del_match, delete = true, content = nil }
        content_accum = {}
      elseif write_match then
        if keep_current and (current_fileinfo.path ~= nil) then
          if #content_accum > 0 then
            current_fileinfo.content = table.concat(content_accum, "\n")
          end
          table.insert(final_changes, current_fileinfo)
        end
        keep_current = true
        current_fileinfo = { path = write_match, delete = false, content = nil }
        content_accum = {}
      else
        if keep_current then
          table.insert(content_accum, line:gsub("^%s*", ""))
        end
      end
      ::continue::
    end

    if keep_current and (current_fileinfo.path ~= nil) then
      if #content_accum > 0 then
        current_fileinfo.content = table.concat(content_accum, "\n")
      end
      table.insert(final_changes, current_fileinfo)
    end

    vim.api.nvim_buf_set_option(bufnr, "modified", false)
  end

  vim.api.nvim_create_autocmd("BufWriteCmd", {
    buffer = bufnr,
    once = true,
    callback = function()
      on_write()
      vim.cmd("bd! " .. bufnr)
    end
  })

  vim.cmd("split")
  vim.cmd("buffer " .. bufnr)

  vim.wait(60000, function()
    local winids = vim.api.nvim_tabpage_list_wins(0)
    for _, w in ipairs(winids) do
      local b = vim.api.nvim_win_get_buf(w)
      if b == bufnr then
        return false
      end
    end
    return true
  end)

  return final_changes
end

-- Utility to store generated prompt in a scratch buffer
local function store_prompt_for_reference(prompt)
  close_existing_buffer_by_name("ChatGPT_Generated_Prompt$")

  local bufnr = vim.api.nvim_create_buf(false, true)
  vim.api.nvim_buf_set_name(bufnr, "ChatGPT_Generated_Prompt")
  vim.api.nvim_buf_set_option(bufnr, "filetype", "markdown")

  local lines = {
    "# Below is the generated prompt. You can keep it for reference:",
    ""
  }
  local prompt_lines = vim.split(prompt, "\n")
  for _, line in ipairs(prompt_lines) do
    table.insert(lines, line)
  end

  vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, lines)
  vim.cmd("vsplit")
  vim.cmd("buffer " .. bufnr)
end

----------------------------------------------------------------------------
-- run_chatgpt_command
----------------------------------------------------------------------------
function M.run_chatgpt_command()
  local conf = config.load()
  ui.debug_log("Running :ChatGPT command.")

  -- Possibly let user select directories if interactive_file_selection is true
  local dirs = conf.directories or {"."}
  if conf.interactive_file_selection then
    dirs = ui.pick_directories(dirs)
    if #dirs == 0 then
      dirs = conf.directories
    end
  end

  -- Close existing prompt buffer if open
  close_existing_buffer_by_name("ChatGPT_Prompt.md$")

  local bufnr = vim.api.nvim_create_buf(false, false)
  vim.api.nvim_buf_set_name(bufnr, "ChatGPT_Prompt.md")
  vim.api.nvim_buf_set_option(bufnr, "filetype", "markdown")
  vim.api.nvim_buf_set_option(bufnr, "bufhidden", "wipe")
  vim.api.nvim_buf_set_option(bufnr, "buftype", "")
  vim.api.nvim_buf_set_option(bufnr, "modifiable", true)

  vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, {
    "# Enter your prompt below.",
    "",
    "Save & close with :wq, :x, or :bd to finalize your prompt."
  })

  vim.api.nvim_create_autocmd("BufWriteCmd", {
    buffer = bufnr,
    callback = function()
      local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
      local user_input = table.concat(lines, "\n")
      if user_input == "" or user_input:find("^# Enter your prompt below.") then
        vim.api.nvim_out_write("No valid input provided.\n")
        vim.api.nvim_buf_set_option(bufnr, "modified", false)
        return
      end

      local project_structure = context.get_project_structure(dirs)
      local initial_files = conf.initial_files or {}
      local included_sections = {}

      for _, item in ipairs(initial_files) do
        local root = vim.fn.getcwd()
        local full_path = root .. "/" .. item
        if is_directory(full_path) then
          local dir_files = context.get_project_files({item})
          for _, f in ipairs(dir_files) do
            local path = root .. "/" .. f
            local data = read_file(path)
            if data then
              table.insert(included_sections, "\nFile: `" .. f .. "`\n```\n" .. data .. "\n```\n")
            end
          end
        else
          local data = read_file(full_path)
          if data then
            table.insert(included_sections, "\nFile: `" .. item .. "`\n```\n" .. data .. "\n```\n")
          end
        end
      end

      local initial_sections = {
        "### Basic Prompt Instructions:\n",
        conf.initial_prompt .. "\n\n\n",
        "### User Instructions:\n",
        user_input .. "\n\n\n",
        "### Context/Data:\n",
        "Project name: " .. (conf.project_name or "") .. "\n",
        "Project Structure:\n",
        project_structure,
        table.concat(included_sections, "\n")
      }

      local prompt = table.concat(initial_sections, "\n")
      store_prompt_for_reference(prompt)

      local estimate_fn = get_estimate_fn()
      local chunks = handle_step_by_step_if_needed(prompt, estimate_fn)

      copy_to_clipboard(chunks[1])
      if #chunks == 1 then
        vim.api.nvim_out_write("Prompt copied to clipboard! Paste into ChatGPT.\n")
      else
        vim.api.nvim_out_write("Step-by-step prompt copied to clipboard!\n")
      end

      vim.api.nvim_buf_set_option(bufnr, "modified", false)
    end
  })

  vim.cmd("buffer " .. bufnr)
end

----------------------------------------------------------------------------
-- run_chatgpt_paste_command
----------------------------------------------------------------------------
function M.run_chatgpt_paste_command()
  local conf = config.load()
  ui.debug_log("Running :ChatGPTPaste command.")
  print("Reading ChatGPT YAML response from clipboard...")
  local raw = handler.get_clipboard_content()
  if raw == "" then
    vim.api.nvim_err_writeln("Clipboard is empty. Please copy the YAML response from ChatGPT first.")
    return
  end

  local data = parse_response(raw)
  if not data then
    return
  end

  if data.project_name and data.files then
    ui.debug_log("Received project_name and files in response.")
    if data.project_name ~= conf.project_name then
      vim.api.nvim_err_writeln("Project name mismatch. Aborting.")
      return
    end

    local is_final = false
    for _, fileinfo in ipairs(data.files) do
      if fileinfo.content or fileinfo.delete == true then
        is_final = true
        break
      end
    end

    if is_final then
      if conf.preview_changes then
        preview_changes(data.files)
        print("Close the preview window to apply changes, or use :q to cancel.")
        local closed = vim.wait(60000, function()
          local bufs = vim.api.nvim_list_bufs()
          for _, b in ipairs(bufs) do
            local name = vim.api.nvim_buf_get_name(b)
            if name:match("ChatGPT_Changes_Preview$") then
              return false
            end
          end
          return true
        end)
        if not closed then
          vim.api.nvim_err_writeln("Preview not closed in time. Aborting.")
          return
        end

      local final_files = data.files
      if conf.partial_acceptance then
        final_files = partial_accept(data.files)
        if #final_files == 0 then
          vim.api.nvim_err_writeln("No changes remain after partial acceptance. Aborting.")
          return
        end
      end

      local root = vim.fn.getcwd()
      for _, fileinfo in ipairs(final_files) do
        if not fileinfo.path then
          vim.api.nvim_err_writeln("Invalid file entry. Must have 'path'.")
          goto continue
        end

        if not is_subpath(root, fileinfo.path) then
          vim.api.nvim_err_writeln("Invalid path outside project root: " .. fileinfo.path)
          goto continue
        end

        if fileinfo.delete == true then
          ui.debug_log("Deleting file: " .. fileinfo.path)
          handler.delete_file(fileinfo.path)
          print("Deleted: " .. fileinfo.path)
        elseif fileinfo.content then
          ui.debug_log("Writing file: " .. fileinfo.path)
          handler.write_file(fileinfo.path, fileinfo.content)
          print("Wrote: " .. fileinfo.path)
        else
          vim.api.nvim_err_writeln("Invalid file entry. Must have 'content' or 'delete'.")
        end
        ::continue::
      end

    else
      -- Intermediate request for more files
      local requested_paths = {}
      for _, fileinfo in ipairs(data.files) do
        if fileinfo.path then
          table.insert(requested_paths, fileinfo.path)
        end
      end

      local file_sections = {}
      local root = vim.fn.getcwd()
      for _, f in ipairs(requested_paths) do
        local path = root .. "/" .. f
        local content = read_file(path)
        if content then
          table.insert(file_sections, "\nFile: `" .. f .. "`\n```\n" .. content .. "\n```\n")
        else
          table.insert(file_sections, "\nFile: `" .. f .. "`\n```\n(Could not read file)\n```\n")
        end
      end

      local sections = {
        config.load().initial_prompt,
        "\n\nProject name: " .. (conf.project_name or ""),
        "\n\nBelow are the requested files from the project, each preceded by its filename in backticks and enclosed in triple backticks.\n",
        table.concat(file_sections, "\n"),
        "\n\nIf you need more files, please respond again in YAML listing additional files. If you have all information you need, provide the final YAML with `project_name` and `files` (with `content` or `delete`) to apply changes.\n"
      }

      local prompt = table.concat(sections, "\n")
      local estimate_fn = get_estimate_fn()

      local token_count = estimate_fn(prompt)
      ui.debug_log("Returning requested files. Token count: " .. token_count)

      if token_count > (conf.token_limit or 8000) and conf.enable_step_by_step then
        local step_message = [[
          It appears this requested data is quite large. Please split the task into smaller steps
          and continue step by step.
          Which files would you need for the first step?
        ]]
        copy_to_clipboard(step_message)
        print("Step-by-step guidance copied to clipboard!")
        return
      elseif token_count > (conf.token_limit or 8000) then
        vim.api.nvim_err_writeln("Requested files exceed token limit. No step-by-step support enabled.")
        return
      end

      copy_to_clipboard(prompt)
      print("Prompt (with requested files) copied to clipboard! Paste it into the ChatGPT O1 model.")
    end
  else
    vim.api.nvim_err_writeln("Invalid response. Expected 'project_name' and 'files'.")
  end
end

----------------------------------------------------------------------------
-- run_chatgpt_current_buffer_command
----------------------------------------------------------------------------
function M.run_chatgpt_current_buffer_command()
  local conf = config.load()
  ui.debug_log("Running :ChatGPTCurrentBuffer command.")

  -- Get the content of the current buffer as user instructions
  local lines = vim.api.nvim_buf_get_lines(0, 0, -1, false)
  local user_input = table.concat(lines, "\n")

  -- Possibly let user select directories if interactive_file_selection is true
  local dirs = conf.directories or {"."}
  if conf.interactive_file_selection then
    dirs = ui.pick_directories(dirs)
    if #dirs == 0 then
      dirs = conf.directories
    end
  end

  local project_structure = context.get_project_structure(dirs)
  local initial_files = conf.initial_files or {}
  local included_sections = {}

  for _, item in ipairs(initial_files) do
    local root = vim.fn.getcwd()
    local full_path = root .. "/" .. item
    if is_directory(full_path) then
      local dir_files = context.get_project_files({item})
      for _, f in ipairs(dir_files) do
        local path = root .. "/" .. f
        local data = read_file(path)
        if data then
          table.insert(included_sections, "\nFile: `" .. f .. "`\n```\n" .. data .. "\n```\n")
        end
      end
    else
      local data = read_file(full_path)
      if data then
        table.insert(included_sections, "\nFile: `" .. item .. "`\n```\n" .. data .. "\n```\n")
      end
    end
  end

  local initial_sections = {
    "### Basic Prompt Instructions:\n",
    conf.initial_prompt .. "\n\n\n",
    "### User Instructions:\n",
    user_input .. "\n\n\n",
    "### Context/Data:\n",
    "Project name: " .. (conf.project_name or "") .. "\n",
    "Project Structure:\n",
    project_structure,
    table.concat(included_sections, "\n")
  }

  local prompt = table.concat(initial_sections, "\n")
  store_prompt_for_reference(prompt)

  local estimate_fn = get_estimate_fn()
  local chunks = handle_step_by_step_if_needed(prompt, estimate_fn)

  copy_to_clipboard(chunks[1])
  if #chunks == 1 then
    vim.api.nvim_out_write("Prompt (from current buffer) copied to clipboard! Paste into ChatGPT.\n")
  else
    vim.api.nvim_out_write("Step-by-step prompt (from current buffer) copied to clipboard!\n")
  end
end

return M