require "./protocol"
require "../crystal/psi/psi_file"

module Lattice
  # InlayHintsProvider walks the Crystal PSI tree and emits inlay hints for:
  #
  # - Inferred types of local variables and instance variables whose type is
  #   not explicitly annotated (`: TypeName` suffix after the identifier).
  # - Method return types inferred from the AST when no explicit return
  #   annotation is present on the `def`.
  # - Parameter names at call sites — emitted as `name:` labels before each
  #   positional argument when the callee is defined in the same file or is
  #   indexed in the workspace, matching the most-wanted inlay hint kind.
  #
  # The structure here is intentionally extensible — add more `process_node`
  # branches as inference matures.

  class InlayHintsProvider
    # Default maximum source lines before hints are suppressed.
    # Overridden at construction time via the inlay_hints_max_lines config key.
    DEFAULT_MAX_LINES = 2_000

    def initialize(@psi : Psi::PsiFile, @max_lines : Int32 = DEFAULT_MAX_LINES)
      # Accumulated type map: variable/ivar name → inferred type string.
      # Built during the hints pass so that later assignments can reuse types
      # inferred from earlier ones (e.g. `b = a` after `a = 42` → Int32).
      @type_map = Hash(String, String).new

      # Method parameter name map: method_name → Array(arg_name)
      # Populated on the first walk so call-site parameter-name hints can
      # resolve parameter names without a second full AST walk per call.
      @param_map = Hash(String, Array(String)).new

      # Abstract return type map: method_name → return type string.
      # Built from `abstract def` nodes so concrete override defs whose bodies
      # cannot be cheaply typed can still receive a return-type hint.
      @abstract_return_map = Hash(String, String).new
    end

    # Build @param_map from all `def` nodes in the file.
    private def build_param_map
      @psi.walk do |node|
        next unless node.is_a?(::Crystal::Def)
        names = node.args.map(&.name)
        # Store under both the bare name and `receiver.name` for class methods.
        @param_map[node.name] = names
      end
    end

    # Build @abstract_return_map from all `abstract def` nodes in the file.
    # An `abstract def` in Crystal is a `::Crystal::Def` with `abstract = true`
    # and an explicit return_type annotation.  We collect name → return_type_string
    # so concrete override defs that lack their own annotation can inherit the hint.
    private def build_abstract_return_map
      @psi.walk do |node|
        next unless node.is_a?(::Crystal::Def)
        next unless node.abstract?
        rt = node.return_type
        next unless rt
        @abstract_return_map[node.name] = rt.to_s
      end
    end

    # Return all inlay hints for lines within *range*.
    def hints(range : ::LSP::Range) : Array(::LSP::InlayHint)
      result = [] of ::LSP::InlayHint
      lines  = @psi.source.split('\n')
      return result if lines.size > @max_lines

      # Populate the parameter-name map before the main walk so call-site
      # hints can look up parameter names for locally-defined methods.
      build_param_map

      # Populate the abstract-return map so concrete override defs that lack
      # their own return type annotation can receive an inherited hint.
      build_abstract_return_map

      # AST pass — collect typed local variable declarations.
      @psi.walk do |node|
        loc = node.location
        next unless loc

        ln = loc.line_number - 1
        next if ln < range.start.line || ln > range.ending.line

        case node
        when ::Crystal::Assign
          # Emit a hint when the RHS is a literal or a recognizable call pattern.
          # Only emit if the target is an untyped local variable or instance variable.
          target = node.target
          next unless target.is_a?(::Crystal::Var) || target.is_a?(::Crystal::InstanceVar)
          target_name = target.is_a?(::Crystal::Var) ? target.name : "@#{target.as(::Crystal::InstanceVar).name}"
          line_text = lines[ln]? || ""
          next if line_text.match(/#{Regex.escape(target_name)}\s*:/)

          inferred = infer_type(node.value)
          next unless inferred

          # Record into the type map so later references can resolve.
          @type_map[target_name] = inferred

          col = loc.column_number - 1 + target_name.size
          result << ::LSP::InlayHint.new(
            position: ::LSP::Position.new(ln, col),
            label: " : #{inferred}",
            kind: ::LSP::InlayHintKind::Type,
            padding_left: false,
            padding_right: false,
          )

        when ::Crystal::Def
          # Emit a return-type hint on the `def` line when no explicit return
          # type annotation is present and the body's last expression can be
          # cheaply typed.  Also fall back to the abstract-return map when the
          # body cannot be inferred — this annotates concrete overrides of
          # abstract defs that declared an explicit return type.
          next if node.return_type
          next if node.abstract?
          inferred = nil
          if body = node.body
            inferred = infer_type(body)
          end
          # Fall back to the abstract def's declared return type when body
          # inference yields nothing.
          inferred ||= @abstract_return_map[node.name]?
          next unless inferred

          line_text = lines[ln]? || ""
          hint_col = if m = line_text.match(/def\s+\w[\w?!]*(\([^)]*\))?/)
                       m.end(0)
                     else
                       line_text.size
                     end

          result << ::LSP::InlayHint.new(
            position: ::LSP::Position.new(ln, hint_col),
            label: " : #{inferred}",
            kind: ::LSP::InlayHintKind::Type,
            padding_left: false,
            padding_right: false,
          )

        when ::Crystal::MultiAssign
          # Tuple destructuring: `a, b, c = some_tuple_expr`
          # Emit a hint for the RHS showing it is Tuple, then individual hints
          # for each target when their types can be inferred.
          rhs_inferred = infer_type(node.values.first? || node)
          node.targets.each_with_index do |target, idx|
            next unless target.is_a?(::Crystal::Var)
            line_text = lines[ln]? || ""
            # Skip if the target already has a type annotation in source.
            next if line_text.match(/#{Regex.escape(target.name)}\s*:/)
            # Individual element type is unknown without generics, so use Tuple index.
            hint_label = rhs_inferred ? " : (#{rhs_inferred}[#{idx}])" : " : ?"
            col = loc.column_number - 1 + target.name.size
            result << ::LSP::InlayHint.new(
              position: ::LSP::Position.new(ln, col),
              label: hint_label,
              kind: ::LSP::InlayHintKind::Type,
              padding_left: false,
              padding_right: false,
            )
          end

        when ::Crystal::Call
          # Parameter-name hints at call sites.
          # Only emit for positional arguments (not named arguments, which
          # already carry their own label in source).
          args = node.args
          next if args.empty?

          # Look up parameter names for the callee.
          param_names = @param_map[node.name]?
          next unless param_names && !param_names.empty?

          # Skip calls where all args are already named in source — detected
          # by checking that the call node has no named args covering all args.
          named_count = node.named_args.try(&.size) || 0
          positional_count = args.size
          next if named_count >= positional_count

          # Emit one `name:` hint per positional argument for which a param
          # name is known. Hints are placed at the start of each argument.
          args.each_with_index do |arg, idx|
            pname = param_names[idx]?
            next unless pname
            # Skip single-character or underscore-prefixed parameter names —
            # they add noise without useful information.
            next if pname.size <= 1 || pname.starts_with?('_')

            arg_loc = arg.location
            next unless arg_loc

            arg_ln  = arg_loc.line_number - 1
            arg_col = arg_loc.column_number - 1
            next if arg_ln < range.start.line || arg_ln > range.ending.line

            result << ::LSP::InlayHint.new(
              position: ::LSP::Position.new(arg_ln, arg_col),
              label: "#{pname}:",
              kind: ::LSP::InlayHintKind::Parameter,
              padding_left: false,
              padding_right: true,
            )
          end
        end
      end

      result
    end

    # ---------------------------------------------------------------------------
    # Type inference — literals + common call patterns + variables
    # ---------------------------------------------------------------------------

    # Attempt to infer a type from an AST node without invoking the compiler.
    # Returns nil when inference is not possible.
    private def infer_type(node : ::Crystal::ASTNode) : String?
      # Fast path: direct literal
      if t = infer_literal_type(node)
        return t
      end

      case node
      when ::Crystal::Var
        # Look up previously recorded type for this variable name.
        return @type_map[node.name]?

      when ::Crystal::InstanceVar
        return @type_map["@#{node.name}"]?

      when ::Crystal::Path
        # A bare constant/type path used as a value — return its last segment
        # so that `x = SomeModule::CONST` gets a hint of "CONST" rather than
        # nothing.  Not a type name, but still more useful than silence.
        return node.names.last?

      when ::Crystal::Generic
        # e.g. Array(Int32), Hash(String, Int32) — return the outer type name.
        return node.name.to_s.split("::").last

      when ::Crystal::Cast
        # x = expr.as(T) — the type is always T.
        return node.to.to_s.split("::").last

      when ::Crystal::NilableCast
        # x = expr.as?(T) — the type is T? (nilable).
        inner = node.to.to_s.split("::").last
        return "#{inner}?"

      when ::Crystal::If, ::Crystal::Unless
        # Try to infer from the then-branch; fall back to else.
        # Only emit a hint when both branches agree on a type.
        then_type = node.then ? infer_type(node.then) : nil
        else_type = node.else ? infer_type(node.else) : nil
        if then_type && else_type
          return then_type == else_type ? then_type : nil
        end
        return then_type || else_type

      when ::Crystal::While, ::Crystal::Until
        # Loop bodies are Nil-typed by Crystal convention.
        return "Nil"

      when ::Crystal::Call
        # SomeType.new
        if node.name == "new"
          if obj = node.obj
            return obj.to_s.split("::").last
          end
        end
        # Common conversion/builder methods
        case node.name
        when "to_s"                          then return "String"
        when "to_i", "to_i32"               then return "Int32"
        when "to_i16"                        then return "Int16"
        when "to_i64"                        then return "Int64"
        when "to_i128"                       then return "Int128"
        when "to_u8"                         then return "UInt8"
        when "to_u16"                        then return "UInt16"
        when "to_u32"                        then return "UInt32"
        when "to_u64"                        then return "UInt64"
        when "to_f", "to_f64"               then return "Float64"
        when "to_f32"                        then return "Float32"
        when "to_big_i"                      then return "BigInt"
        when "to_big_f"                      then return "BigFloat"
        when "to_big_d"                      then return "BigDecimal"
        when "to_r"                          then return "Rational"
        when "to_c"                          then return "Complex"
        when "to_sym"                        then return "Symbol"
        when "to_a", "to_ary"               then return "Array"
        when "to_h"                          then return "Hash"
        when "size", "length", "count",
             "index", "rindex"               then return "Int32"
        when "empty?", "includes?",
             "starts_with?", "ends_with?",
             "has_key?", "has_value?",
             "any?", "all?", "none?",
             "one?", "zero?",
             "nil?", "is_a?", "responds_to?" then return "Bool"
        when "keys", "values", "chars",
             "lines", "split", "map",
             "flat_map", "compact_map",
             "select", "reject", "flatten",
             "uniq", "sort", "sort_by",
             "reverse", "rotate", "sample",
             "take", "first", "last",
             "zip", "each_with_object",
             "combination", "permutation"    then return "Array"
        when "upcase", "downcase", "strip",
             "chomp", "lstrip", "rstrip",
             "gsub", "sub", "tr", "squeeze",
             "delete", "encode", "decode",
             "center", "ljust", "rjust",
             "insert", "slice", "succ",
             "inspect", "dump"               then return "String"
        when "round", "ceil", "floor",
             "truncate", "abs"               then return "Float64"
        when "not_nil!"
          # x = maybe_nil.not_nil! — unwrap the receiver's inferred type.
          if obj = node.obj
            inner = infer_type(obj)
            return inner ? inner.rchop('?') : nil
          end
        end
        # Receiver-type propagation: receiver.dup / .clone / .itself
        if %w[dup clone itself].includes?(node.name)
          if obj = node.obj
            return infer_type(obj)
          end
        end

      when ::Crystal::Expressions
        # Last expression of a block determines the type.
        last = node.expressions.last?
        return last ? infer_type(last) : nil

      when ::Crystal::Assign
        # Nested assignment (e.g. `a = b = 1`) — infer from the RHS.
        return infer_type(node.value)

      when ::Crystal::MultiAssign
        # Multi-assign yields a Tuple of the RHS values — return Tuple.
        return "Tuple"

      when ::Crystal::Select
        # Crystal `select` on channels is always Nil-typed.
        return "Nil"

      when ::Crystal::TypeOf
        # typeof(expr) — the compiler resolves this; return a hint string.
        return "typeof(...)"
      end

      # Last attempt: walk a method-call chain and infer from the outermost receiver.
      # E.g. `foo.bar.baz` — try to infer the type of the full chain by peeling
      # calls from right to left until we reach a literal or known type.
      if node.is_a?(::Crystal::Call)
        receiver_type = infer_chain_type(node)
        return receiver_type if receiver_type
      end

      nil
    end

    # Infer the type of a possibly-chained call expression by peeling the chain
    # from right to left.  Returns nil when the chain cannot be resolved.
    private def infer_chain_type(node : ::Crystal::Call) : String?
      # Named constructors and well-known call patterns are already handled
      # in the main infer_type dispatch above, so this only handles chains
      # where the outermost node is not directly recognizable.
      return nil if node.name == "new"

      # Try to infer the receiver's type and apply the method name to it.
      if obj = node.obj
        recv_type = infer_type(obj)
        return nil unless recv_type

        # Strip nilability before dispatching on method name so `Array?.size`
        # still resolves correctly.
        bare = recv_type.rchop('?')

        case bare
        when "String"
          case node.name
          when "chars", "bytes", "split", "lines"
            return "Array"
          when "size", "bytesize", "count"
            return "Int32"
          when "to_f", "to_f64"
            return "Float64"
          when "to_i", "to_i32"
            return "Int32"
          when "to_i64"
            return "Int64"
          end
        when "Array"
          case node.name
          when "size", "count"
            return "Int32"
          when "join"
            return "String"
          when "first", "last", "sample"
            return nil  # element type unknown without generics
          end
        when "Hash"
          case node.name
          when "size", "count"
            return "Int32"
          when "keys"
            return "Array"
          when "values"
            return "Array"
          end
        when "Int32", "Int64", "Float64", "Float32"
          case node.name
          when "to_s"
            return "String"
          when "abs"
            return bare
          when "to_f", "to_f64"
            return "Float64"
          when "to_i", "to_i32"
            return "Int32"
          end
        end
      end

      nil
    end

    # ---------------------------------------------------------------------------
    # Cheap literal type inference (no compiler invocation required)
    # ---------------------------------------------------------------------------

    private def infer_literal_type(node : ::Crystal::ASTNode) : String?
      case node
      when ::Crystal::StringLiteral   then "String"
      when ::Crystal::CharLiteral     then "Char"
      when ::Crystal::BoolLiteral     then "Bool"
      when ::Crystal::NilLiteral      then "Nil"
      when ::Crystal::ArrayLiteral    then "Array"
      when ::Crystal::HashLiteral     then "Hash"
      when ::Crystal::RegexLiteral    then "Regex"
      when ::Crystal::SymbolLiteral   then "Symbol"
      when ::Crystal::TupleLiteral    then "Tuple"
      when ::Crystal::RangeLiteral    then "Range"
      when ::Crystal::NumberLiteral
        case node.kind
        when .i8?            then "Int8"
        when .i16?           then "Int16"
        when .i32?           then "Int32"
        when .i64?           then "Int64"
        when .u8?            then "UInt8"
        when .u16?           then "UInt16"
        when .u32?           then "UInt32"
        when .u64?           then "UInt64"
        when .f32?           then "Float32"
        when .f64?           then "Float64"
        else                      "Int32"
        end
      else
        nil
      end
    end
  end
end
