from pathlib import Path

from typing import Any, Dict, Optional, Tuple

import re

from codegraphcontext.utils.debug_log import debug_log, info_logger, error_logger, warning_logger

from codegraphcontext.utils.tree_sitter_manager import execute_query

KOTLIN_QUERIES = {

"functions": """

(function_declaration

(simple_identifier) @name

(function_value_parameters) @params

) @function_node

""",

"classes": """

[

(class_declaration (type_identifier) @name)

(object_declaration (type_identifier) @name)

(companion_object (type_identifier)? @name)

] @class

""",

"imports": """

(import_header) @import

""",

"calls": """

(call_expression) @call_node

""",

"variables": """

(property_declaration

(variable_declaration

(simple_identifier) @name

)

) @variable

""",

}

class KotlinTreeSitterParser:

def __init__(self, generic_parser_wrapper: Any):

self.generic_parser_wrapper = generic_parser_wrapper

self.language_name = "kotlin"

self.language = generic_parser_wrapper.language

self.parser = generic_parser_wrapper.parser

def parse(self, path: Path, is_dependency: bool = False, index_source: bool = False) -> Dict[str, Any]:

try:

self.index_source = index_source

with open(path, 'r', encoding='utf-8', errors='ignore') as f:

source_code = f.read()

if not source_code.strip():

warning_logger(f"Empty or whitespace-only file: {path}")

return {

"path": str(path),

"functions": [],

"classes": [],

"variables": [],

"imports": [],

"function_calls": [],

"is_dependency": is_dependency,

"lang": self.language_name,

}

tree = self.parser.parse(bytes(source_code, "utf8"))

parsed_functions = []

parsed_classes = []

parsed_variables = []

parsed_imports = []

parsed_calls = []

# Parse Variables first to populate for inference

if 'variables' in KOTLIN_QUERIES:

results = execute_query(self.language, KOTLIN_QUERIES['variables'], tree.root_node)

parsed_variables = self._parse_variables(results, source_code, path)

for capture_name, query in KOTLIN_QUERIES.items():

if capture_name == 'variables': continue # Already done

results = execute_query(self.language, query, tree.root_node)

if capture_name == "functions":

parsed_functions.extend(self._parse_functions(results, source_code, path))

elif capture_name == "classes":

parsed_classes.extend(self._parse_classes(results, source_code, path))

elif capture_name == "imports":

parsed_imports.extend(self._parse_imports(results, source_code))

elif capture_name == "calls":

parsed_calls.extend(self._parse_calls(results, source_code, path, parsed_variables))

return {

"path": str(path),

"functions": parsed_functions,

"classes": parsed_classes,

"variables": parsed_variables,

"imports": parsed_imports,

"function_calls": parsed_calls,

"is_dependency": is_dependency,

"lang": self.language_name,

}

except Exception as e:

error_logger(f"Error parsing Kotlin file {path}: {e}")

return {

"path": str(path),

"functions": [],

"classes": [],

"variables": [],

"imports": [],

"function_calls": [],

"is_dependency": is_dependency,

"lang": self.language_name,

}

def _get_parent_context(self, node: Any) -> Tuple[Optional[str], Optional[str], Optional[int]]:

curr = node.parent

while curr:

if curr.type in ("function_declaration",):

name_node = None

for child in curr.children:

if child.type == "simple_identifier":

name_node = child

break

return (

self._get_node_text(name_node) if name_node else None,

curr.type,

curr.start_point[0] + 1,

)

if curr.type in ("class_declaration", "object_declaration"):

for child in curr.children:

if child.type in ("simple_identifier", "type_identifier"):

return (

self._get_node_text(child),

curr.type,

curr.start_point[0] + 1,

)

# Check for secondary constructors

if curr.type == "secondary_constructor":

return (

"constructor",

curr.type,

curr.start_point[0] + 1

)

if curr.type == "companion_object":

name = "Companion"

for child in curr.children:

if child.type in ("simple_identifier", "type_identifier"):

name = self._get_node_text(child)

break

return (

name,

curr.type,

curr.start_point[0] + 1,

)

# Handle anonymous objects (object_literal)

if curr.type == "object_literal":

# checking if it is assigned to a variable to get a name?

# or simply "AnonymousObject"

# It's usually hard to name them without variable context.

# We can check if parent is property/variable declaration

