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 prompts = require('chatgpt_nvim.prompts')

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

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

local function parse_response(raw, conf)
  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 handle_step_by_step_if_needed(prompt, conf)
  local length = #prompt
  if not conf.enable_step_by_step or length <= (conf.prompt_char_limit or 8000) then
    return { prompt }
  end
  return { prompts["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, conf)
  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

local function partial_accept(changes, conf)
  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

local function store_prompt_for_reference(prompt, conf)
  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

-- Updated to allow flags/args for 'ls' and 'grep'
local function execute_debug_command(cmd, conf)
  if type(cmd) ~= "table" or not cmd.command then
    return "Invalid command object."
  end

  local command = cmd.command
  if command == "ls" then
    -- Accept optional flags/args
    local dir = cmd.dir or "."
    local args = cmd.args or {}
    local cmd_str = "ls"
    if #args > 0 then
      cmd_str = cmd_str .. " " .. table.concat(args, " ")
    end
    cmd_str = cmd_str .. " " .. dir

    local handle = io.popen(cmd_str)
    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
    -- Accept optional flags/args
    -- If the user wants to do something like:
    --   args: ["-r", "somePattern", "someDir"]
    -- we just pass them all to grep.
    local args = cmd.args or {}
    if #args == 0 then
      -- fallback to old usage
      local pattern = cmd.pattern
      local target = cmd.target
      if not pattern or not target then
        return "Usage for grep: {command='grep', args=['-r','pattern','target']} or {pattern='<text>', target='<path>'}"
      end
      local stat = vim.loop.fs_stat(target)
      if not stat then
        return "Cannot grep: target path does not exist"
      end
      -- old logic remains for backward compatibility
      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 = {}
        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 lines = {}
          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(lines, filepath .. ":" .. line_num .. ":" .. line)
            end
          end
          return (#lines == 0) and ("No matches in " .. filepath) or table.concat(lines, "\n")
        end
        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
        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 lines = {}
          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(lines, filepath .. ":" .. line_num .. ":" .. line)
            end
          end
          return (#lines == 0) and ("No matches in " .. filepath) or table.concat(lines, "\n")
        end
        return grep_in_file(pattern, target)
      end
    else
      -- new approach with flags/args
      local cmd_str = "grep " .. table.concat(args, " ")
      local handle = io.popen(cmd_str)
      if not handle then
        return "Failed to run grep command."
      end
      local result = handle:read("*a") or ""
      handle:close()
      return result
    end
  else
    return "Unknown command: " .. command
  end
end

local function run_chatgpt_command()
  local conf = config.load()
  ui.setup_ui(conf)
  ui.debug_log("Running :ChatGPT command.")
  local dirs = conf.directories or {"."}
  if conf.interactive_file_selection then
    dirs = ui.pick_directories(dirs, conf)
    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, conf)
      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}, conf)
          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, prompts["debug-commands-info"])
      end

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

      local chunks = handle_step_by_step_if_needed(prompt, conf)
      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

local function run_chatgpt_paste_command()
  local conf = config.load()
  ui.setup_ui(conf)
  ui.debug_log("Running :ChatGPTPaste command.")
  print("Reading ChatGPT YAML response from clipboard...")
  local raw = handler.get_clipboard_content(conf)
  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, conf)
  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, conf))
    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 then
        is_final = true
        break
      end
    end

    if is_final then
      if conf.preview_changes then
        preview_changes(data.files, conf)
        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, conf)
        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, conf)
          print("Deleted: " .. fileinfo.path)
        elseif fileinfo.content then
          ui.debug_log("Writing file: " .. fileinfo.path)
          handler.write_file(fileinfo.path, fileinfo.content, conf)
          print("Wrote: " .. fileinfo.path)
        else
          vim.api.nvim_err_writeln("Invalid file entry. Must have '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 = {
        conf.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 length = #prompt
      ui.debug_log("Returning requested files. Character count: " .. length)

      if length > (conf.prompt_char_limit or 8000) and conf.enable_step_by_step then
        local large_step = prompts["step-prompt"]
        copy_to_clipboard(large_step)
        print("Step-by-step guidance copied to clipboard!")
        return
      elseif length > (conf.prompt_char_limit or 8000) then
        vim.api.nvim_err_writeln("Requested files exceed prompt character 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

local function run_chatgpt_current_buffer_command()
  local conf = config.load()
  ui.setup_ui(conf)
  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, conf)
    if #dirs == 0 then
      dirs = conf.directories
    end
  end

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

  local function is_directory(path)
    local stat = vim.loop.fs_stat(path)
    return stat and stat.type == "directory"
  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

  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}, conf)
      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, prompts["debug-commands-info"])
  end

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

  local function store_prompt_for_reference(pr)
    close_existing_buffer_by_name("ChatGPT_Generated_Prompt$")
    local bufnr_ref = vim.api.nvim_create_buf(false, true)
    vim.api.nvim_buf_set_name(bufnr_ref, "ChatGPT_Generated_Prompt")
    vim.api.nvim_buf_set_option(bufnr_ref, "filetype", "markdown")

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

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

  store_prompt_for_reference(prompt, conf)

  local chunks = handle_step_by_step_if_needed(prompt, conf)
  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

M.run_chatgpt_command = run_chatgpt_command
M.run_chatgpt_paste_command = run_chatgpt_paste_command
M.run_chatgpt_current_buffer_command = run_chatgpt_current_buffer_command
M.execute_debug_command = execute_debug_command

return M