GitFOSS

require "tree_sitter"
require "json"

# JSON-serializable types as classes with explicit constructors
class CodeArg
  include JSON::Serializable
  property type : String
  property name : String

  def initialize(@type : String = "", @name : String = "")
  end
end

class CodeProperty
  include JSON::Serializable
  property name : String
  property type : String

  def initialize(@name : String = "", @type : String = "")
  end
end

class CodeConstant
  include JSON::Serializable
  property name : String
  property value : String

  def initialize(@name : String = "", @value : String = "")
  end
end

class CodeMethod
  include JSON::Serializable
  property name : String
  property args : Array(CodeArg) = [] of CodeArg
  property returnType : String

  def initialize(@name : String = "", @args : Array(CodeArg) = [] of CodeArg, @returnType : String = "")
  end
end

class CodeFunction
  include JSON::Serializable
  property type : String = "function"
  property name : String
  property args : Array(CodeArg) = [] of CodeArg
  property returnType : String
  property visibility : String = "public"

  def initialize(@type : String = "function", @name : String = "", @args : Array(CodeArg) = [] of CodeArg, @returnType : String = "", @visibility : String = "public")
  end
end

class CodeClass
  include JSON::Serializable
  property type : String = "class"
  property name : String
  property properties : Array(CodeProperty) = [] of CodeProperty
  property methods : Array(CodeMethod) = [] of CodeMethod
  property constants : Array(CodeConstant) = [] of CodeConstant

  def initialize(@type : String = "class", @name : String = "", @properties : Array(CodeProperty) = [] of CodeProperty, @methods : Array(CodeMethod) = [] of CodeMethod, @constants : Array(CodeConstant) = [] of CodeConstant)
  end
end

class CodeTopConstant
  include JSON::Serializable
  property type : String = "constant"
  property name : String
  property value : String

  def initialize(@type : String = "constant", @name : String = "", @value : String = "")
  end
end

alias CodeSymbol = CodeFunction | CodeClass | CodeTopConstant

# RepositoryFile for demonstration purposes
class RepositoryFile
  property name : String
  property content : String

  def initialize(@name, @content)
  end
end

def get_file_content(file : RepositoryFile) : String
  file.content
end

# Helpers to iterate children (prefer named children)
def each_child(node : TreeSitter::Node)
  count = node.named_child_count
  if count > 0
    i = 0
    while i < count
      yield node.named_child(i.to_i)
      i += 1
    end
  else
    count = node.child_count
    i = 0
    while i < count
      yield node.child(i)
      i += 1
    end
  end
end

# Find first child matching predicate
def find_child(node : TreeSitter::Node, source : String, &block : (TreeSitter::Node -> Bool)) : TreeSitter::Node?
  found = nil : TreeSitter::Node?
  each_child(node) do |child|
    if block.call(child)
      found = child
      break
    end
  end
  found
end

def node_text_or_empty(node : TreeSitter::Node?, source : String) : String
  if node
    node.text(source)
  else
    ""
  end
end

def extract_symbols(parser : TreeSitter::Parser, node : TreeSitter::Node, file : RepositoryFile) : Array(CodeSymbol)
  source = get_file_content(file)
  symbols = [] of CodeSymbol

  case node.type
  when "function"
    args = [] of CodeArg
    each_child(node) do |child|
      if child.type == "parameter"
        param_type_node = find_child(child, source) { |n| n.type == "type" }
        param_name_node = find_child(child, source) { |n| n.type == "identifier" }
        args << CodeArg.new(type: node_text_or_empty(param_type_node, source), name: node_text_or_empty(param_name_node, source))
      end
    end

    return_type_node = find_child(node, source) { |c| c.type == "type" }
    name_node = find_child(node, source) { |c| c.type == "identifier" }

    symbols << CodeFunction.new(
      name: node_text_or_empty(name_node, source),
      args: args,
      returnType: node_text_or_empty(return_type_node, source),
      visibility: "public"
    )
  when "class"
    properties = [] of CodeProperty
    methods = [] of CodeMethod
    constants = [] of CodeConstant

    each_child(node) do |child|
      case child.type
      when "property"
        prop_name_node = find_child(child, source) { |n| n.type == "identifier" }
        prop_type_node = find_child(child, source) { |n| n.type == "type" }
        properties << CodeProperty.new(
          name: node_text_or_empty(prop_name_node, source),
          type: node_text_or_empty(prop_type_node, source)
        )
      when "function"
        method_args = [] of CodeArg
        each_child(child) do |grandchild|
          if grandchild.type == "parameter"
            ma_type_node = find_child(grandchild, source) { |g| g.type == "type" }
            ma_name_node = find_child(grandchild, source) { |g| g.type == "identifier" }
            method_args << CodeArg.new(type: node_text_or_empty(ma_type_node, source), name: node_text_or_empty(ma_name_node, source))
          end
        end
        method_return_node = find_child(child, source) { |n| n.type == "type" }
        method_name_node = find_child(child, source) { |n| n.type == "identifier" }
        methods << CodeMethod.new(
          name: node_text_or_empty(method_name_node, source),
          args: method_args,
          returnType: node_text_or_empty(method_return_node, source)
        )
      when "constant"
        const_name_node = find_child(child, source) { |n| n.type == "identifier" }
        const_value_node = find_child(child, source) { |n| n.type == "type" }
        constants << CodeConstant.new(
          name: node_text_or_empty(const_name_node, source),
          value: node_text_or_empty(const_value_node, source)
        )
      end
    end

    class_name_node = find_child(node, source) { |c| c.type == "identifier" }
    symbols << CodeClass.new(
      name: node_text_or_empty(class_name_node, source),
      properties: properties,
      methods: methods,
      constants: constants
    )
  when "constant"
    name_node = find_child(node, source) { |c| c.type == "identifier" }
    value_node = find_child(node, source) { |c| c.type == "type" }
    symbols << CodeTopConstant.new(
      name: node_text_or_empty(name_node, source),
      value: node_text_or_empty(value_node, source)
    )
  end

  # Recurse into children
  each_child(node) do |child|
    symbols.concat(extract_symbols(parser, child, file))
  end

  symbols
end