name = "AnonymousObject"

return (

name,

curr.type,

curr.start_point[0] + 1

)

curr = curr.parent

return None, None, None

def _get_node_text(self, node: Any) -> str:

if not node: return ""

return node.text.decode("utf-8")

def _parse_functions(self, captures: list, source_code: str, path: Path) -> list[Dict[str, Any]]:

functions = []

seen_nodes = set()

for node, capture_name in captures:

if capture_name == "function_node":

node_id = (node.start_byte, node.end_byte, node.type)

if node_id in seen_nodes:

continue

seen_nodes.add(node_id)

try:

start_line = node.start_point[0] + 1

end_line = node.end_point[0] + 1

# Manual child lookup

name_node = None

for child in node.children:

if child.type == "simple_identifier":

name_node = child

break

if name_node:

func_name = self._get_node_text(name_node)

params_node = None

for child in node.children:

if child.type == "function_value_parameters":

params_node = child

break

parameters = []

if params_node:

params_text = self._get_node_text(params_node)

parameters = self._extract_parameter_names(params_text)

source_text = self._get_node_text(node)

context_name, context_type, context_line = self._get_parent_context(node)

func_data = {

"name": func_name,

"args": parameters,

"line_number": start_line,

"end_line": end_line,

"path": str(path),

"lang": self.language_name,

"context": context_name,

"class_context": context_name if context_type and ("class" in context_type or "object" in context_type) else None

}

if self.index_source:

func_data["source"] = source_text

functions.append(func_data)

except Exception as e:

error_logger(f"Error parsing function in {path}: {e}")

continue

return functions

def _parse_classes(self, captures: list, source_code: str, path: Path) -> list[Dict[str, Any]]:

classes = []

seen_nodes = set()

for node, capture_name in captures:

if capture_name == "class":

node_id = (node.start_byte, node.end_byte, node.type)

if node_id in seen_nodes:

continue

seen_nodes.add(node_id)

try:

start_line = node.start_point[0] + 1

end_line = node.end_point[0] + 1

# Find name child (type_identifier or simple_identifier)

class_name = "Anonymous"

if node.type == "companion_object":

class_name = "Companion" # Default name

for child in node.children:

if child.type in ("type_identifier", "simple_identifier"):

class_name = self._get_node_text(child)

break

source_text = self._get_node_text(node)

bases = []

# Check for delegation specifiers

# class_declaration -> delegation_specifier

for child in node.children:

if child.type == "delegation_specifier":

# children: constructor_invocation or user_type

for specifier in child.children:

# constructor_invocation -> user_type -> type_identifier

# user_type -> type_identifier

# We want the text of the type

if specifier.type == "constructor_invocation":

# child 0 is typically user_type

for sub in specifier.children:

if sub.type == "user_type":

bases.append(self._get_node_text(sub))

break

elif specifier.type == "user_type":

bases.append(self._get_node_text(specifier))

elif specifier.type == "explicit_delegation":

# Not handling simple yet, uses 'by'

pass

class_data = {

"name": class_name,

"line_number": start_line,

"end_line": end_line,

"bases": bases,

"path": str(path),

"lang": self.language_name,

}

if self.index_source:

class_data["source"] = source_text

classes.append(class_data)

except Exception as e:

error_logger(f"Error parsing class in {path}: {e}")

continue

return classes

def _parse_variables(self, captures: list, source_code: str, path: Path) -> list[Dict[str, Any]]:

variables = []

seen_vars = set()

for node, capture_name in captures:

if capture_name == "variable":

try:

start_line = node.start_point[0] + 1

ctx_name, ctx_type, ctx_line = self._get_parent_context(node)

# Destructuring declaration

if "destructuring" in node.type:

pass

# Regular property/variable

var_name = "unknown"

var_type = "Unknown"

var_decl = None

for child in node.children:

if child.type == "variable_declaration":

var_decl = child

break

if var_decl:

# Check for name and type in variable_declaration

for child in var_decl.children:

if child.type == "simple_identifier":

var_name = self._get_node_text(child)

if child.type == "user_type":

var_type = self._get_node_text(child)

# Attempt inference from initializer if type is unknown

if var_type == "Unknown":

# property_declaration -> expression (e.g. call_expression)

for child in node.children:

if child.type == "call_expression":

