require "./inspection"
require "compiler/crystal/syntax"

module Lattice
  module Inspections
    # Flags `foo.not_nil!` calls when the variable `foo` can be determined to be
    # non-nilable from a simple type-map scan of the enclosing scope.
    #
    # Detection strategy: walk `Assign` nodes to build a per-method type map
    # (variable name → inferred type string, same approach as InlayHintsProvider).
    # When a `Call` node with name "not_nil!" is encountered, check whether the
    # receiver's name appears in the type map with a non-nilable type.  If so,
    # the call is redundant.
    #
    # This is a best-effort, conservative check: when type information is absent
    # or ambiguous the inspection stays silent.
    class UnnecessaryNotNilInspection < Inspection
      def run(
        uri : String,
        source : String,
        lines : Array(String),
        ast : ::Crystal::ASTNode,
      ) : Array(::LSP::Diagnostic)
        diagnostics = [] of ::LSP::Diagnostic
        visitor = Visitor.new(diagnostics)
        begin
          ast.accept(visitor)
        rescue ex : Exception
          # Never let a visitor crash block other inspections.
        end
        diagnostics
      end

      private class Visitor < ::Crystal::Visitor
        # Stack of type maps, one per method scope.
        def initialize(@diagnostics : Array(::LSP::Diagnostic))
          @type_stack = [] of Hash(String, String)
        end

        def visit(node : ::Crystal::Def) : Bool
          @type_stack.push(Hash(String, String).new)
          true
        end

        def end_visit(node : ::Crystal::Def) : Nil
          @type_stack.pop
        end

        def visit(node : ::Crystal::Assign) : Bool
          if (map = @type_stack.last?) && node.target.is_a?(::Crystal::Var)
            name = node.target.as(::Crystal::Var).name
            inferred = infer_literal_type(node.value)
            map[name] = inferred if inferred
          end
          true
        end

        def visit(node : ::Crystal::Call) : Bool
          return true unless node.name == "not_nil!"

          if obj = node.obj
            if obj.is_a?(::Crystal::Var)
              var_name = obj.as(::Crystal::Var).name
              if (map = @type_stack.last?) && (t = map[var_name]?)
                # Only flag when the stored type is non-nilable (no trailing `?`).
                unless t.ends_with?('?') || t == "Nil"
                  loc = node.location
                  return true unless loc
                  ln  = loc.line_number - 1
                  col = loc.column_number - 1
                  # Highlight the `.not_nil!` call including the dot.
                  call_text = ".not_nil!"
                  full_col = col + var_name.size
                  @diagnostics << ::LSP::Diagnostic.new(
                    ::LSP::Range.new(
                      ::LSP::Position.new(ln, full_col),
                      ::LSP::Position.new(ln, full_col + call_text.size),
                    ),
                    "Unnecessary .not_nil! — '#{var_name}' is already #{t} (non-nilable) in this scope",
                    ::LSP::DiagnosticSeverity::Warning,
                    "lattice",
                  )
                end
              end
            end
          end
          true
        end

        def visit(node : ::Crystal::ASTNode) : Bool
          true
        end

        # Infer the type of a literal node (subset of InlayHintsProvider logic).
        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::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::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
          when ::Crystal::Call
            case node.name
            when "new"
              node.obj.try { |o| o.to_s.split("::").last }
            when "to_s" then "String"
            when "to_i", "to_i32" then "Int32"
            when "to_i64" then "Int64"
            when "to_f", "to_f64" then "Float64"
            else nil
            end
          else nil
          end
        end
      end
    end
  end
end
