require "../protocol"

module Lattice
  # Generates LSP code actions (quick-fixes and source actions) derived from
  # inspection diagnostics and cursor context for a document.
  #
  # The provider maps the `source` and message prefix of each diagnostic to a
  # concrete `WorkspaceEdit`-backed `::LSP::CodeAction`.  Only diagnostics whose
  # source tag is `"lattice"` are considered; Ameba diagnostics (source
  # `"ameba"`) are passed through as-is without a quick-fix because the
  # corrective edit is ameba's responsibility.
  #
  # Supported quick-fixes
  #
  # - **Unused variable** — removes the assignment line.
  # - **Unused parameter** — prepends `_` to the parameter name to suppress
  #   the warning conventionally.
  # - **Missing return** — inserts a bare `return` at the end of the method body.
  # - **Wrap nilable** — replaces `foo.not_nil!` with `if val = foo; ...; end`.
  # - **If to modifier** — converts `if x\n  y\nend` to `y if x`.
  # - **Add missing require** — prepends `require "foo"` when a diagnostic
  #   names an unknown constant.
  # - **Sort requires** — groups and sorts all `require` lines in the file.
  class CodeActionsProvider
    # Produce code actions for *uri* given its *source* text, the
    # *diagnostics* already computed for that document, and the cursor *range*
    # (used for context-sensitive actions such as nil-check rewriting and
    # single-statement if-to-modifier conversion).
    def actions(
      uri : String,
      source : String,
      diagnostics : Array(::LSP::Diagnostic),
      range : ::LSP::Range? = nil,
    ) : Array(::LSP::CodeAction)
      lines = source.split('\n')
      result = [] of ::LSP::CodeAction

      diagnostics.each do |diag|
        next unless diag.source == "lattice"

        message = diag.message

        if message.starts_with?("Unused variable '")
          if action = remove_unused_variable_action(uri, lines, diag)
            result << action
          end
        elsif message.starts_with?("Unused parameter '")
          if action = prefix_unused_parameter_action(uri, lines, diag)
            result << action
          end
        elsif message.starts_with?("Method '") && message.includes?("may not return a value")
          if action = add_missing_return_action(uri, lines, diag)
            result << action
          end
        end

        # "Add missing require" fires on any diagnostic naming an unknown constant.
        if action = add_missing_require_action(uri, lines, diag)
          result << action
        end
      end

      # Context-sensitive actions that depend on cursor position.
      if r = range
        cursor_line = r.start.line
        if cursor_line < lines.size
          if action = wrap_not_nil_action(uri, lines, cursor_line)
            result << action
          end
          if action = if_to_modifier_action(uri, lines, cursor_line)
            result << action
          end
        end
      end

      # Source actions offered on any require line.
      if r = range
        cursor_line = r.start.line
        if cursor_line < lines.size && lines[cursor_line].strip.starts_with?("require ")
          if action = sort_requires_action(uri, lines)
            result << action
          end
        end
      end

      result
    end

    # ---------------------------------------------------------------------------
    # Quick-fix: replace `expr.not_nil!` with a safe `if val = expr` check
    # ---------------------------------------------------------------------------

    private def wrap_not_nil_action(
      uri : String,
      lines : Array(String),
      cursor_line : Int32,
    ) : ::LSP::CodeAction?
      line = lines[cursor_line]

      # Match `something.not_nil!` anywhere on the line; capture the receiver.
      m = line.match(/\b(\w[\w.\[\]()]*)\s*\.not_nil!/) ||
          line.match(/(\w[\w.]*)\s*\.not_nil!/)
      return unless m

      receiver  = m[1]
      full_call = "#{receiver}.not_nil!"
      col_start = line.index(full_call) || 0
      col_end   = col_start + full_call.size

      indent = line[/^\s*/]
      inner  = "#{indent}  # TODO: use val\n"

      replacement = "if val = #{receiver}\n#{inner}#{indent}end"

      edit = ::LSP::TextEdit.new(
        ::LSP::Range.new(
          ::LSP::Position.new(cursor_line, col_start),
          ::LSP::Position.new(cursor_line, col_end),
        ),
        replacement,
      )
      workspace_edit = ::LSP::WorkspaceEdit.new({uri => [edit]})

      ::LSP::CodeAction.new(
        title: "Replace .not_nil! with safe nil check",
        kind: "quickfix",
        edit: workspace_edit,
        is_preferred: false,
      )
    end

    # ---------------------------------------------------------------------------
    # Quick-fix: convert single-statement `if x\n  y\nend` to `y if x`
    # ---------------------------------------------------------------------------

    private def if_to_modifier_action(
      uri : String,
      lines : Array(String),
      cursor_line : Int32,
    ) : ::LSP::CodeAction?
      line = lines[cursor_line]
      stripped = line.strip

      # Only act on lines that start with `if ` and have no `else`/`elsif`.
      return unless m = stripped.match(/^if\s+(.+)$/)
      condition = m[1].strip
      return if condition.empty?

      indent = line[/^\s*/]

      # The next line should be the single body statement.
      body_line_num = cursor_line + 1
      return if body_line_num >= lines.size
      body_line = lines[body_line_num]
      body_stripped = body_line.strip
      return if body_stripped.empty? || body_stripped.starts_with?('#')

      # The line after that should be `end` at the same indent level.
      end_line_num = cursor_line + 2
      return if end_line_num >= lines.size
      end_line = lines[end_line_num]
      return unless end_line.strip == "end" && end_line[/^\s*/] == indent

      # There must be no else/elsif between if and end.
      return if body_stripped.starts_with?("else") ||
                body_stripped.starts_with?("elsif")

      modifier_form = "#{indent}#{body_stripped} if #{condition}"

      # Replace the three lines (if / body / end) with the modifier form.
      delete_range = ::LSP::Range.new(
        ::LSP::Position.new(cursor_line, 0),
        ::LSP::Position.new(end_line_num + 1, 0),
      )
      edit = ::LSP::TextEdit.new(delete_range, "#{modifier_form}\n")
      workspace_edit = ::LSP::WorkspaceEdit.new({uri => [edit]})

      ::LSP::CodeAction.new(
        title: "Convert to modifier form: `#{body_stripped} if #{condition}`",
        kind: "refactor.rewrite",
        edit: workspace_edit,
        is_preferred: false,
      )
    end

    # ---------------------------------------------------------------------------
    # Quick-fix: prepend `require "foo"` for unknown constants named in diagnostics
    # ---------------------------------------------------------------------------

    private def add_missing_require_action(
      uri : String,
      lines : Array(String),
      diag : ::LSP::Diagnostic,
    ) : ::LSP::CodeAction?
      # Crystal compiler error messages that name an undefined constant look like:
      #   "undefined constant Foo"
      #   "undefined constant Foo::Bar"
      # We support both the compiler source and our own source tag.
      msg = diag.message
      m = msg.match(/undefined constant ([A-Z][\w:]*)/)
      return unless m

      const_name = m[1]

      # Derive a require path guess: split on `::`, downcase each segment.
      # e.g. "JSON::Any" → "json", "MyApp::Utils" → "my_app/utils"
      segments = const_name.split("::")
      path_guess = segments.map { |s|
        # CamelCase → snake_case
        s.gsub(/([A-Z]+)([A-Z][a-z])/, "\\1_\\2")
         .gsub(/([a-z\d])([A-Z])/, "\\1_\\2")
         .downcase
      }.join("/")

      # Use only the first segment as the require path when it is a
      # well-known stdlib or shard root (avoids "json/any" for JSON::Any).
      require_path = segments.size == 1 ? path_guess : path_guess.split("/").first

      require_stmt = "require \"#{require_path}\""

      # Do not offer the action if the require is already present.
      return if lines.any? { |l| l.strip == require_stmt }

      # Insert after the last existing require block (or at line 0).
      insert_line = last_require_line(lines) + 1

      insert_pos = ::LSP::Position.new(insert_line, 0)
      edit = ::LSP::TextEdit.new(
        ::LSP::Range.new(insert_pos, insert_pos),
        "#{require_stmt}\n",
      )
      workspace_edit = ::LSP::WorkspaceEdit.new({uri => [edit]})

      ::LSP::CodeAction.new(
        title: "Add #{require_stmt}",
        kind: "quickfix",
        diagnostics: [diag],
        edit: workspace_edit,
        is_preferred: false,
      )
    end

    # ---------------------------------------------------------------------------
    # Source action: sort and group all `require` statements in the file
    # ---------------------------------------------------------------------------

    # Require grouping order:
    #   1. stdlib requires  — bare word, no `/`, no `.`  (e.g. `require "json"`)
    #   2. shard requires   — bare word possibly with `/` but no leading `.`
    #   3. relative requires — starts with `./` or `../`
    # Groups are separated by a single blank line.
    private def sort_requires_action(
      uri : String,
      lines : Array(String),
    ) : ::LSP::CodeAction?
      # Collect all require lines with their line numbers.
      require_lines = [] of {Int32, String}
      lines.each_with_index do |line, idx|
        require_lines << {idx, line} if line.strip.starts_with?("require ")
      end
      return if require_lines.size < 2

      # Determine the contiguous block of lines to replace.  We expand the
      # selection to include blank lines that sit only between require lines,
      # so the sort result can fill in the canonical grouping blank lines.
      first_line = require_lines.first[0]
      last_line  = require_lines.last[0]

      # Check that every line in [first_line, last_line] is either a require
      # or blank — otherwise we cannot safely reorder.
      (first_line..last_line).each do |idx|
        l = lines[idx].strip
        return if !l.empty? && !l.starts_with?("require ")
      end

      # Classify and sort within each group.
      stdlib   = [] of String
      shards   = [] of String
      relative = [] of String

      require_lines.each do |_, raw|
        stripped = raw.strip
        if m = stripped.match(/require\s+["']([^"']+)["']/)
          path = m[1]
          if path.starts_with?(".")
            relative << stripped
          elsif !path.includes?("/") && !path.includes?(".")
            # Heuristic: single-segment bare word → stdlib
            stdlib << stripped
          else
            shards << stripped
          end
        end
      end

      stdlib.sort!
      shards.sort!
      relative.sort!

      groups = [stdlib, shards, relative].reject(&.empty?)
      sorted_block = groups.map(&.join("\n")).join("\n\n") + "\n"

      # Replacement range covers first_line through the line after last_line.
      replace_range = ::LSP::Range.new(
        ::LSP::Position.new(first_line, 0),
        ::LSP::Position.new(last_line + 1, 0),
      )
      edit = ::LSP::TextEdit.new(replace_range, sorted_block)
      workspace_edit = ::LSP::WorkspaceEdit.new({uri => [edit]})

      ::LSP::CodeAction.new(
        title: "Sort require statements",
        kind: "source",
        edit: workspace_edit,
        is_preferred: false,
      )
    end

    # Return the 0-based index of the last `require` line, or -1 if none.
    private def last_require_line(lines : Array(String)) : Int32
      idx = -1
      lines.each_with_index do |line, i|
        idx = i if line.strip.starts_with?("require ")
      end
      idx
    end

    private def remove_unused_variable_action(
      uri : String,
      lines : Array(String),
      diag : ::LSP::Diagnostic,
    ) : ::LSP::CodeAction?
      line_num = diag.range.start.line
      return if line_num >= lines.size

      # Delete the entire line (including the trailing newline).
      delete_range = ::LSP::Range.new(
        ::LSP::Position.new(line_num, 0),
        ::LSP::Position.new(line_num + 1, 0),
      )
      edit = ::LSP::TextEdit.new(delete_range, "")
      workspace_edit = ::LSP::WorkspaceEdit.new({uri => [edit]})

      ::LSP::CodeAction.new(
        title: "Remove unused variable",
        kind: "quickfix",
        diagnostics: [diag],
        edit: workspace_edit,
        is_preferred: true,
      )
    end

    # ---------------------------------------------------------------------------
    # Quick-fix: prefix unused parameter with underscore
    # ---------------------------------------------------------------------------

    private def prefix_unused_parameter_action(
      uri : String,
      lines : Array(String),
      diag : ::LSP::Diagnostic,
    ) : ::LSP::CodeAction?
      line_num = diag.range.start.line
      col = diag.range.start.character
      return if line_num >= lines.size

      line = lines[line_num]
      return if col >= line.size

      # Extract the parameter name from the diagnostic message.
      if m = diag.message.match(/Parameter '(\w+)'/)
        param_name = m[1]

        # Find the exact column of the name in the line.
        start_col = line.index(param_name, col) || col
        range = ::LSP::Range.new(
          ::LSP::Position.new(line_num, start_col),
          ::LSP::Position.new(line_num, start_col),
        )
        edit = ::LSP::TextEdit.new(range, "_")
        workspace_edit = ::LSP::WorkspaceEdit.new({uri => [edit]})

        ::LSP::CodeAction.new(
          title: "Prefix '#{param_name}' with underscore to suppress warning",
          kind: "quickfix",
          diagnostics: [diag],
          edit: workspace_edit,
          is_preferred: true,
        )
      end
    end

    # ---------------------------------------------------------------------------
    # Quick-fix: insert explicit `return` at end of method body
    # ---------------------------------------------------------------------------

    private def add_missing_return_action(
      uri : String,
      lines : Array(String),
      diag : ::LSP::Diagnostic,
    ) : ::LSP::CodeAction?
      def_line = diag.range.start.line
      return if def_line >= lines.size

      def_text = lines[def_line]
      indent_size = def_text[/^\s*/].size
      body_indent = " " * (indent_size + 2)

      # Walk forward to find the matching `end` at the same indent level.
      end_line = nil
      (def_line + 1...lines.size).each do |i|
        if lines[i].strip == "end" && lines[i][/^\s*/].size == indent_size
          end_line = i
          break
        end
      end
      return unless end_line

      insert_pos = ::LSP::Position.new(end_line, 0)
      edit = ::LSP::TextEdit.new(
        ::LSP::Range.new(insert_pos, insert_pos),
        "#{body_indent}return\n",
      )
      workspace_edit = ::LSP::WorkspaceEdit.new({uri => [edit]})

      ::LSP::CodeAction.new(
        title: "Add explicit return",
        kind: "quickfix",
        diagnostics: [diag],
        edit: workspace_edit,
        is_preferred: false,
      )
    end
  end
end

module LSP
  # ---------------------------------------------------------------------------
  # Protocol types for textDocument/codeAction
  # ---------------------------------------------------------------------------

  struct CodeActionContext
    include JSON::Serializable

    property diagnostics : Array(Diagnostic)
    property? only : Array(String)?
    @[JSON::Field(key: "triggerKind")]
    property? trigger_kind : Int32?

    def initialize(@diagnostics : Array(Diagnostic), @only : Array(String)? = nil)
    end
  end

  struct CodeActionParams
    include JSON::Serializable

    @[JSON::Field(key: "textDocument")]
    property text_document : TextDocumentIdentifier
    property range : Range
    property context : CodeActionContext
  end

  struct CodeAction
    include JSON::Serializable

    property title : String
    property kind : String?
    property? diagnostics : Array(Diagnostic)?
    property? edit : WorkspaceEdit?
    @[JSON::Field(key: "isPreferred")]
    property? is_preferred : Bool?
    property? data : JSON::Any?

    def initialize(
      @title : String,
      @kind : String? = nil,
      @diagnostics : Array(Diagnostic)? = nil,
      @edit : WorkspaceEdit? = nil,
      @is_preferred : Bool? = nil,
    )
    end
  end
end