# call_expression -> simple_identifier (constructor)

for sub in child.children:

if sub.type == "simple_identifier":

var_type = self._get_node_text(sub)

break

if var_type != "Unknown": break

if var_name != "unknown":

variables.append({

"name": var_name,

"type": var_type,

"line_number": start_line,

"path": str(path),

"lang": self.language_name,

"context": ctx_name,

"class_context": ctx_name if ctx_type and ("class" in ctx_type or "object" in ctx_type) else None

})

except Exception as e:

continue

return variables

def _parse_imports(self, captures: list, source_code: str) -> list[dict]:

imports = []

for node, capture_name in captures:

if capture_name == "import":

try:

# import_header -> "import" identifier (import_alias)?

text = self._get_node_text(node)

# remove 'import '

path = text.replace('import ', '').strip().split(' as ')[0].strip()

alias = None

if ' as ' in text:

alias = text.split(' as ')[1].strip()

imports.append({

"name": path,

"full_import_name": path,

"line_number": node.start_point[0] + 1,

"alias": alias,

"context": (None, None),

"lang": self.language_name,

"is_dependency": False,

})

except Exception as e:

continue

return imports

def _parse_calls(self, captures: list, source_code: str, path: Path, variables: list[Dict[str, Any]] = []) -> list[Dict[str, Any]]:

calls = []

seen_calls = set()

# Index variables for fast lookup: (name, context) -> type

var_map = {}

for v in variables:

key = (v['name'], v['context'])

var_map[key] = v['type']

# Fallback for null context or partial match could be added

# For class props: (name, class_context) might work if local lookup fails?

for node, capture_name in captures:

if capture_name == "call_node":

try:

# navigation_expression check

start_line = node.start_point[0] + 1

call_name = "unknown"

base_obj = None

# call_expression usually has children:

# simple_identifier (func name)

# or navigation_expression (obj.method)

# Heuristic for base object:

# If navigation_expression -> child[0] is base, child[1] is suffix (method)

# We need to look deeper into the call_expression structure.

# call_expression -> (simple_identifier)

# OR call_expression -> (navigation_expression (simple_identifier) (navigation_suffix (simple_identifier) ...))

# OR call_expression -> (navigation_expression (call_expression) ...) (chained)

# Simplified traversal to find the "function name" and "receiver"

# If it's a direct call: foo()

# If it's a method call: x.foo()

# Tree-sitter struct:

# (call_expression (simple_identifier) (call_suffix ...)) -> name = simple_identifier

# (call_expression (navigation_expression (simple_identifier) (navigation_suffix (simple_identifier))) (call_suffix))

# -> name = 2nd simple_identifier, base = 1st simple_identifier

# Let's verify children

children = node.children

first_child = children[0]

if first_child.type == "simple_identifier":

call_name = self._get_node_text(first_child)

# No explicit base object

elif first_child.type == "navigation_expression":

# x.foo

# children: operand (x), operator (.), suffix (foo)

# Usually 3 children?

# Let's inspect nav expression children

nav_children = first_child.children

if len(nav_children) >= 2:

# operand is 0

operand = nav_children[0]

# last one is suffix?

suffix = nav_children[-1]

# Suffix usually contains the method name in a navigation_suffix node or directly?

# Suffix is (navigation_suffix (simple_identifier)) usually.

if suffix.type == "navigation_suffix":

# (navigation_suffix (simple_identifier))

for c in suffix.children:

if c.type == "simple_identifier":

call_name = self._get_node_text(c)

break

elif suffix.type == "simple_identifier":

call_name = self._get_node_text(suffix)

# Base object

base_obj = self._get_node_text(operand)

if call_name == "unknown":

continue

full_name = f"{base_obj}.{call_name}" if base_obj else call_name

ctx_name, ctx_type, ctx_line = self._get_parent_context(node)

# Inference

inferred_type = None

if base_obj:

# Lookup base_obj in variables

# Try exact context

inferred_type = var_map.get((base_obj, ctx_name))

if not inferred_type:

# Try class context if we are in a method

# This logic is approximate.

# If we are in method 'foo' of 'ClassA', and 'base_obj' refers to a property of 'ClassA',

# var_map entry would have context 'ClassA'.

# But our 'variables' parsing puts context as 'ClassA' for props.

# But 'ctx_name' here is 'foo'.

