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

-----------------------------------------------------------------------------
-- CHUNKING
-----------------------------------------------------------------------------
local function handle_chunking_if_needed(prompt, estimate_fn)
  local conf = config.load()
  local token_count = estimate_fn(prompt)
  if not conf.enable_chunking or token_count <= (conf.token_limit or 8000) then
    return { prompt }
  end

  local chunks = ui.chunkify(prompt, estimate_fn, conf.token_limit or 8000)
  return chunks
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)  -- auto-close buffer
    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_chunking_if_needed(prompt, estimate_fn)

      if #chunks == 1 then
        copy_to_clipboard(chunks[1])
        vim.api.nvim_out_write("Prompt copied to clipboard! Paste into ChatGPT.\n")
      else
        -- Multiple chunks. We'll create separate scratch buffers
        for i, chunk in ipairs(chunks) do
          close_existing_buffer_by_name("ChatGPT_Generated_Chunk_" .. i .. "$")

          local cbuf = vim.api.nvim_create_buf(false, true)
          vim.api.nvim_buf_set_name(cbuf, "ChatGPT_Generated_Chunk_" .. i)
          vim.api.nvim_buf_set_option(cbuf, "filetype", "markdown")

          local lines = {
            "# Chunk " .. i .. " of " .. #chunks .. ":",
            "# Copy/paste this chunk into ChatGPT, then come back and copy the next chunk as needed.",
            ""
          }
          vim.list_extend(lines, vim.split(chunk, "\n"))
          vim.api.nvim_buf_set_lines(cbuf, 0, -1, false, lines)

          if i == 1 then
            copy_to_clipboard(chunk)
            vim.api.nvim_out_write("Copied chunk #1 of " .. #chunks .. " to clipboard!\n")
          end
        end
      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
      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) then
        if conf.enable_chunking then
          local chunks = ui.chunkify(prompt, estimate_fn, conf.token_limit or 8000)
          for i, chunk in ipairs(chunks) do
            local cbuf = vim.api.nvim_create_buf(false, true)
            vim.api.nvim_buf_set_name(cbuf, "ChatGPT_Requested_Files_Chunk_" .. i)
            local lines = {
              "# Chunk " .. i .. " of " .. #chunks .. ":",
              "# Copy/paste this chunk into ChatGPT, then come back and copy next chunk as needed.",
              ""
            }
            vim.list_extend(lines, vim.split(chunk, "\n"))
            vim.api.nvim_buf_set_lines(cbuf, 0, -1, false, lines)

            if i == 1 then
              copy_to_clipboard(chunk)
              print("Copied chunk #1 of " .. #chunks .. " to clipboard!")
            end
          end
        else
          vim.api.nvim_err_writeln("Too many requested files. Exceeds token limit.")
        end
        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

return M