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

module Lattice
  # SemanticTokensProvider walks the Crystal PSI tree for a single file and
  # emits LSP semantic tokens encoded in the five-integer delta format required
  # by the LSP 3.17 specification.
  #
  # Token types (index == legend position):
  #   0  namespace  – module / lib names
  #   1  type       – generic type references
  #   2  class      – class names at definition sites
  #   3  enum       – enum names at definition sites
  #   4  struct     – struct names at definition sites
  #   5  parameter  – method / block parameters
  #   6  variable   – local variables and instance variables
  #   7  property   – property / getter / setter declarations
  #   8  enumMember – enum member names
  #   9  method     – method names at definition sites
  #  10  macro      – macro definition names
  #  11  keyword    – Crystal keywords (supplemental)
  #  12  comment    – doc-comment lines
  #  13  string     – string literals (supplemental)
  #  14  number     – numeric literals (supplemental)
  #
  # Token modifiers (bit-field index):
  #   0  declaration – the name is being declared here
  #   1  definition  – alias for declaration in Crystal
  #   2  readonly    – constants and enum members
  #   3  static      – class-level (self.) methods
  #   4  deprecated  – annotated with @[Deprecated]

  TOKEN_TYPES = [
    "namespace", "type", "class", "enum", "struct",
    "parameter", "variable", "property", "enumMember",
    "method", "macro", "keyword", "comment",
    "string", "number",
  ]

  TOKEN_MODIFIERS = [
    "declaration", "definition", "readonly", "static", "deprecated", "modification",
  ]

  # A single raw token before delta-encoding.
  private record RawToken,
    line : Int32,
    start : Int32,
    length : Int32,
    token_type : Int32,
    token_modifiers : Int32

  class SemanticTokensProvider
    MAX_LINES = 10_000

    def initialize(@psi : Psi::PsiFile)
    end

    # Compute semantic tokens for the entire file.
    def full : ::LSP::SemanticTokens
      ::LSP::SemanticTokens.new(encode(collect(nil)))
    end

    # Compute semantic tokens restricted to *range*.
    def range(range : ::LSP::Range) : ::LSP::SemanticTokens
      ::LSP::SemanticTokens.new(encode(collect(range)))
    end

    # ---------------------------------------------------------------------------
    # Token collection
    # ---------------------------------------------------------------------------

    private def collect(range : ::LSP::Range?) : Array(RawToken)
      tokens = [] of RawToken
      lines  = @psi.source.split('\n')
      return tokens if lines.size > MAX_LINES

      lines.each_with_index do |line, line_num|
        next if range && (line_num < range.start.line || line_num > range.ending.line)
        scan_line(line, line_num, tokens)
      end

      ast_pass(tokens, range)
      tokens.sort_by! { |t| {t.line, t.start} }
      tokens
    end

    # Regex-based line scanner for literals and comments.
    private def scan_line(line : String, line_num : Int32, tokens : Array(RawToken))
      # Doc-comment / regular comment
      if m = line.match(/^(\s*)(#.*)$/)
        tokens << RawToken.new(line_num, m[1].size, m[2].size, type_idx("comment"), 0)
        return
      end

      # Double-quoted string literals (single-line)
      offset = 0
      while m = line.match(/"(?:[^"\\]|\\.)*"/, offset)
        tokens << RawToken.new(line_num, m.begin(0), m[0].size, type_idx("string"), 0)
        offset = m.end(0)
      end

      # Numeric literals
      offset = 0
      while m = line.match(/\b(\d[\d_]*(?:\.\d[\d_]*)?(?:_[iu]\d+|_f\d+)?)\b/, offset)
        tokens << RawToken.new(line_num, m.begin(1), m[1].size, type_idx("number"), 0)
        offset = m.end(0)
      end
    end

    # AST-based pass: annotates declaration sites via the Crystal compiler AST.
    private def ast_pass(tokens : Array(RawToken), range : ::LSP::Range?)
      @psi.walk do |node|
        loc = node.location
        next unless loc

        ln  = loc.line_number - 1
        col = loc.column_number - 1
        next if range && (ln < range.start.line || ln > range.ending.line)

        case node
        when ::Crystal::ClassDef
          short = node.name.to_s.split("::").last
          emit(tokens, short, ln, col, "class", mod_declaration | mod_definition)
        when ::Crystal::ModuleDef
          short = node.name.to_s.split("::").last
          emit(tokens, short, ln, col, "namespace", mod_declaration | mod_definition)
        when ::Crystal::EnumDef
          short = node.name.to_s.split("::").last
          emit(tokens, short, ln, col, "enum", mod_declaration | mod_definition)
        when ::Crystal::Def
          mods = mod_declaration | mod_definition
          mods |= mod_static if node.receiver
          emit(tokens, node.name, ln, col, "method", mods)
          node.args.each do |arg|
            if ploc = arg.location
              emit(tokens, arg.name,
                ploc.line_number - 1, ploc.column_number - 1,
                "parameter", mod_declaration)
            end
          end
        when ::Crystal::Macro
          emit(tokens, node.name, ln, col, "macro", mod_declaration | mod_definition)
        when ::Crystal::Assign
          target = node.target
          if target.is_a?(::Crystal::Var)
            emit(tokens, target.name, ln, col, "variable", mod_declaration | mod_modification)
          elsif target.is_a?(::Crystal::InstanceVar)
            emit(tokens, target.as(::Crystal::InstanceVar).name, ln, col, "variable", mod_modification)
          end
        when ::Crystal::OpAssign
          # Compound assignment (+=, -=, etc.) — also a write.
          target = node.target
          if target.is_a?(::Crystal::Var)
            emit(tokens, target.name, ln, col, "variable", mod_modification)
          elsif target.is_a?(::Crystal::InstanceVar)
            emit(tokens, target.as(::Crystal::InstanceVar).name, ln, col, "variable", mod_modification)
          end
        when ::Crystal::InstanceVar
          # Read reference to an instance variable (not an assignment target).
          emit(tokens, node.name, ln, col, "variable", 0)
        end
      end
    end

    private def emit(
      tokens : Array(RawToken),
      name : String,
      line : Int32,
      col : Int32,
      type_name : String,
      mods : Int32,
    )
      return if name.empty?
      tokens << RawToken.new(line, col, name.size, type_idx(type_name), mods)
    end

    # ---------------------------------------------------------------------------
    # Delta encoding (LSP §3.17 — five u32 integers per token)
    # ---------------------------------------------------------------------------

    private def encode(tokens : Array(RawToken)) : Array(UInt32)
      result  = [] of UInt32
      prev_ln = 0
      prev_st = 0

      tokens.each do |tok|
        delta_line  = tok.line - prev_ln
        delta_start = delta_line == 0 ? tok.start - prev_st : tok.start

        result << delta_line.to_u32
        result << delta_start.to_u32
        result << tok.length.to_u32
        result << tok.token_type.to_u32
        result << tok.token_modifiers.to_u32

        prev_ln = tok.line
        prev_st = tok.start
      end

      result
    end

    # ---------------------------------------------------------------------------
    # Helpers
    # ---------------------------------------------------------------------------

    private def type_idx(name : String) : Int32
      TOKEN_TYPES.index(name) || 0
    end

    private def mod_declaration  : Int32; 1 << 0; end
    private def mod_definition   : Int32; 1 << 1; end
    private def mod_readonly     : Int32; 1 << 2; end
    private def mod_static       : Int32; 1 << 3; end
    private def mod_deprecated   : Int32; 1 << 4; end
    private def mod_modification : Int32; 1 << 5; end
  end
end
