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

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)
  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

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

  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

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

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
    if fileinfo.delete then
      vim.api.nvim_buf_set_lines(bufnr, -1, -1, false, {
        string.format("=== Delete file: %s ===", fileinfo.path or "<no path>"),
        ""
      })
    elseif fileinfo.diff then
      vim.api.nvim_buf_set_lines(bufnr, -1, -1, false, {
        string.format("=== Diff for file: %s ===", fileinfo.path or "<no path>")
      })
      local lines = vim.split(fileinfo.diff, "\n")
      for _, line in ipairs(lines) do
        vim.api.nvim_buf_set_lines(bufnr, -1, -1, false, { line })
      end
      vim.api.nvim_buf_set_lines(bufnr, -1, -1, false, { "" })
    elseif fileinfo.content then
      vim.api.nvim_buf_set_lines(bufnr, -1, -1, false, {
        string.format("=== New file: %s ===", fileinfo.path or "<no path>")
      })
      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
      vim.api.nvim_buf_set_lines(bufnr, -1, -1, false, { "" })
    end
  end

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

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
    if fileinfo.delete then
      table.insert(lines, string.format("[DELETE] %s", fileinfo.path or "<no path>"))
    elseif fileinfo.diff then
      table.insert(lines, string.format("[DIFF] %s", fileinfo.path or "<no path>"))
      local diff_lines = vim.split(fileinfo.diff, "\n")
      for _, dl in ipairs(diff_lines) do
        table.insert(lines, "  " .. dl)
      end
    elseif fileinfo.content then
      table.insert(lines, string.format("[WRITE] %s", fileinfo.path or "<no path>"))
      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, delete = false, diff = nil, content = nil }
    local accum = {}

    for _, line in ipairs(edited_lines) do
      if line:match("^#") or line == "" then
        goto continue
      end

      local del_match = line:match("^%[DELETE%] (.+)")
      local diff_match = line:match("^%[DIFF%] (.+)")
      local write_match = line:match("^%[WRITE%] (.+)")

      if del_match or diff_match or write_match then
        -- store previous if any
        if keep_current and (current_fileinfo.path ~= nil) then
          if current_fileinfo.diff then
            current_fileinfo.diff = table.concat(accum, "\n")
          elseif current_fileinfo.content then
            current_fileinfo.content = table.concat(accum, "\n")
          end
          table.insert(final_changes, current_fileinfo)
        end

        accum = {}
        keep_current = true

        if del_match then
          current_fileinfo = { path = del_match, delete = true, diff = nil, content = nil }
        elseif diff_match then
          current_fileinfo = { path = diff_match, delete = false, diff = "", content = nil }
        elseif write_match then
          current_fileinfo = { path = write_match, delete = false, diff = nil, content = "" }
        end

      else
        if keep_current then
          table.insert(accum, line:gsub("^%s*", ""))
        end
      end
      ::continue::
    end

    if keep_current and (current_fileinfo.path ~= nil) then
      if current_fileinfo.diff ~= nil then
        current_fileinfo.diff = table.concat(accum, "\n")
      elseif current_fileinfo.content ~= nil then
        current_fileinfo.content = table.concat(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

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

local function grep_in_file(search_string, filepath)
  local content = read_file(filepath)
  if not content then
    return "Could not read file: " .. filepath
  end
  local results = {}
  local line_num = 0
  for line in content:gmatch("([^\n]*)\n?") do
    line_num = line_num + 1
    if line:find(search_string, 1, true) then
      table.insert(results, filepath .. ":" .. line_num .. ":" .. line)
    end
  end
  if #results == 0 then
    return "No matches in " .. filepath
  else
    return table.concat(results, "\n")
  end
end

local function execute_debug_command(cmd)
  if type(cmd) ~= "table" or not cmd.command then
    return "Invalid command object."
  end

  local command = cmd.command
  if command == "ls" then
    local dir = cmd.dir or "."
    local handle = io.popen("ls " .. dir)
    if not handle then
      return "Failed to run ls command."
    end
    local result = handle:read("*a") or ""
    handle:close()
    return "Listing files in: " .. dir .. "\n" .. result
  elseif command == "grep" then
    local pattern = cmd.pattern
    local target = cmd.target
    if not pattern or not target then
      return "Usage for grep: {command='grep', pattern='<text>', target='<file_or_directory>'}"
    end
    local stat = vim.loop.fs_stat(target)
    if not stat then
      return "Cannot grep: target path does not exist"
    end
    if stat.type == "directory" then
      local handle = io.popen("ls -p " .. target .. " | grep -v /")
      if not handle then
        return "Failed to read directory contents for grep."
      end
      local all_files = {}
      for file in handle:read("*a"):gmatch("[^\n]+") do
        table.insert(all_files, target .. "/" .. file)
      end
      handle:close()
      local results = {}
      for _, f in ipairs(all_files) do
        local fstat = vim.loop.fs_stat(f)
        if fstat and fstat.type == "file" then
          table.insert(results, grep_in_file(pattern, f))
        end
      end
      return table.concat(results, "\n")
    else
      return grep_in_file(pattern, target)
    end
  else
    return "Unknown command: " .. command
  end
end

function M.run_chatgpt_command()
  local conf = config.load()
  ui.debug_log("Running :ChatGPT command.")
  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_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")
      }

      if conf.enable_debug_commands then
        table.insert(initial_sections, "\n### Debug Commands Info:\n")
        table.insert(initial_sections, [[
          If you need debugging commands, include them in your YAML response as follows:

          ```yaml
          commands:
            - command: "list"
              dir: "some/directory"

            - command: "grep"
              pattern: "searchString"
              target: "path/to/file/or/directory"
          ```

          When these commands are present and enable_debug_commands is true, I'll execute them and return the results in the clipboard.
        ]])
      end

      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

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.commands and conf.enable_debug_commands then
    local results = {}
    for _, cmd in ipairs(data.commands) do
      table.insert(results, execute_debug_command(cmd))
    end
    local output = table.concat(results, "\n\n")
    copy_to_clipboard(output)
    print("Debug command results copied to clipboard!")
    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 or fileinfo.diff 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.diff then
          ui.debug_log("Applying diff to file: " .. fileinfo.path)
          local success, err = handler.apply_diff(fileinfo.path, fileinfo.diff)
          if not success then
            vim.api.nvim_err_writeln("Error applying diff: " .. (err or "unknown"))
          else
            print("Patched: " .. fileinfo.path)
          end
        elseif fileinfo.content then
          ui.debug_log("Writing new 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 'diff', 'content', or 'delete'.")
        end
        ::continue::
      end

    else
      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` (use `diff`, `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

function M.run_chatgpt_current_buffer_command()
  local conf = config.load()
  ui.debug_log("Running :ChatGPTCurrentBuffer command.")
  local lines = vim.api.nvim_buf_get_lines(0, 0, -1, false)
  local user_input = table.concat(lines, "\n")
  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")
  }

  if conf.enable_debug_commands then
    table.insert(initial_sections, "\n### Debug Commands Info:\n")
    table.insert(initial_sections, [[
      If you need debugging commands, include them in your YAML response as follows:

      ```yaml
      commands:
        - command: "list"
          dir: "some/directory"

        - command: "grep"
          pattern: "searchString"
          target: "path/to/file/or/directory"
      ```

      When these commands are present and enable_debug_commands is true, I'll execute them and return the results in the clipboard.
    ]])
  end

  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