# We need to know 'foo' is in 'ClassA'.

# 'get_parent_context' returns immediate parent.

pass

# Fallback: check global/file scope (context=None)

if not inferred_type:

inferred_type = var_map.get((base_obj, None))

# Fallback: check if any variable named base_obj exists (loose match)

if not inferred_type:

for (vname, vctx), vtype in var_map.items():

if vname == base_obj:

inferred_type = vtype

break

calls.append({

"name": call_name,

"full_name": full_name,

"line_number": start_line,

"args": [], # Simplified

"inferred_obj_type": inferred_type,

"context": [None, ctx_type, ctx_line], # Keeping format compatible

"class_context": [None, None],

"lang": self.language_name,

"is_dependency": False

})

except Exception as e:

continue

return calls

def _extract_parameter_names(self, params_text: str) -> list[str]:

"""

Extracts parameter names from a Kotlin parameter list string.

Handles nested generics like Map<String, Int>.

Args:

params_text (str): The text content of function_value_parameters node, e.g. "(a: Int, b: Map<String, Int>)"

Returns:

list[str]: List of parameter names.

"""

params = []

if not params_text: return params

# Remove outer parentheses

clean = params_text.strip()

if clean.startswith('(') and clean.endswith(')'):

clean = clean[1:-1]

if not clean.strip(): return params

# Robust splitting by comma, respecting brackets <>, (), [], {}

current_param = []

depth_angle = 0 # < >

depth_round = 0 # ( )

depth_square = 0 # [ ]

depth_curly = 0 # { }

raw_params = []

for char in clean:

if char == '<': depth_angle += 1

elif char == '>': depth_angle -= 1

elif char == '(': depth_round += 1

elif char == ')': depth_round -= 1

elif char == '[': depth_square += 1

elif char == ']': depth_square -= 1

elif char == '{': depth_curly += 1

elif char == '}': depth_curly -= 1

if char == ',' and depth_angle == 0 and depth_round == 0 and depth_square == 0 and depth_curly == 0:

raw_params.append("".join(current_param).strip())

current_param = []

else:

current_param.append(char)

if current_param:

raw_params.append("".join(current_param).strip())

# Process each raw parameter string to extract name

# Format: "val x: Int", "override var y: String", "@Ann z: Int", "a: Int = 5"

for p in raw_params:

if not p: continue

# Remove default value if present

# Be careful with '=' inside strings or generic defaults, but usually param defaults are at top level

# A simple split by '=' might be risky if default value has '=', but for name extraction it's usually safe

# as the name is on the LHS.

# But wait, "val x: Type = ..." -> name is before ':'

# Split by ':' to separate name/modifiers from Type

# Using the first ':' usually works, assuming name doesn't contain ':'

colon_index = p.find(':')

if colon_index != -1:

lhs = p[:colon_index].strip()

else:

# Could be a parameter without type? (not common in Kotlin unless lambda destructuring)

# Or "var x = 5" (unlikely in func params)

# Just take the whole string if no colon?

lhs = p.strip()

# LHS contains keywords (val, var), annotations (@Foo), modifiers (crossinline, noinline, vararg)

# and the parameter name. The parameter name is usually the LAST identifier.

if not lhs: continue

tokens = lhs.split()

if tokens:

# The name is the last token

params.append(tokens[-1])

return params

def pre_scan_kotlin(files: list[Path], parser_wrapper) -> dict:

name_to_files = {}

for path in files:

try:

with open(path, 'r', encoding='utf-8', errors='ignore') as f:

content = f.read()

# 1. Extract package

# package com.example.project

package_name = ""

pkg_match = re.search(r'^\s*package\s+([\w\.]+)', content, re.MULTILINE)

if pkg_match:

package_name = pkg_match.group(1)

# 2. Extract classes/objects/interfaces/typealiases

matches = re.finditer(r'\b(class|interface|object|typealias)\s+(\w+)', content)

for match in matches:

name = match.group(2)

# Map simple name

if name not in name_to_files:

name_to_files[name] = []

name_to_files[name].append(str(path))

# If package exists, map FQN

if package_name:

fqn = f"{package_name}.{name}"

if fqn not in name_to_files:

name_to_files[fqn] = []

name_to_files[fqn].append(str(path))

except Exception:

pass

return name_to_files