ragmcp/mcp_codebase.py

import os
import signal
import json
from collections import defaultdict, deque
import ast

import re
import subprocess
import logging
from pathlib import Path
from typing import List, Dict, Tuple, Any, Optional
from threading import RLock
from contextlib import contextmanager
from functools import lru_cache
import time
import hashlib

import pathspec
import requests
import numpy as np

from mcp.server.fastmcp import FastMCP
from langchain_ollama import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from rank_bm25 import BM25Okapi
import javalang
import sqlparse
import sqlglot

from enhanced_toon import EnhancedToon
from graph.graph import LocalGraph


os.environ["ANONYMIZED_TELEMETRY"] = "False"

# -----------------------------
# Configuration
# -----------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.FileHandler("mcp_codebase.log"),  # Log to file
        logging.StreamHandler()  # Also log to console
    ]
)
logger = logging.getLogger("rag-mcp")

OLLAMA_BASE_URL = "http://127.0.0.1:11434"
os.environ["OLLAMA_HOST"] = OLLAMA_BASE_URL
os.environ["OLLAMA_API_BASE"] = OLLAMA_BASE_URL

CODEBASE_PATH = Path("./working_repo")
CODEBASE_PATH.mkdir(exist_ok=True)

VECTOR_DB_PATH = Path(os.environ.get("VECTOR_DB_PATH", "./chroma_db"))
BM25_INDEX_PATH = Path(os.environ.get("BM25_INDEX_PATH", "./bm25_index.json"))
RAGIGNORE_PATH = CODEBASE_PATH / ".ragignore"

# Models / behavior
EMBEDDING_MODEL = "bge-m3"  # Ollama local embedding model
RERANKER_MODEL_OLLAMA = "dengcao/Qwen3-Reranker-8B:Q4_K_M"  # prompt-based reranker via Ollama
ENABLE_RERANK = True  # toggle reranking
RERANK_TOP_N = 10

# Hybrid weights (vector first: 80%, BM25: 20%)
VECTOR_WEIGHT = 0.8
BM25_WEIGHT = 0.2

# Embedding batch configuration
EMBEDDING_BATCH_SIZE = 200

# Contextual weighting configuration
CONTEXTUAL_WEIGHTS = {
    # SQL weights
    "sql_schema": 1.3,      # CREATE TABLE, ALTER TABLE, schema definitions
    "sql_function": 1.2,    # Function/procedure definitions
    "sql_view": 1.15,       # VIEW definitions
    "sql_index": 1.1,       # INDEX definitions

    # Python weights
    "python_class": 1.25,   # Class definitions (enhanced)
    "python_function": 1.15, # Function definitions
    "python_method": 1.2,   # Method definitions (enhanced)
    "python_async": 1.1,    # Async functions

    # Go weights
    "go_type": 1.25,        # Type definitions (struct, interface)
    "go_function": 1.15,    # Function definitions
    "go_method": 1.2,       # Method definitions (enhanced)
    "go_struct": 1.3,       # Struct definitions (enhanced)

    # Java weights
    "java_class": 1.3,      # Class definitions (enhanced)
    "java_interface": 1.25, # Interface definitions (enhanced)
    "java_method": 1.2,     # Method definitions (enhanced)
    "java_constructor": 1.15, # Constructor definitions (enhanced)
    "java_enum": 1.2,       # Enum definitions (enhanced)
    "java_field": 1.1,      # Field definitions (enhanced)

    # Rust weights (new)
    "rust_struct": 1.3,     # Struct definitions
    "rust_enum": 1.25,      # Enum definitions
    "rust_trait": 1.35,     # Trait definitions (very important in Rust)
    "rust_impl": 1.2,       # Implementation blocks
    "rust_function": 1.15,  # Function definitions
    "rust_method": 1.2,     # Method definitions
    "rust_async": 1.1,      # Async functions
    "rust_unsafe": 1.15,    # Unsafe blocks/functions

    # TypeScript/JavaScript weights (for Svelte scripts)
    "typescript_interface": 1.25, # Interface definitions
    "typescript_class": 1.2,      # Class definitions
    "typescript_function": 1.15,  # Function definitions
    "typescript_type": 1.2,       # Type definitions

    # Svelte weights (new)
    "svelte_component": 1.4,      # Component definitions (very important)
    "svelte_script": 1.1,         # Script blocks
    "svelte_style": 1.05,         # Style blocks
    "svelte_markup": 1.15,        # Markup/template
    "svelte_prop": 1.25,          # Component props (enhanced)
    "svelte_reactive": 1.2,       # Reactive declarations

    # Default weight for unspecified types
    "default": 1.0
}

SUPPORTED_EXTENSIONS = {
    ".py", ".go", ".rs", ".ts", ".tsx", ".js", ".jsx",
    ".svelte", ".sql", ".pgsql", ".sh", ".bash", ".zsh",
    ".rb", ".java", ".cpp", ".c", ".h", ".hpp"
}

# MCP
mcp = FastMCP("codebase-rag")

# Global state
vectorstore: Optional[Chroma] = None
bm25: Optional[BM25Okapi] = None
bm25_corpus: List[str] = []
chunks_metadata: List[Dict[str, Any]] = []
embeddings: Optional[OllamaEmbeddings] = None
_startup_lock = RLock()
index_build_time: float = 0.0

# -----------------------------
# Thread-Safe Context Manager
# -----------------------------
@contextmanager
def safe_search():
    """Thread-safe context manager for search operations."""
    with _startup_lock:
        if vectorstore is None or bm25 is None:
            raise RuntimeError("Indexes not ready. Please wait for initialization or rebuild.")
        yield

# -----------------------------
# File Hashing for Cache Invalidation
# -----------------------------
def get_file_hash(filepath: Path) -> str:
    """Generate hash of file content + mtime for cache key."""
    try:
        stat = filepath.stat()
        content_sample = filepath.read_bytes()[:1024]  # First 1KB for speed
        hash_input = f"{stat.st_mtime}:{stat.st_size}:{content_sample}".encode()
        return hashlib.md5(hash_input).hexdigest()
    except Exception:
        return ""

# -----------------------------
# Contextual Weight Calculator
# -----------------------------
def calculate_contextual_weight(metadata: Dict[str, Any]) -> float:
    """
    Calculate contextual weight multiplier based on chunk metadata.
    Higher weights for schema-defining statements, function definitions, etc.
    """
    language = metadata.get("language", "").lower()
    chunk_type = metadata.get("type", "").lower()
    block_type = metadata.get("block_type", "").lower()

    # SQL-specific weights
    if language == "sql":
        if chunk_type in ("create", "alter", "create table", "alter table"):
            return CONTEXTUAL_WEIGHTS.get("sql_schema", 1.0)
        elif chunk_type in ("function", "procedure", "create function", "create procedure"):
            return CONTEXTUAL_WEIGHTS.get("sql_function", 1.0)
        elif chunk_type in ("view", "create view"):
            return CONTEXTUAL_WEIGHTS.get("sql_view", 1.0)
        elif chunk_type in ("index", "create index"):
            return CONTEXTUAL_WEIGHTS.get("sql_index", 1.0)

    # Python-specific weights (enhanced)
    elif language == "python":
        if chunk_type == "classdef":
            return CONTEXTUAL_WEIGHTS.get("python_class", 1.0)
        elif chunk_type in ("functiondef", "asyncfunctiondef"):
            if metadata.get("is_method"):
                return CONTEXTUAL_WEIGHTS.get("python_method", 1.0)
            elif chunk_type == "asyncfunctiondef" or metadata.get("is_async"):
                return CONTEXTUAL_WEIGHTS.get("python_async", 1.0)
            else:
                return CONTEXTUAL_WEIGHTS.get("python_function", 1.0)

    # Go-specific weights (enhanced)
    elif language == "go":
        if chunk_type in ("struct", "interface"):
            return CONTEXTUAL_WEIGHTS.get("go_struct", 1.0)
        elif chunk_type == "type":
            return CONTEXTUAL_WEIGHTS.get("go_type", 1.0)
        elif chunk_type == "method":
            return CONTEXTUAL_WEIGHTS.get("go_method", 1.0)
        elif chunk_type == "func":
            return CONTEXTUAL_WEIGHTS.get("go_function", 1.0)

    # Java-specific weights (enhanced)
    elif language == "java":
        if chunk_type == "class":
            return CONTEXTUAL_WEIGHTS.get("java_class", 1.0)
        elif chunk_type == "interface":
            return CONTEXTUAL_WEIGHTS.get("java_interface", 1.0)
        elif chunk_type == "method":
            return CONTEXTUAL_WEIGHTS.get("java_method", 1.0)
        elif chunk_type == "constructor":
            return CONTEXTUAL_WEIGHTS.get("java_constructor", 1.0)
        elif chunk_type == "enum":
            return CONTEXTUAL_WEIGHTS.get("java_enum", 1.0)
        elif chunk_type == "field":
            return CONTEXTUAL_WEIGHTS.get("java_field", 1.0)

    # Rust-specific weights (new)
    elif language == "rust":
        if chunk_type == "struct":
            return CONTEXTUAL_WEIGHTS.get("rust_struct", 1.0)
        elif chunk_type == "enum":
            return CONTEXTUAL_WEIGHTS.get("rust_enum", 1.0)
        elif chunk_type == "trait":
            return CONTEXTUAL_WEIGHTS.get("rust_trait", 1.0)
        elif chunk_type == "impl":
            return CONTEXTUAL_WEIGHTS.get("rust_impl", 1.0)
        elif chunk_type == "function":
            if metadata.get("is_async"):
                return CONTEXTUAL_WEIGHTS.get("rust_async", 1.0)
            elif metadata.get("is_unsafe"):
                return CONTEXTUAL_WEIGHTS.get("rust_unsafe", 1.0)
            elif metadata.get("is_method", False):
                return CONTEXTUAL_WEIGHTS.get("rust_method", 1.0)
            else:
                return CONTEXTUAL_WEIGHTS.get("rust_function", 1.0)

    # TypeScript/JavaScript weights (for Svelte scripts)
    elif language == "typescript":
        if chunk_type == "interface":
            return CONTEXTUAL_WEIGHTS.get("typescript_interface", 1.0)
        elif chunk_type == "class":
            return CONTEXTUAL_WEIGHTS.get("typescript_class", 1.0)
        elif chunk_type == "function":
            return CONTEXTUAL_WEIGHTS.get("typescript_function", 1.0)
        elif chunk_type == "type":
            return CONTEXTUAL_WEIGHTS.get("typescript_type", 1.0)

    # Svelte-specific weights (new)
    elif language == "svelte":
        if chunk_type == "component":
            return CONTEXTUAL_WEIGHTS.get("svelte_component", 1.0)
        elif chunk_type == "prop":
            return CONTEXTUAL_WEIGHTS.get("svelte_prop", 1.0)
        elif chunk_type == "reactive":
            return CONTEXTUAL_WEIGHTS.get("svelte_reactive", 1.0)
        elif block_type == "markup":
            return CONTEXTUAL_WEIGHTS.get("svelte_markup", 1.0)
        elif block_type == "script":
            return CONTEXTUAL_WEIGHTS.get("svelte_script", 1.0)
        elif block_type == "style":
            return CONTEXTUAL_WEIGHTS.get("svelte_style", 1.0)

    # CSS weights (for Svelte styles)
    elif language == "css":
        if block_type == "style":
            return CONTEXTUAL_WEIGHTS.get("svelte_style", 1.0)

    return CONTEXTUAL_WEIGHTS.get("default", 1.0)

def apply_contextual_weights_to_embeddings(embeddings_list: List[List[float]],
                                          metadata_list: List[Dict[str, Any]]) -> List[List[float]]:
    """
    Apply contextual weights to embedding vectors by scaling them.
    This biases the vector space without changing the embedding model.
    """
    weighted_embeddings = []

    if len(embeddings_list) != len(metadata_list):
        raise ValueError(
            f"Embedding/metadata mismatch: {len(embeddings_list)} embeddings vs {len(metadata_list)} metadata. "
            "This will break downstream systems (Chroma)."
        )

    for embedding, metadata in zip(embeddings_list, metadata_list):
        weight = calculate_contextual_weight(metadata)

        # Convert to numpy for easier manipulation
        emb_array = np.array(embedding)

        # Scale the embedding vector by the weight
        # This effectively increases the magnitude, making it more "important"
        weighted_emb = emb_array * weight

        # Optionally normalize to maintain consistent vector magnitudes
        # Comment out if you want the raw weighted vectors
        norm = np.linalg.norm(weighted_emb)
        if norm > 0:
            weighted_emb = weighted_emb / norm * np.linalg.norm(emb_array)

        weighted_embeddings.append(weighted_emb.tolist())

        if weight != 1.0:
            logger.debug(f"Applied weight {weight:.2f} to {metadata.get('file')}:{metadata.get('type')}")

    return weighted_embeddings


# -----------------------------
# Utility: .ragignore
# -----------------------------
def load_ragignore(base_path: Path) -> pathspec.PathSpec:
    if RAGIGNORE_PATH.exists():
        with open(RAGIGNORE_PATH, "r", encoding="utf-8") as f:
            return pathspec.PathSpec.from_lines("gitwildmatch", f)
    default_patterns = [
        ".git/", "node_modules/", "__pycache__/", "*.pyc",
        "*.so", "*.dll", "*.exe", "venv/", ".venv/", "dist/",
        "build/", "*.log", "*.tmp", ".DS_Store"
    ]
    return pathspec.PathSpec.from_lines("gitwildmatch", default_patterns)

# -----------------------------
# Python AST chunking with caching
# -----------------------------
@lru_cache(maxsize=1000)
def parse_python_file_cached(filepath_str: str, file_hash: str) -> tuple:
    """Cached Python AST parsing with enhanced metadata for graph relationships.
    ALWAYS returns a tuple of (text, metadata_tuple) entries. If parsing
    fails or produces no chunks, return a single fallback chunk containing
    the whole file text and minimal metadata.
    """
    filepath = Path(filepath_str)
    chunks = []

    logger.info(f"🐍 Parsing Python file: {filepath}")

    try:
        code = filepath.read_text(encoding="utf-8")
        tree = ast.parse(code, filename=str(filepath))
        lines = code.splitlines(keepends=True)

        # Track imports for module relationships
        imports = []
        module_imports = {}  # Maps names to modules for better tracking

        # First pass: Extract all imports
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    import_name = alias.asname if alias.asname else alias.name
                    imports.append(alias.name)
                    module_imports[import_name] = alias.name
            elif isinstance(node, ast.ImportFrom):
                module = node.module or ""
                for alias in node.names:
                    import_name = alias.asname if alias.asname else alias.name
                    full_import = f"{module}.{alias.name}" if module else alias.name
                    imports.append(full_import)
                    module_imports[import_name] = full_import

        logger.info(f"  Found {len(imports)} imports")

        # Second pass: Process declarations
        decl_count = 0
        for node in ast.walk(tree):
            chunk_text = None
            metadata = None

            # FUNCTION DEFINITIONS
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                docstring = ast.get_docstring(node) or ""
                start_line = node.lineno - 1
                end_line = getattr(node, "end_lineno", start_line + 1)

                # Guard the end_line
                if end_line <= start_line:
                    end_line = start_line + 1

                chunk_code = "".join(lines[start_line:end_line])

                metadata = {
                    "file": str(filepath),
                    "name": node.name,
                    "type": "function" if isinstance(node, ast.FunctionDef) else "async_function",
                    "line": node.lineno,
                    "language": "python",
                }

                # Extract function metadata
                func_meta = _extract_function_metadata(node, code)
                metadata.update(func_meta)

                # Include imports
                if imports:
                    metadata["imports"] = imports

                # Build chunk text
                chunk_text = f"File: {filepath}\nType: {metadata['type']}\nName: {node.name}\n"

                # Add decorators
                if func_meta.get("decorators"):
                    chunk_text += f"Decorators: {', '.join(func_meta['decorators'])}\n"

                # Add arguments
                args = _format_arguments(node.args)
                chunk_text += f"Arguments: {args}\n"

                # Add return type
                returns = _extract_return_annotation(node)
                if returns:
                    chunk_text += f"Returns: {returns}\n"

                # Add async/generator info
                if isinstance(node, ast.AsyncFunctionDef):
                    chunk_text += "Async: yes\n"
                if func_meta.get("is_generator"):
                    chunk_text += "Generator: yes\n"

                # Add docstring
                if docstring:
                    chunk_text += f"Docstring:\n{docstring}\n\n"

                chunk_text += f"Code:\n{chunk_code}"

            # CLASS DEFINITIONS
            elif isinstance(node, ast.ClassDef):
                docstring = ast.get_docstring(node) or ""
                start_line = node.lineno - 1
                end_line = getattr(node, "end_lineno", start_line + 1)

                if end_line <= start_line:
                    end_line = start_line + 1

                chunk_code = "".join(lines[start_line:end_line])

                metadata = {
                    "file": str(filepath),
                    "name": node.name,
                    "type": "class",
                    "line": node.lineno,
                    "language": "python",
                }

                # Extract class metadata
                class_meta = _extract_class_metadata(node, code)
                metadata.update(class_meta)

                # Include imports
                if imports:
                    metadata["imports"] = imports

                # Build chunk text
                chunk_text = f"File: {filepath}\nType: class\nName: {node.name}\n"

                # Add decorators
                if class_meta.get("decorators"):
                    chunk_text += f"Decorators: {', '.join(class_meta['decorators'])}\n"

                # Add base classes
                if class_meta.get("bases"):
                    chunk_text += f"Base Classes: {', '.join(class_meta['bases'])}\n"

                # Add methods
                if class_meta.get("methods"):
                    chunk_text += f"Methods: {', '.join(class_meta['methods'])}\n"

                # Add properties
                if class_meta.get("properties"):
                    chunk_text += f"Properties: {', '.join(class_meta['properties'])}\n"

                # Add class variables
                if class_meta.get("class_variables"):
                    chunk_text += f"Class Variables: {', '.join(class_meta['class_variables'])}\n"

                # Add docstring
                if docstring:
                    chunk_text += f"Docstring:\n{docstring}\n\n"

                chunk_text += f"Code:\n{chunk_code}"

            # MODULE-LEVEL ASSIGNMENTS (constants, globals)
            elif isinstance(node, ast.Assign) and node.lineno:
                # Only process top-level assignments (not inside functions/classes)
                if _is_module_level(node, tree):
                    for target in node.targets:
                        if isinstance(target, ast.Name):
                            # Check if it looks like a constant (UPPER_CASE)
                            if target.id.isupper():
                                start_line = node.lineno - 1
                                end_line = getattr(node, "end_lineno", start_line + 1)

                                if end_line <= start_line:
                                    end_line = start_line + 1

                                chunk_code = "".join(lines[start_line:end_line])

                                metadata = {
                                    "file": str(filepath),
                                    "name": target.id,
                                    "type": "constant",
                                    "line": node.lineno,
                                    "language": "python",
                                }

                                # Try to extract value
                                try:
                                    if hasattr(ast, 'unparse'):
                                        value_str = ast.unparse(node.value)
                                        metadata["value"] = value_str
                                except Exception:
                                    pass

                                chunk_text = f"File: {filepath}\nType: constant\nName: {target.id}\n"
                                chunk_text += f"Code:\n{chunk_code}"

            # ANNOTATED ASSIGNMENTS (type-annotated variables)
            elif isinstance(node, ast.AnnAssign) and node.lineno:
                if _is_module_level(node, tree) and isinstance(node.target, ast.Name):
                    start_line = node.lineno - 1
                    end_line = getattr(node, "end_lineno", start_line + 1)

                    if end_line <= start_line:
                        end_line = start_line + 1

                    chunk_code = "".join(lines[start_line:end_line])

                    metadata = {
                        "file": str(filepath),
                        "name": node.target.id,
                        "type": "variable",
                        "line": node.lineno,
                        "language": "python",
                    }

                    # Extract type annotation
                    if node.annotation:
                        try:
                            if hasattr(ast, 'unparse'):
                                type_str = ast.unparse(node.annotation)
                                metadata["var_type"] = type_str
                        except Exception:
                            pass

                    chunk_text = f"File: {filepath}\nType: variable\nName: {node.target.id}\n"
                    if metadata.get("var_type"):
                        chunk_text += f"Type: {metadata['var_type']}\n"
                    chunk_text += f"Code:\n{chunk_code}"

            # Add the chunk if we created one
            if chunk_text and metadata:
                chunks.append({
                    "text": chunk_text,
                    "metadata": metadata
                })
                decl_count += 1

        logger.info(f"  Parsed {decl_count} declarations from {filepath}")

    except SyntaxError as e:
        logger.warning(f"Python syntax error in {filepath}: {e}")
    except Exception as e:
        logger.warning(f"Failed to parse Python file {filepath}: {e}")

    # Return parsed chunks or fallback
    if chunks:
        logger.info(f"✅ Successfully parsed {len(chunks)} chunks from {filepath}")
        return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)

    # Fallback: emit a whole-file chunk
    logger.warning(f"❌ No chunks created from Python parsing for {filepath}. Emitting fallback whole-file chunk.")
    try:
        full_code = filepath.read_text(encoding="utf-8")
    except Exception:
        full_code = ""

    fallback_meta = {
        "file": str(filepath),
        "name": filepath.stem,
        "type": "module",
        "line": 1,
        "language": "python",
    }
    return ((full_code, tuple(fallback_meta.items())),)


def _is_module_level(node, tree):
    """Check if a node is at module level (not inside a function or class)."""
    for parent in ast.walk(tree):
        if isinstance(parent, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            for child in ast.walk(parent):
                if child is node:
                    return False
    return True


def _extract_import_info(node):
    """Extract import information from Import or ImportFrom nodes."""
    if isinstance(node, ast.Import):
        return [alias.name for alias in node.names]
    elif isinstance(node, ast.ImportFrom):
        module = node.module or ""
        return [f"{module}.{alias.name}" if module else alias.name for alias in node.names]
    return []


def _extract_function_metadata(node, code):
    """Extract enhanced metadata from function nodes."""
    metadata = {}

    # Extract decorators
    if node.decorator_list:
        decorators = []
        for dec in node.decorator_list:
            try:
                if hasattr(ast, 'unparse'):
                    decorators.append(ast.unparse(dec))
                elif isinstance(dec, ast.Name):
                    decorators.append(dec.id)
                elif isinstance(dec, ast.Call) and isinstance(dec.func, ast.Name):
                    decorators.append(dec.func.id)
            except Exception:
                pass
        if decorators:
            metadata["decorators"] = decorators

    # Extract parameters with type hints
    parameters = []
    if node.args.args:
        for arg in node.args.args:
            param_info = arg.arg
            if arg.annotation:
                try:
                    if hasattr(ast, 'unparse'):
                        param_info += f": {ast.unparse(arg.annotation)}"
                except Exception:
                    pass
            parameters.append(param_info)

    if parameters:
        metadata["parameters"] = parameters

    # Check for special function types
    if node.args.vararg:
        metadata["has_varargs"] = True
    if node.args.kwarg:
        metadata["has_kwargs"] = True

    # Check if it's a generator (contains yield)
    for child in ast.walk(node):
        if isinstance(child, (ast.Yield, ast.YieldFrom)):
            metadata["is_generator"] = True
            break

    # Check for property decorator
    if any(isinstance(dec, ast.Name) and dec.id == "property" for dec in node.decorator_list):
        metadata["is_property"] = True

    # Check for staticmethod/classmethod
    for dec in node.decorator_list:
        if isinstance(dec, ast.Name):
            if dec.id == "staticmethod":
                metadata["is_static"] = True
            elif dec.id == "classmethod":
                metadata["is_classmethod"] = True

    # Extract return type annotation
    if node.returns:
        try:
            if hasattr(ast, 'unparse'):
                metadata["return_type"] = ast.unparse(node.returns)
        except Exception:
            pass

    return metadata


def _extract_class_metadata(node, code):
    """Extract enhanced metadata from class nodes."""
    metadata = {}

    # Extract decorators
    if node.decorator_list:
        decorators = []
        for dec in node.decorator_list:
            try:
                if hasattr(ast, 'unparse'):
                    decorators.append(ast.unparse(dec))
                elif isinstance(dec, ast.Name):
                    decorators.append(dec.id)
            except Exception:
                pass
        if decorators:
            metadata["decorators"] = decorators

    # Extract base classes
    bases = []
    for base in node.bases:
        try:
            if hasattr(ast, 'unparse'):
                bases.append(ast.unparse(base))
            elif isinstance(base, ast.Name):
                bases.append(base.id)
        except Exception:
            pass
    if bases:
        metadata["bases"] = bases

    # Extract methods and properties
    methods = []
    properties = []
    static_methods = []
    class_methods = []
    class_variables = []

    for item in node.body:
        if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
            # Check decorators
            is_property = any(
                isinstance(dec, ast.Name) and dec.id == "property"
                for dec in item.decorator_list
            )
            is_static = any(
                isinstance(dec, ast.Name) and dec.id == "staticmethod"
                for dec in item.decorator_list
            )
            is_classmethod = any(
                isinstance(dec, ast.Name) and dec.id == "classmethod"
                for dec in item.decorator_list
            )

            if is_property:
                properties.append(item.name)
            elif is_static:
                static_methods.append(item.name)
            elif is_classmethod:
                class_methods.append(item.name)
            else:
                methods.append(item.name)

        elif isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name):
            # Class variable with type annotation
            class_variables.append(item.target.id)

        elif isinstance(item, ast.Assign):
            # Class variable without type annotation
            for target in item.targets:
                if isinstance(target, ast.Name):
                    class_variables.append(target.id)

    if methods:
        metadata["methods"] = methods
    if properties:
        metadata["properties"] = properties
    if static_methods:
        metadata["static_methods"] = static_methods
    if class_methods:
        metadata["class_methods"] = class_methods
    if class_variables:
        metadata["class_variables"] = class_variables

    # Check for metaclass
    for keyword in node.keywords:
        if keyword.arg == "metaclass":
            try:
                if hasattr(ast, 'unparse'):
                    metadata["metaclass"] = ast.unparse(keyword.value)
                elif isinstance(keyword.value, ast.Name):
                    metadata["metaclass"] = keyword.value.id
            except Exception:
                pass

    return metadata


def _format_arguments(args):
    """Format function arguments with type hints."""
    arg_strs = []

    # Regular arguments
    for arg in args.args:
        arg_str = arg.arg
        if arg.annotation:
            try:
                if hasattr(ast, 'unparse'):
                    arg_str += f": {ast.unparse(arg.annotation)}"
            except Exception:
                pass
        arg_strs.append(arg_str)

    # *args
    if args.vararg:
        vararg_str = f"*{args.vararg.arg}"
        if args.vararg.annotation:
            try:
                if hasattr(ast, 'unparse'):
                    vararg_str += f": {ast.unparse(args.vararg.annotation)}"
            except Exception:
                pass
        arg_strs.append(vararg_str)

    # **kwargs
    if args.kwarg:
        kwarg_str = f"**{args.kwarg.arg}"
        if args.kwarg.annotation:
            try:
                if hasattr(ast, 'unparse'):
                    kwarg_str += f": {ast.unparse(args.kwarg.annotation)}"
            except Exception:
                pass
        arg_strs.append(kwarg_str)

    return ", ".join(arg_strs) if arg_strs else "None"


def _extract_return_annotation(node):
    """Extract return type annotation from function node."""
    if not node.returns:
        return None

    try:
        if hasattr(ast, 'unparse'):
            return ast.unparse(node.returns)
    except Exception:
        pass

    return None

def parse_python_file(filepath: Path) -> List[Dict]:
    """Parse Python file with caching."""
    file_hash = get_file_hash(filepath)
    cached_result = parse_python_file_cached(str(filepath), file_hash)
    return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]

@lru_cache(maxsize=1000)
def parse_sql_file_cached(filepath_str: str, file_hash: str) -> tuple:
    """Cached SQL parsing. Returns tuple for hashability."""
    filepath = Path(filepath_str)
    try:
        sql_code = filepath.read_text(encoding="utf-8")
    except Exception as e:
        logger.warning(f"parse_sql_file: cannot read {filepath}: {e}")
        return tuple()

    chunks: List[Dict[str, Any]] = []
    statements = [s for s in sqlparse.split(sql_code) if s and s.strip()]
    lines = sql_code.splitlines()
    search_pos = 0

    for stmt in statements:
        start_idx = sql_code.find(stmt, search_pos)
        if start_idx == -1:
            start_idx = sql_code.find(stmt)
        if start_idx == -1:
            start_line = len(lines)
        else:
            start_line = sql_code[:start_idx].count("\n") + 1
            search_pos = start_idx + len(stmt)

        leading_comments = include_leading_comments(lines, start_line)
        stmt_clean = stmt.strip()

        meta: Dict[str, Any] = {
            "file": str(filepath),
            "name": None,
            "type": None,
            "line": start_line,
            "language": "sql"
        }

        try:
            # Special handling for CREATE TYPE ENUM (PostgreSQL)
            enum_match = re.match(
                r'CREATE\s+TYPE\s+(\w+)\s+AS\s+ENUM\s*\((.*?)\)',
                stmt_clean,
                re.IGNORECASE | re.DOTALL
            )
            if enum_match:
                enum_name = enum_match.group(1)
                enum_values = enum_match.group(2)
                meta["type"] = "create type enum"
                meta["name"] = enum_name
                meta["enum_values"] = ",".join([v.strip().strip("'\"") for v in enum_values.split(",") if v.strip()])

                chunk_text = f"File: {filepath}\nType: {meta['type']}\nName: {meta['name']}\n"
                if leading_comments:
                    chunk_text += f"\nComments:\n{leading_comments}\n\n"
                chunk_text += f"SQL:\n{stmt_clean}"

                chunks.append({"text": chunk_text, "metadata": meta})
                continue

            # Special handling for CREATE EXTENSION
            extension_match = re.match(
                r'CREATE\s+EXTENSION\s+(?:IF\s+NOT\s+EXISTS\s+)?(\w+)',
                stmt_clean,
                re.IGNORECASE
            )
            if extension_match:
                extension_name = extension_match.group(1)
                meta["type"] = "create extension"
                meta["name"] = extension_name

                chunk_text = f"File: {filepath}\nType: {meta['type']}\nName: {meta['name']}\n"
                if leading_comments:
                    chunk_text += f"\nComments:\n{leading_comments}\n\n"
                chunk_text += f"SQL:\n{stmt_clean}"

                chunks.append({"text": chunk_text, "metadata": meta})
                continue

            # Special handling for CREATE TRIGGER
            trigger_match = re.match(
                r'CREATE\s+(?:OR\s+REPLACE\s+)?TRIGGER\s+(\w+)',
                stmt_clean,
                re.IGNORECASE
            )
            if trigger_match:
                trigger_name = trigger_match.group(1)
                meta["type"] = "create trigger"
                meta["name"] = trigger_name

                chunk_text = f"File: {filepath}\nType: {meta['type']}\nName: {meta['name']}\n"
                if leading_comments:
                    chunk_text += f"\nComments:\n{leading_comments}\n\n"
                chunk_text += f"SQL:\n{stmt_clean}"

                chunks.append({"text": chunk_text, "metadata": meta})
                continue

            # Try sqlglot for standard SQL
            parsed = sqlglot.parse_one(stmt_clean, read="postgres")
            stmt_type = getattr(parsed, "key", None) or parsed.token_type if hasattr(parsed, "token_type") else None
            meta["type"] = str(stmt_type).lower() if stmt_type else "statement"

            tables = [t.this for t in parsed.find_all(sqlglot.exp.Table)]
            meta["tables"] = ",".join([str(t) for t in tables if isinstance(t, str)])

            ctes = []
            for cte in parsed.find_all(sqlglot.exp.CTE):
                try:
                    alias = cte.alias_or_name
                    if alias:
                        ctes.append(alias)
                except Exception:
                    pass
            meta["ctes"] = ",".join(ctes)

            funcs = []
            for f in parsed.find_all(sqlglot.exp.Func):
                try:
                    name = f.name
                    if name:
                        funcs.append(name)
                except Exception:
                    pass
            meta["functions"] = ",".join(funcs)

            if meta.get("tables") and meta["tables"]:
                meta["name"] = meta["tables"].split(",")[0]
            elif ctes:
                meta["name"] = ctes[0]
            elif funcs:
                meta["name"] = funcs[0]

            chunk_text = f"File: {filepath}\nType: {meta.get('type')}\nName: {meta.get('name')}\n"
            if leading_comments:
                chunk_text += f"\nComments:\n{leading_comments}\n\n"
            chunk_text += f"SQL:\n{stmt_clean}"

            chunks.append({"text": chunk_text, "metadata": meta})
            continue

        except Exception:
            pass

        # Fallback: simple heuristics if sqlglot fails
        try:
            parsed_tok = sqlparse.parse(stmt_clean)[0]
            first_token = parsed_tok.token_first(skip_cm=True)
            stmt_type = first_token.value.upper() if first_token else "UNKNOWN"
        except Exception:
            stmt_type = "UNKNOWN"

        meta["type"] = stmt_type.lower()
        chunk_text = f"File: {filepath}\nType: {meta['type']}\nName: {meta.get('name')}\n"
        if leading_comments:
            chunk_text += f"\nComments:\n{leading_comments}\n\n"
        chunk_text += f"SQL:\n{stmt_clean}"

        chunks.append({"text": chunk_text, "metadata": meta})

    return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)

def parse_sql_file(filepath: Path) -> List[Dict[str, Any]]:
    """Parse SQL file with caching."""
    file_hash = get_file_hash(filepath)
    cached_result = parse_sql_file_cached(str(filepath), file_hash)
    return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]


def include_leading_comments(lines: List[str], stmt_start_line: int, max_context_lines: int = 8) -> str:
    """
    Collect contiguous comment lines immediately above stmt_start_line (1-based).
    Supports single-line comments (--), and block comments (/* ... */).
    Returns a string containing the comment block (no leading comment markers).
    """
    comments: List[str] = []
    idx = stmt_start_line - 2  # convert to 0-based index of the line above the statement
    collected = 0

    # Collect single-line comments (-- ...) and blank lines (include blanks between comments)
    while idx >= 0 and collected < max_context_lines:
        raw = lines[idx]
        stripped = raw.strip()
        if stripped.startswith("--"):
            # remove leading -- and optional space
            comments.insert(0, stripped[2:].lstrip())
            idx -= 1
            collected += 1
            continue
        # block comment end
        if stripped.endswith("*/"):
            # gather entire block
            block_lines = []
            while idx >= 0:
                block_line = lines[idx].rstrip()
                block_lines.insert(0, block_line)
                if block_line.strip().startswith("/*"):
                    break
                idx -= 1
            # remove /* and */ markers and join
            cleaned = []
            for bl in block_lines:
                s = bl.strip()
                if s.startswith("/*"):
                    s = s[2:].lstrip()
                if s.endswith("*/"):
                    s = s[:-2].rstrip()
                cleaned.append(s)
            comments = cleaned + comments
            break
        # blank line - allow if we already have comments (so comments may be separated by one blank)
        if stripped == "":
            if comments:
                comments.insert(0, "")
                idx -= 1
                collected += 1
                continue
            else:
                break
        # otherwise stop when non-comment found
        break

    # join into paragraph
    if comments:
        # remove possible leading/trailing empty lines
        while comments and comments[0] == "":
            comments.pop(0)
        while comments and comments[-1] == "":
            comments.pop()
        return "\n".join(comments).strip()
    return ""


# -----------------------------
# Go AST via helper binary with caching
# -----------------------------
# Replace the parse_go_file_cached function in mcp_codebase.py (around line 685)
# The issue is that the Go binary outputs lowercase JSON keys, but Python expects uppercase

@lru_cache(maxsize=1000)
def parse_go_file_cached(filepath_str: str, file_hash: str) -> tuple:
    """Cached Go parsing with enhanced metadata for graph relationships.
    ALWAYS returns a tuple of (text, metadata_tuple) entries. If the Go helper
    fails or produces no chunks, return a single fallback chunk containing the
    whole file text and minimal metadata.
    """
    filepath = Path(filepath_str)
    chunks = []

    logger.info(f"🐹 Parsing Go file: {filepath}")

    go_helper = Path("./tools/parse_go_ast")
    if not go_helper.is_file():
        logger.warning(f"Go parser binary not found: {go_helper}. Using fallback regex parsing for {filepath}")
        return _parse_go_fallback(filepath_str, file_hash)

    try:
        code = filepath.read_text(encoding="utf-8")
        lines = code.splitlines(keepends=True)

        proc = subprocess.run(
            [str(go_helper), str(filepath)],
            capture_output=True,
            text=True,
            check=True,
            timeout=20
        )
        decls = json.loads(proc.stdout)
        logger.info(f"  AST helper found {len(decls)} declarations")

        for i, d in enumerate(decls):
            # Use lowercase keys (matching the Go JSON output)
            start_line = max(0, d.get("start_line", 1) - 1)
            end_line = d.get("end_line", start_line + 1)

            # Guard the end_line to be at least start_line+1
            if end_line <= start_line:
                end_line = start_line + 1

            # Make slice safe even if end_line > len(lines)
            chunk_code = "".join(lines[start_line:end_line])

            name = d.get("name", "")
            typ = d.get("type", "unknown")
            doc_comment = d.get("doc_comment", "")
            receiver = d.get("receiver", "")
            full_name = d.get("full_name", name)
            fields = d.get("fields", [])
            methods = d.get("methods", [])

            # Skip package declarations (they don't contribute to code graph)
            if typ == "package":
                continue

            # Skip empty or invalid declarations
            if not name or typ == "unknown":
                continue

            metadata = {
                "file": str(filepath),
                "name": name,
                "type": typ,
                "line": start_line + 1,
                "language": "go",
            }

            # Build comprehensive chunk text
            chunk_text = f"File: {filepath}\nType: {typ}\nName: {name}\n"

            # Add receiver information for methods
            if receiver:
                metadata["receiver"] = receiver
                # Extract receiver type without pointer for type matching
                receiver_base = receiver.replace('*', '').strip()
                metadata["receiver_base_type"] = receiver_base
                chunk_text += f"Receiver: {receiver}\n"

            # Add full name if different from name (methods have qualified names)
            if full_name and full_name != name:
                metadata["full_name"] = full_name

            # Add documentation
            if doc_comment:
                metadata["docstring"] = doc_comment
                chunk_text += f"Doc:\n{doc_comment}\n\n"

            # Add fields for structs
            if fields:
                metadata["fields"] = fields
                chunk_text += f"Fields: {', '.join(fields)}\n"

            # Add methods for interfaces
            if methods:
                metadata["interface_methods"] = methods
                chunk_text += f"Methods: {', '.join(methods)}\n"

            chunk_text += f"Code:\n{chunk_code}"

            logger.info(f"    Declaration {i+1}: {typ} {name}")

            chunks.append({
                "text": chunk_text,
                "metadata": metadata
            })

    except subprocess.CalledProcessError as e:
        stderr = e.stderr if e.stderr else "unknown error"
        logger.warning(f"Go AST helper failed for {filepath}: {stderr}")
    except json.JSONDecodeError as e:
        logger.warning(f"Go AST helper output invalid JSON for {filepath}: {e}")
    except Exception as e:
        logger.warning(f"Go AST helper failed for {filepath}: {e}")

    # If parser produced chunks, return them in the (text, metadata_tuple) format
    if chunks:
        logger.info(f"✅ Successfully parsed {len(chunks)} chunks from {filepath}")
        return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)

    # Fallback to regex-based parsing
    logger.warning(f"⚠️  Go AST helper produced no chunks for {filepath}. Trying regex fallback.")
    return _parse_go_fallback(filepath_str, file_hash)


def _parse_go_fallback(filepath_str: str, file_hash: str) -> tuple:
    """Fallback regex-based Go parser when AST helper fails or is unavailable."""
    filepath = Path(filepath_str)
    chunks = []

    logger.info(f"  Using regex-based Go parsing for {filepath}")

    try:
        content = filepath.read_text(encoding="utf-8")
        lines = content.splitlines(keepends=True)

        # Helper function to find matching closing brace
        def find_brace_block_end(start_line: int) -> int:
            """Find the line with the matching closing brace."""
            brace_count = 0
            for i in range(start_line, len(lines)):
                for char in lines[i]:
                    if char == '{':
                        brace_count += 1
                    elif char == '}':
                        brace_count -= 1
                        if brace_count == 0:
                            return i + 1  # Include closing brace line
            return len(lines)

        # Helper function to extract preceding comments
        def extract_comments(start_line: int) -> str:
            """Extract comments immediately before a declaration."""
            comments = []
            for i in range(max(0, start_line - 10), start_line):
                line = lines[i].strip()
                if line.startswith('//'):
                    comments.append(line[2:].strip())
                elif line.startswith('/*'):
                    # Multi-line comment
                    comment_lines = []
                    for j in range(i, start_line):
                        comment_lines.append(lines[j].strip())
                        if '*/' in lines[j]:
                            break
                    comment_text = ' '.join(comment_lines)
                    comment_text = comment_text.replace('/*', '').replace('*/', '').strip()
                    comments.append(comment_text)
                elif line and not line.startswith('package') and not line.startswith('import'):
                    # Non-empty, non-comment line before declaration - stop looking
                    break
            return '\n'.join(comments) if comments else ""

        # Pattern to match Go function declarations (including methods)
        func_pattern = re.compile(
            r'^func\s+(?:\(([^)]+)\)\s+)?(\w+)\s*\([^)]*\)(?:\s*\([^)]*\)|\s+[\w\[\].*]+)?\s*\{',
            re.MULTILINE
        )

        for match in func_pattern.finditer(content):
            receiver_part = match.group(1)  # Receiver (optional)
            func_name = match.group(2)
            start_pos = match.start()
            start_line = content[:start_pos].count('\n')
            end_line = find_brace_block_end(start_line)
            func_code = "".join(lines[start_line:end_line])
            comments = extract_comments(start_line)

            # Parse receiver if present
            receiver_type = ""
            receiver_base_type = ""
            if receiver_part:
                # Extract type from receiver (e.g., "u *User" -> "*User")
                parts = receiver_part.strip().split()
                if len(parts) >= 2:
                    receiver_type = parts[-1]
                elif len(parts) == 1:
                    receiver_type = parts[0]
                receiver_base_type = receiver_type.replace('*', '').strip()

            metadata = {
                "file": str(filepath),
                "name": func_name,
                "type": "method" if receiver_type else "function",
                "line": start_line + 1,
                "language": "go",
            }

            chunk_text = f"File: {filepath}\nType: {metadata['type']}\nName: {func_name}\n"

            # Add receiver information for methods
            if receiver_type:
                metadata["receiver"] = receiver_type
                metadata["receiver_base_type"] = receiver_base_type
                chunk_text += f"Receiver: {receiver_type}\n"

            if comments:
                metadata["docstring"] = comments
                chunk_text += f"Comments:\n{comments}\n\n"

            chunk_text += f"Code:\n{func_code}"

            chunks.append({
                "text": chunk_text,
                "metadata": metadata
            })

        # Pattern to match struct declarations
        struct_pattern = re.compile(r'^type\s+(\w+)\s+struct\s*\{', re.MULTILINE)
        for match in struct_pattern.finditer(content):
            struct_name = match.group(1)
            start_pos = match.start()
            start_line = content[:start_pos].count('\n')
            end_line = find_brace_block_end(start_line)
            struct_code = "".join(lines[start_line:end_line])
            comments = extract_comments(start_line)

            # Extract field names from struct
            field_names = []
            struct_body_match = re.search(r'struct\s*\{([^}]+)\}', struct_code, re.DOTALL)
            if struct_body_match:
                field_lines = struct_body_match.group(1).split('\n')
                for line in field_lines:
                    line = line.strip()
                    if not line or line.startswith('//'):
                        continue
                    # Match field declarations: FieldName Type or FieldName, FieldName2 Type
                    field_match = re.match(r'(\w+(?:\s*,\s*\w+)*)\s+[\w\[\]*.]+', line)
                    if field_match:
                        # Handle multiple fields on one line
                        field_str = field_match.group(1)
                        fields = [f.strip() for f in field_str.split(',')]
                        field_names.extend(fields)
                    else:
                        # Check for embedded fields (just a type, no name)
                        embedded_match = re.match(r'(\*?[\w.]+)\s*(?://|$)', line)
                        if embedded_match and not line.startswith('type'):
                            field_names.append(f"embedded:{embedded_match.group(1)}")

            metadata = {
                "file": str(filepath),
                "name": struct_name,
                "type": "struct",
                "line": start_line + 1,
                "language": "go",
            }

            chunk_text = f"File: {filepath}\nType: struct\nName: {struct_name}\n"

            if comments:
                metadata["docstring"] = comments
                chunk_text += f"Comments:\n{comments}\n\n"

            if field_names:
                metadata["fields"] = field_names
                chunk_text += f"Fields: {', '.join(field_names)}\n"

            chunk_text += f"Code:\n{struct_code}"

            chunks.append({
                "text": chunk_text,
                "metadata": metadata
            })

        # Pattern to match interface declarations
        interface_pattern = re.compile(r'^type\s+(\w+)\s+interface\s*\{', re.MULTILINE)
        for match in interface_pattern.finditer(content):
            interface_name = match.group(1)
            start_pos = match.start()
            start_line = content[:start_pos].count('\n')
            end_line = find_brace_block_end(start_line)
            interface_code = "".join(lines[start_line:end_line])
            comments = extract_comments(start_line)

            # Extract method signatures from interface
            method_names = []
            interface_body_match = re.search(r'interface\s*\{([^}]+)\}', interface_code, re.DOTALL)
            if interface_body_match:
                method_lines = interface_body_match.group(1).split('\n')
                for line in method_lines:
                    line = line.strip()
                    if not line or line.startswith('//'):
                        continue
                    # Match method declarations: MethodName(params) returnType
                    method_match = re.match(r'(\w+)\s*\([^)]*\)', line)
                    if method_match:
                        method_names.append(method_match.group(1))

            metadata = {
                "file": str(filepath),
                "name": interface_name,
                "type": "interface",
                "line": start_line + 1,
                "language": "go",
            }

            chunk_text = f"File: {filepath}\nType: interface\nName: {interface_name}\n"

            if comments:
                metadata["docstring"] = comments
                chunk_text += f"Comments:\n{comments}\n\n"

            if method_names:
                metadata["interface_methods"] = method_names
                chunk_text += f"Methods: {', '.join(method_names)}\n"

            chunk_text += f"Code:\n{interface_code}"

            chunks.append({
                "text": chunk_text,
                "metadata": metadata
            })

        # Pattern to match type aliases and other type declarations
        type_alias_pattern = re.compile(r'^type\s+(\w+)\s+(?!struct|interface)(.+?)(?:\n|$)', re.MULTILINE)
        for match in type_alias_pattern.finditer(content):
            type_name = match.group(1)
            type_def = match.group(2).strip()
            start_pos = match.start()
            start_line = content[:start_pos].count('\n')

            # For simple type aliases, just capture the line
            type_code = lines[start_line] if start_line < len(lines) else match.group(0)
            comments = extract_comments(start_line)

            metadata = {
                "file": str(filepath),
                "name": type_name,
                "type": "type",
                "line": start_line + 1,
                "language": "go",
            }

            chunk_text = f"File: {filepath}\nType: type alias\nName: {type_name}\n"

            if comments:
                metadata["docstring"] = comments
                chunk_text += f"Comments:\n{comments}\n\n"

            chunk_text += f"Definition: {type_def}\n"
            chunk_text += f"Code:\n{type_code}"

            chunks.append({
                "text": chunk_text,
                "metadata": metadata
            })

        # Pattern to match const and var declarations
        const_var_pattern = re.compile(r'^(const|var)\s+(?:\(|(\w+))', re.MULTILINE)
        for match in const_var_pattern.finditer(content):
            decl_type = match.group(1)  # 'const' or 'var'
            single_name = match.group(2)  # Name if single declaration
            start_pos = match.start()
            start_line = content[:start_pos].count('\n')

            if single_name:
                # Single declaration
                var_code = lines[start_line] if start_line < len(lines) else match.group(0)
                comments = extract_comments(start_line)

                metadata = {
                    "file": str(filepath),
                    "name": single_name,
                    "type": decl_type,
                    "line": start_line + 1,
                    "language": "go",
                }

                chunk_text = f"File: {filepath}\nType: {decl_type}\nName: {single_name}\n"

                if comments:
                    metadata["docstring"] = comments
                    chunk_text += f"Comments:\n{comments}\n\n"

                chunk_text += f"Code:\n{var_code}"

                chunks.append({
                    "text": chunk_text,
                    "metadata": metadata
                })
            else:
                # Block declaration (const ( ... ) or var ( ... ))
                end_line = find_brace_block_end(start_line) if '(' in lines[start_line] else start_line + 1
                # For simplicity in fallback, just capture the whole block
                block_code = "".join(lines[start_line:end_line])
                comments = extract_comments(start_line)

                # Extract individual names from block
                names = re.findall(r'^\s*(\w+)\s+', block_code, re.MULTILINE)

                if names:
                    for name in names:
                        metadata = {
                            "file": str(filepath),
                            "name": name,
                            "type": decl_type,
                            "line": start_line + 1,
                            "language": "go",
                        }

                        chunk_text = f"File: {filepath}\nType: {decl_type}\nName: {name}\n"

                        if comments:
                            metadata["docstring"] = comments

                        chunk_text += f"Code:\n{block_code}"

                        chunks.append({
                            "text": chunk_text,
                            "metadata": metadata
                        })

    except Exception as e:
        logger.warning(f"Regex-based Go parsing failed for {filepath}: {e}")

    # Return parsed chunks or final fallback
    if chunks:
        logger.info(f"✅ Regex fallback parsed {len(chunks)} chunks from {filepath}")
        return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)

    # Final fallback: emit a whole-file chunk
    logger.warning(f"❌ No chunks created from Go parsing for {filepath}. Emitting fallback whole-file chunk.")
    try:
        full_code = filepath.read_text(encoding="utf-8")
    except Exception:
        full_code = ""

    fallback_meta = {
        "file": str(filepath),
        "name": filepath.name,
        "type": "file",
        "line": 1,
        "language": "go",
    }
    return ((full_code, tuple(fallback_meta.items())),)


def find_brace_block_end_go(lines: List[str], start_line: int) -> int:
    """Find the end of a Go brace block starting at start_line."""
    brace_count = 0
    in_block = False

    for i in range(start_line, len(lines)):
        line = lines[i]

        for char in line:
            if char == '{':
                brace_count += 1
                in_block = True
            elif char == '}':
                brace_count -= 1
                if in_block and brace_count == 0:
                    return i

    return len(lines) - 1


def extract_go_comments(lines: List[str], func_line_index: int, max_lines: int = 10) -> str:
    """Extract leading comments above a Go function/struct/interface."""
    comments = []
    idx = func_line_index - 1

    while idx >= 0 and len(comments) < max_lines:
        line = lines[idx].strip()

        if line.startswith('//'):
            # Remove // and trim
            comments.insert(0, line[2:].strip())
            idx -= 1
        elif line.startswith('/*') or '*/' in line:
            # Block comment - collect entire block
            if '*/' in line and not line.startswith('/*'):
                # End of block, go backwards to find start
                block_lines = [line]
                idx -= 1
                while idx >= 0:
                    block_line = lines[idx].strip()
                    block_lines.insert(0, block_line)
                    if block_line.startswith('/*'):
                        break
                    idx -= 1
                # Clean up block comment markers
                block_text = ' '.join(block_lines)
                block_text = block_text.replace('/*', '').replace('*/', '').replace('*', '').strip()
                comments.insert(0, block_text)
                idx -= 1
            else:
                # Single line block comment
                clean = line.replace('/*', '').replace('*/', '').strip()
                comments.insert(0, clean)
                idx -= 1
        elif line == '':
            # Allow one blank line
            if comments:
                idx -= 1
            else:
                break
        else:
            break

    # Clean up
    while comments and comments[0] == '':
        comments.pop(0)
    while comments and comments[-1] == '':
        comments.pop(-1)

    return '\n'.join(comments)

def parse_go_file(filepath: Path) -> List[Dict]:
    """Parse Go file with caching."""
    file_hash = get_file_hash(filepath)
    cached_result = parse_go_file_cached(str(filepath), file_hash)
    if not cached_result:
        return chunk_text_file(filepath, "go")
    return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]

# -----------------------------
# Java parsing (javalang + manual body extraction) with caching
# -----------------------------
@lru_cache(maxsize=1000)
def parse_java_file_cached(filepath_str: str, file_hash: str) -> tuple:
    """Cached Java parsing with enhanced metadata for graph relationships.
    ALWAYS returns a tuple of (text, metadata_tuple) entries. If parsing
    fails or produces no chunks, return a single fallback chunk containing
    the whole file text and minimal metadata.
    """
    filepath = Path(filepath_str)
    chunks = []

    logger.info(f"☕ Parsing Java file: {filepath}")

    try:
        code = filepath.read_text(encoding="utf-8")
        lines = code.splitlines(keepends=True)

        try:
            tree = javalang.parse.parse(code)
        except javalang.parser.JavaSyntaxError as e:
            logger.warning(f"Java syntax error in {filepath}: {e}. Using fallback chunk.")
            fallback_meta = {
                "file": str(filepath),
                "name": filepath.name,
                "type": "file",
                "line": 1,
                "language": "java",
            }
            return ((code, tuple(fallback_meta.items())),)

        def find_brace_block_end(start_idx: int) -> int:
            """Find the closing brace for a block starting at start_idx."""
            brace_count = 0
            for i in range(start_idx, len(lines)):
                for char in lines[i]:
                    if char == '{':
                        brace_count += 1
                    elif char == '}':
                        brace_count -= 1
                        if brace_count == 0:
                            return i + 1  # Include the closing brace line
            return len(lines)  # If no closing brace found, take rest of file

        def extract_docstring(node) -> str | None:
            """Extract documentation/Javadoc from a node."""
            if hasattr(node, 'documentation') and node.documentation:
                return node.documentation.strip()
            return None

        def get_visibility(modifiers) -> str:
            """Extract visibility from modifiers set."""
            if not modifiers:
                return "package-private"
            if 'public' in modifiers:
                return "public"
            if 'protected' in modifiers:
                return "protected"
            if 'private' in modifiers:
                return "private"
            return "package-private"

        def extract_imports(tree) -> list[str]:
            """Extract all imports from the compilation unit."""
            imports = []
            if hasattr(tree, 'imports') and tree.imports:
                for imp in tree.imports:
                    imports.append(imp.path)
            return imports

        def extract_annotations(node) -> list[str]:
            """Extract annotations from a node."""
            annotations = []
            if hasattr(node, 'annotations') and node.annotations:
                for anno in node.annotations:
                    if hasattr(anno, 'name'):
                        annotations.append(anno.name)
            return annotations

        def type_to_string(type_obj) -> str:
            """Convert a javalang type object to string."""
            if not type_obj:
                return "void"
            if hasattr(type_obj, 'name'):
                base_type = type_obj.name
                # Handle generic types
                if hasattr(type_obj, 'arguments') and type_obj.arguments:
                    args = [type_to_string(arg.type) if hasattr(arg, 'type') else str(arg)
                           for arg in type_obj.arguments]
                    return f"{base_type}<{', '.join(args)}>"
                # Handle array dimensions
                if hasattr(type_obj, 'dimensions') and type_obj.dimensions:
                    return base_type + "[]" * len(type_obj.dimensions)
                return base_type
            return str(type_obj)

        # Extract package and imports once
        package_name = tree.package.name if hasattr(tree, 'package') and tree.package else None
        imports = extract_imports(tree)

        decl_count = 0

        for path, node in tree:
            metadata = None
            chunk_text = None
            start_line = 0
            end_line = 0

            # CLASS DECLARATION
            if isinstance(node, javalang.tree.ClassDeclaration):
                start_line = node.position.line - 1 if node.position else 0
                end_line = find_brace_block_end(start_line)
                chunk_text = "".join(lines[start_line:end_line])

                extends_class = node.extends.name if node.extends else None
                implements_interfaces = []
                if hasattr(node, 'implements') and node.implements:
                    implements_interfaces = [type_to_string(impl) for impl in node.implements]

                metadata = {
                    "file": str(filepath),
                    "name": node.name,
                    "type": "class",
                    "line": start_line + 1,
                    "language": "java",
                    "visibility": get_visibility(node.modifiers),
                    "package": package_name,
                    "imports": imports,
                    "extends": extends_class,
                    "implements": implements_interfaces,
                    "annotations": extract_annotations(node),
                }

                # Extract generics
                if hasattr(node, 'type_parameters') and node.type_parameters:
                    metadata["generics"] = [tp.name for tp in node.type_parameters]

                # Check for abstract/final/static modifiers
                if node.modifiers:
                    if 'abstract' in node.modifiers:
                        metadata["is_abstract"] = True
                    if 'final' in node.modifiers:
                        metadata["is_final"] = True
                    if 'static' in node.modifiers:
                        metadata["is_static"] = True

                docstring = extract_docstring(node)
                if docstring:
                    metadata["docstring"] = docstring

            # METHOD DECLARATION
            elif isinstance(node, javalang.tree.MethodDeclaration):
                if not hasattr(node, 'position') or not node.position:
                    continue

                start_line = node.position.line - 1
                end_line = find_brace_block_end(start_line)
                chunk_text = "".join(lines[start_line:end_line])

                # Find enclosing class
                class_name = None
                for p in reversed(path):
                    if isinstance(p, javalang.tree.ClassDeclaration):
                        class_name = p.name
                        break

                return_type = type_to_string(node.return_type)
                parameters = []
                if node.parameters:
                    for param in node.parameters:
                        param_type = type_to_string(param.type)
                        param_str = f"{param.name}: {param_type}"
                        if param.varargs:
                            param_str = f"{param.name}: {param_type}..."
                        parameters.append(param_str)

                metadata = {
                    "file": str(filepath),
                    "name": node.name,
                    "type": "method",
                    "class": class_name,
                    "line": start_line + 1,
                    "language": "java",
                    "visibility": get_visibility(node.modifiers),
                    "package": package_name,
                    "return_type": return_type,
                    "parameters": parameters,
                    "annotations": extract_annotations(node),
                }

                # Extract generics
                if hasattr(node, 'type_parameters') and node.type_parameters:
                    metadata["generics"] = [tp.name for tp in node.type_parameters]

                # Check modifiers
                if node.modifiers:
                    if 'abstract' in node.modifiers:
                        metadata["is_abstract"] = True
                    if 'final' in node.modifiers:
                        metadata["is_final"] = True
                    if 'static' in node.modifiers:
                        metadata["is_static"] = True
                    if 'synchronized' in node.modifiers:
                        metadata["is_synchronized"] = True
                    if 'native' in node.modifiers:
                        metadata["is_native"] = True

                docstring = extract_docstring(node)
                if docstring:
                    metadata["docstring"] = docstring

            # FIELD DECLARATION
            elif isinstance(node, javalang.tree.FieldDeclaration):
                if not hasattr(node, 'position') or not node.position:
                    continue

                start_line = node.position.line - 1
                end_line = start_line + 1
                chunk_text = lines[start_line] if start_line < len(lines) else ""

                field_type = type_to_string(node.type)

                # Process each declarator (there can be multiple fields on one line)
                for declarator in node.declarators:
                    field_metadata = {
                        "file": str(filepath),
                        "name": declarator.name,
                        "type": "field",
                        "line": start_line + 1,
                        "language": "java",
                        "visibility": get_visibility(node.modifiers),
                        "package": package_name,
                        "field_type": field_type,
                        "annotations": extract_annotations(node),
                    }

                    # Check modifiers
                    if node.modifiers:
                        if 'static' in node.modifiers:
                            field_metadata["is_static"] = True
                        if 'final' in node.modifiers:
                            field_metadata["is_final"] = True
                        if 'volatile' in node.modifiers:
                            field_metadata["is_volatile"] = True
                        if 'transient' in node.modifiers:
                            field_metadata["is_transient"] = True

                    docstring = extract_docstring(node)
                    if docstring:
                        field_metadata["docstring"] = docstring

                    field_text = f"File: {filepath}\nType: Field\nName: {declarator.name}\n"
                    field_text += f"Field Type: {field_type}\nVisibility: {field_metadata['visibility']}\n"
                    if field_metadata.get("annotations"):
                        field_text += f"Annotations: {', '.join(field_metadata['annotations'])}\n"
                    field_text += f"Code:\n{chunk_text}"

                    chunks.append({
                        "text": field_text,
                        "metadata": field_metadata
                    })
                    decl_count += 1

                # Skip the standard processing since we handled fields inline
                continue

            # CONSTRUCTOR DECLARATION
            elif isinstance(node, javalang.tree.ConstructorDeclaration):
                if not hasattr(node, 'position') or not node.position:
                    continue

                start_line = node.position.line - 1
                end_line = find_brace_block_end(start_line)
                chunk_text = "".join(lines[start_line:end_line])

                # Find enclosing class
                class_name = None
                for p in reversed(path):
                    if isinstance(p, javalang.tree.ClassDeclaration):
                        class_name = p.name
                        break

                parameters = []
                if node.parameters:
                    for param in node.parameters:
                        param_type = type_to_string(param.type)
                        parameters.append(f"{param.name}: {param_type}")

                metadata = {
                    "file": str(filepath),
                    "name": f"{class_name}",
                    "type": "constructor",
                    "class": class_name,
                    "line": start_line + 1,
                    "language": "java",
                    "visibility": get_visibility(node.modifiers),
                    "package": package_name,
                    "parameters": parameters,
                    "annotations": extract_annotations(node),
                }

                docstring = extract_docstring(node)
                if docstring:
                    metadata["docstring"] = docstring

            # ENUM DECLARATION
            elif isinstance(node, javalang.tree.EnumDeclaration):
                start_line = node.position.line - 1 if node.position else 0
                end_line = find_brace_block_end(start_line)
                chunk_text = "".join(lines[start_line:end_line])

                # Extract enum constants
                variants = []
                if hasattr(node, 'body') and node.body:
                    for item in node.body:
                        if isinstance(item, javalang.tree.EnumConstantDeclaration):
                            variants.append(item.name)

                implements_interfaces = []
                if hasattr(node, 'implements') and node.implements:
                    implements_interfaces = [type_to_string(impl) for impl in node.implements]

                metadata = {
                    "file": str(filepath),
                    "name": node.name,
                    "type": "enum",
                    "line": start_line + 1,
                    "language": "java",
                    "visibility": get_visibility(node.modifiers),
                    "package": package_name,
                    "implements": implements_interfaces,
                    "variants": variants,
                    "annotations": extract_annotations(node),
                }

                docstring = extract_docstring(node)
                if docstring:
                    metadata["docstring"] = docstring

            # INTERFACE DECLARATION
            elif isinstance(node, javalang.tree.InterfaceDeclaration):
                start_line = node.position.line - 1 if node.position else 0
                end_line = find_brace_block_end(start_line)
                chunk_text = "".join(lines[start_line:end_line])

                extends_interfaces = []
                if hasattr(node, 'extends') and node.extends:
                    extends_interfaces = [type_to_string(ext) for ext in node.extends]

                # Extract method signatures
                methods = []
                if hasattr(node, 'body') and node.body:
                    for item in node.body:
                        if isinstance(item, javalang.tree.MethodDeclaration):
                            methods.append(item.name)

                metadata = {
                    "file": str(filepath),
                    "name": node.name,
                    "type": "interface",
                    "line": start_line + 1,
                    "language": "java",
                    "visibility": get_visibility(node.modifiers),
                    "package": package_name,
                    "extends": extends_interfaces,
                    "methods": methods,
                    "annotations": extract_annotations(node),
                }

                # Extract generics
                if hasattr(node, 'type_parameters') and node.type_parameters:
                    metadata["generics"] = [tp.name for tp in node.type_parameters]

                docstring = extract_docstring(node)
                if docstring:
                    metadata["docstring"] = docstring

            # ANNOTATION DECLARATION
            elif isinstance(node, javalang.tree.AnnotationDeclaration):
                start_line = node.position.line - 1 if node.position else 0
                end_line = find_brace_block_end(start_line)
                chunk_text = "".join(lines[start_line:end_line])

                metadata = {
                    "file": str(filepath),
                    "name": node.name,
                    "type": "annotation",
                    "line": start_line + 1,
                    "language": "java",
                    "visibility": get_visibility(node.modifiers),
                    "package": package_name,
                }

                docstring = extract_docstring(node)
                if docstring:
                    metadata["docstring"] = docstring

            # Build chunk text and add to chunks list
            if metadata and chunk_text is not None:
                text_parts = [f"File: {filepath}"]
                text_parts.append(f"Type: {metadata['type']}")
                text_parts.append(f"Name: {metadata['name']}")

                if metadata.get('visibility'):
                    text_parts.append(f"Visibility: {metadata['visibility']}")
                if metadata.get('package'):
                    text_parts.append(f"Package: {metadata['package']}")
                if metadata.get('extends'):
                    text_parts.append(f"Extends: {metadata['extends']}")
                if metadata.get('implements'):
                    text_parts.append(f"Implements: {', '.join(metadata['implements'])}")
                if metadata.get('return_type'):
                    text_parts.append(f"Returns: {metadata['return_type']}")
                if metadata.get('parameters'):
                    text_parts.append(f"Parameters: {', '.join(metadata['parameters'])}")
                if metadata.get('annotations'):
                    text_parts.append(f"Annotations: {', '.join(metadata['annotations'])}")
                if metadata.get('generics'):
                    text_parts.append(f"Generics: {', '.join(metadata['generics'])}")

                text_parts.append(f"Code:\n{chunk_text}")

                chunks.append({
                    "text": "\n".join(text_parts),
                    "metadata": metadata
                })
                decl_count += 1

        logger.info(f"  Parsed {decl_count} declarations from {filepath}")

    except Exception as e:
        logger.warning(f"Failed to parse Java file {filepath}: {e}")

    # Return parsed chunks or fallback
    if chunks:
        logger.info(f"✅ Successfully parsed {len(chunks)} chunks from {filepath}")
        return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)

    # Fallback: emit a whole-file chunk
    logger.warning(f"❌ No chunks created from Java AST for {filepath}. Emitting fallback whole-file chunk.")
    try:
        full_code = filepath.read_text(encoding="utf-8")
    except Exception:
        full_code = ""

    fallback_meta = {
        "file": str(filepath),
        "name": filepath.name,
        "type": "file",
        "line": 1,
        "language": "java",
    }
    return ((full_code, tuple(fallback_meta.items())),)

def parse_java_file(filepath: Path) -> List[Dict]:
    """Parse Java source file with caching."""
    fileHash = hash(filepath.read_bytes())
    cached_result = parse_java_file_cached(str(filepath), fileHash)
    return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]

# -----------------------------
# Svelte support with caching
# -----------------------------
@lru_cache(maxsize=500)
def parse_svelte_file_cached(filepath_str: str, file_hash: str) -> tuple:
    """Cached Svelte parsing with enhanced multi-language graph relationships.
    ALWAYS returns a tuple of (text, metadata_tuple) entries. If parsing
    fails or produces no chunks, return a single fallback chunk containing
    the whole file text and minimal metadata.
    """
    filepath = Path(filepath_str)
    chunks = []

    logger.info(f"🎨 Parsing Svelte file: {filepath}")

    try:
        text = filepath.read_text(encoding="utf-8")
        lines = text.splitlines(keepends=True)

        # Extract component name from filename for graph relationships
        component_name = filepath.stem
        if component_name and component_name[0].islower():
            component_name = component_name[0].upper() + component_name[1:]  # PascalCase

        # Track component props and context for graph relationships
        component_props = []
        component_stores = []
        component_imports = []
        reactive_declarations = []

        # Parse script blocks
        script_matches = list(re.finditer(r"<script(?:\s+([^>]*?))?>(.*?)</script>", text, flags=re.DOTALL))
        script_idx = 0

        for script_match in script_matches:
            script_attrs = script_match.group(1) or ""
            script_content = script_match.group(2)
            script_start_pos = script_match.start(2)
            script_start_line = text[:script_start_pos].count("\n") + 1

            # Determine script context
            is_module = 'context="module"' in script_attrs or "context='module'" in script_attrs
            is_typescript = 'lang="ts"' in script_attrs or "lang='ts'" in script_attrs
            script_type = "script_module" if is_module else "script"
            script_lang = "typescript" if is_typescript else "javascript"

            logger.info(f"  Processing {script_type} block (lang: {script_lang})")

            # Try TypeScript/JavaScript parser
            ts_helper = Path("./tools/parse_ts.js")
            parsed_with_helper = False

            if ts_helper.is_file():
                tmp = None
                try:
                    # Create temporary file for parsing
                    tmp = filepath.parent / f".tmp_{filepath.name}_{script_idx}.ts"
                    tmp.write_text(script_content, encoding="utf-8")

                    proc = subprocess.run(
                        ["node", str(ts_helper), str(tmp)],
                        capture_output=True,
                        text=True,
                        check=True,
                        timeout=15
                    )
                    decls = json.loads(proc.stdout)
                    script_lines = script_content.splitlines(keepends=True)

                    logger.info(f"    TS helper found {len(decls)} declarations")

                    for i, d in enumerate(decls):
                        start_rel = max(0, d.get("start_line", 1) - 1)
                        end_rel = d.get("end_line", start_rel + 1)

                        # Guard the end_line
                        if end_rel <= start_rel:
                            end_rel = start_rel + 1

                        chunk_code = "".join(script_lines[start_rel:end_rel])
                        doc = d.get("doc_comment", "").strip()
                        abs_start = script_start_line + start_rel

                        decl_name = d.get("name", "")
                        decl_type = d.get("type", "unknown")

                        # Skip invalid declarations
                        if not decl_name or decl_type == "unknown":
                            continue

                        # Analyze declaration for Svelte-specific patterns
                        is_exported = "export" in chunk_code
                        is_prop = False
                        is_reactive = False
                        is_store = False

                        if decl_type == "variable" and not is_module:
                            # Check for prop: export let prop = ...
                            if re.search(r'\bexport\s+let\s+' + re.escape(decl_name), chunk_code):
                                component_props.append(decl_name)
                                is_prop = True
                                decl_type = "prop"

                            # Check for reactive declaration: $: reactiveVar = ...
                            if chunk_code.strip().startswith('$:'):
                                reactive_declarations.append(decl_name)
                                is_reactive = True
                                decl_type = "reactive"

                            # Check for store: $storeName
                            if decl_name.startswith('$'):
                                component_stores.append(decl_name)
                                is_store = True

                        # Track imports
                        if "import" in chunk_code and "from" in chunk_code:
                            import_match = re.search(r'from\s+["\']([^"\']+)["\']', chunk_code)
                            if import_match:
                                component_imports.append(import_match.group(1))

                        # Build chunk text
                        chunk_text = f"File: {filepath}\nType: {decl_type}\nName: {decl_name}\n"
                        chunk_text += f"Component: {component_name}\nScript Type: {script_type}\n"
                        if doc:
                            chunk_text += f"Doc:\n{doc}\n\n"
                        chunk_text += f"Code:\n{chunk_code}"

                        metadata = {
                            "file": str(filepath),
                            "name": decl_name,
                            "type": decl_type,
                            "line": abs_start,
                            "language": script_lang,
                            "block_type": script_type,
                            "component_name": component_name,
                        }

                        # Add Svelte-specific metadata
                        if is_prop:
                            metadata["is_exported"] = True
                            metadata["component_prop"] = True
                        if is_reactive:
                            metadata["is_reactive"] = True
                        if is_store:
                            metadata["is_store"] = True
                        if is_exported and not is_prop:
                            metadata["is_exported"] = True

                        # Function-specific metadata
                        if decl_type == "function":
                            if "createEventDispatcher" in chunk_code:
                                metadata["is_event_dispatcher"] = True
                            if "dispatch(" in chunk_code:
                                metadata["dispatches_events"] = True

                        if doc:
                            metadata["docstring"] = doc

                        chunks.append({
                            "text": chunk_text,
                            "metadata": metadata
                        })

                        logger.info(f"      Declaration {i+1}: {decl_type} {decl_name}")

                    parsed_with_helper = True

                except subprocess.CalledProcessError as e:
                    logger.warning(f"    TS parser subprocess failed for {filepath}: {e.stderr}")
                except json.JSONDecodeError as e:
                    logger.warning(f"    TS parser output invalid JSON for {filepath}: {e}")
                except Exception as e:
                    logger.warning(f"    TS parser failed for {filepath}: {e}")
                finally:
                    if tmp and tmp.exists():
                        tmp.unlink(missing_ok=True)

            # Fallback: Create basic script chunk if parser failed
            if not parsed_with_helper:
                logger.info(f"    Using fallback script chunk")
                chunks.append(_make_script_chunk(
                    filepath, script_content, script_start_line,
                    script_type, component_name, script_lang
                ))

            script_idx += 1

        # Parse style blocks
        style_matches = list(re.finditer(r"<style(?:\s+([^>]*?))?>(.*?)</style>", text, flags=re.DOTALL))

        for style_idx, style_match in enumerate(style_matches):
            style_attrs = style_match.group(1) or ""
            style_content = style_match.group(2)
            style_start_pos = style_match.start(2)
            style_start_line = text[:style_start_pos].count("\n") + 1

            # Determine style language
            is_scss = 'lang="scss"' in style_attrs or "lang='scss'" in style_attrs
            is_less = 'lang="less"' in style_attrs or "lang='less'" in style_attrs
            is_scoped = "scoped" not in style_attrs  # Svelte styles are scoped by default

            style_lang = "scss" if is_scss else ("less" if is_less else "css")

            # Extract CSS classes and selectors
            css_classes = list(set(re.findall(r'\.([a-zA-Z][\w-]*)\s*[{,:]', style_content)))
            css_ids = list(set(re.findall(r'#([a-zA-Z][\w-]*)\s*[{,:]', style_content)))
            css_selectors = list(set(re.findall(r'^([a-zA-Z][\w-]*)\s*{', style_content, re.MULTILINE)))

            # Extract CSS custom properties (variables)
            css_vars = list(set(re.findall(r'--([a-zA-Z][\w-]*)', style_content)))

            chunk_text = f"File: {filepath}\nType: style\nComponent: {component_name}\n"
            chunk_text += f"Language: {style_lang}\nScoped: {is_scoped}\n"
            if css_classes:
                chunk_text += f"CSS Classes: {', '.join(css_classes)}\n"
            if css_ids:
                chunk_text += f"CSS IDs: {', '.join(css_ids)}\n"
            if css_vars:
                chunk_text += f"CSS Variables: {', '.join(css_vars)}\n"
            chunk_text += f"Style Content:\n{style_content}"

            chunks.append({
                "text": chunk_text,
                "metadata": {
                    "file": str(filepath),
                    "name": f"style_block_{style_idx}",
                    "type": "style",
                    "line": style_start_line,
                    "language": style_lang,
                    "block_type": "style",
                    "component_name": component_name,
                    "css_classes": css_classes,
                    "css_ids": css_ids,
                    "css_selectors": css_selectors,
                    "css_variables": css_vars,
                    "is_scoped": is_scoped,
                }
            })

        # Parse markup (template)
        markup = text
        # Remove script blocks
        markup = re.sub(r"<script[^>]*>.*?</script>", "", markup, flags=re.DOTALL)
        # Remove style blocks
        markup = re.sub(r"<style[^>]*>.*?</style>", "", markup, flags=re.DOTALL)
        markup = markup.strip()

        if markup:
            # Extract Svelte-specific markup features
            used_components = list(set(re.findall(r'<([A-Z][a-zA-Z0-9]*)', markup)))
            prop_bindings = re.findall(r'\b(\w+)={([^}]+)}', markup)
            event_handlers = list(set(re.findall(r'on:(\w+)=', markup)))

            # Extract slots
            slots = list(set(re.findall(r'<slot\s+name="([^"]+)"', markup)))
            has_default_slot = '<slot' in markup and 'name=' not in markup[:markup.find('<slot') + 100]

            # Extract stores used in template
            template_stores = list(set(re.findall(r'\$([a-zA-Z_]\w*)', markup)))

            # Extract directives
            use_directives = list(set(re.findall(r'use:(\w+)', markup)))
            transition_directives = list(set(re.findall(r'(?:transition|in|out):(\w+)', markup)))

            # Extract bind directives
            bind_directives = list(set(re.findall(r'bind:(\w+)', markup)))

            # Extract reactive statements in template
            reactive_statements = re.findall(r'\{([^}]*\$:\s*[^}]*)\}', markup)

            chunk_text = f"File: {filepath}\nType: markup\nComponent: {component_name}\n"

            if used_components:
                chunk_text += f"Uses Components: {', '.join(used_components)}\n"
            if prop_bindings:
                bindings_preview = [f"{prop}={{{val}}}" for prop, val in prop_bindings[:5]]
                chunk_text += f"Prop Bindings: {', '.join(bindings_preview)}\n"
            if event_handlers:
                chunk_text += f"Event Handlers: {', '.join(event_handlers)}\n"
            if slots or has_default_slot:
                slot_info = []
                if has_default_slot:
                    slot_info.append("default")
                slot_info.extend(slots)
                chunk_text += f"Slots: {', '.join(slot_info)}\n"
            if template_stores:
                chunk_text += f"Template Stores: {', '.join(template_stores)}\n"
            if use_directives:
                chunk_text += f"Use Directives: {', '.join(use_directives)}\n"
            if transition_directives:
                chunk_text += f"Transitions: {', '.join(transition_directives)}\n"
            if bind_directives:
                chunk_text += f"Bind Directives: {', '.join(bind_directives)}\n"

            chunk_text += f"Markup:\n{markup}"

            chunks.append({
                "text": chunk_text,
                "metadata": {
                    "file": str(filepath),
                    "name": "markup",
                    "type": "markup",
                    "line": 1,
                    "language": "svelte",
                    "block_type": "markup",
                    "component_name": component_name,
                    "used_components": used_components,
                    "prop_bindings": [prop for prop, _ in prop_bindings],
                    "event_handlers": event_handlers,
                    "component_props": component_props,
                    "reactive_vars": reactive_declarations,
                    "stores": component_stores,
                    "template_stores": template_stores,
                    "imports": component_imports,
                    "slots": slots,
                    "has_default_slot": has_default_slot,
                    "use_directives": use_directives,
                    "transitions": transition_directives,
                    "bind_directives": bind_directives,
                }
            })

    except Exception as e:
        logger.warning(f"Svelte parsing failed for {filepath}: {e}")

    # Return parsed chunks or fallback
    if chunks:
        logger.info(f"✅ Successfully parsed {len(chunks)} chunks from {filepath}")
        return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)

    # Fallback: emit a whole-file chunk
    logger.warning(f"❌ No chunks created from Svelte parsing for {filepath}. Emitting fallback whole-file chunk.")
    try:
        full_code = filepath.read_text(encoding="utf-8")
    except Exception:
        full_code = ""

    fallback_meta = {
        "file": str(filepath),
        "name": filepath.stem,
        "type": "component",
        "line": 1,
        "language": "svelte",
        "component_name": filepath.stem,
    }
    return ((full_code, tuple(fallback_meta.items())),)


def _make_script_chunk(filepath: Path, script_content: str, start_line: int,
                       script_type: str, component_name: str, language: str = "javascript") -> dict:
    """Create a fallback script chunk when TS parser fails."""
    chunk_text = f"File: {filepath}\nType: {script_type}\nComponent: {component_name}\n"
    chunk_text += f"Language: {language}\nCode:\n{script_content}"

    # Extract basic info with regex
    exports = re.findall(r'export\s+(?:let|const|function|class)\s+(\w+)', script_content)
    imports = re.findall(r'from\s+["\']([^"\']+)["\']', script_content)

    metadata = {
        "file": str(filepath),
        "name": script_type,
        "type": script_type,
        "line": start_line,
        "language": language,
        "block_type": script_type,
        "component_name": component_name,
    }

    if exports:
        metadata["exports"] = exports
    if imports:
        metadata["imports"] = imports

    return {
        "text": chunk_text,
        "metadata": metadata
    }

def parse_svelte_file(filepath: Path) -> List[Dict]:
    """Parse Svelte file with caching."""
    file_hash = get_file_hash(filepath)
    cached_result = parse_svelte_file_cached(str(filepath), file_hash)
    if not cached_result:
        return chunk_text_file(filepath, "svelte")
    return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]

# -----------------------------
# Heuristic chunking for other languages
# -----------------------------
def chunk_text_file(filepath: Path, language: str) -> List[Dict]:
    try:
        lines = filepath.read_text(encoding="utf-8").splitlines()
    except Exception as e:
        logger.warning(f"chunk_text_file read failed {filepath}: {e}")
        return []
    chunks = []
    current_chunk = []
    start_line = 0
    for i, ln in enumerate(lines):
        stripped = ln.strip()
        if stripped.startswith(("def ", "func ", "fn ", "class ", "struct ", "impl ", "export ", "const ", "let ")) or stripped.startswith(("//", "#", "--")) or stripped.endswith("{"):
            if current_chunk:
                text = "\n".join(current_chunk)
                chunks.append({
                    "text": f"File: {filepath}\nLanguage: {language}\n\n{text}",
                    "metadata": {"file": str(filepath), "name": f"chunk_{len(chunks)}", "type": "code_block", "line": start_line + 1, "language": language}
                })
                current_chunk = []
            start_line = i
        current_chunk.append(ln)
    if current_chunk:
        text = "\n".join(current_chunk)
        chunks.append({
            "text": f"File: {filepath}\nLanguage: {language}\n\n{text}",
            "metadata": {"file": str(filepath), "name": f"chunk_{len(chunks)}", "type": "code_block", "line": start_line + 1, "language": language}
        })
    return chunks

def parse_rust_file(filepath: Path) -> List[Dict]:
    """Parse Rust file with caching."""
    file_hash = get_file_hash(filepath)
    cached_result = parse_rust_file_cached(str(filepath), file_hash)
    return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]

@lru_cache(maxsize=1000)
def parse_rust_file_cached(filepath_str: str, file_hash: str) -> tuple:
    """Cached Rust parsing with enhanced metadata for graph relationships.
    ALWAYS returns a tuple of (text, metadata_tuple) entries. If the Rust helper
    fails or produces no chunks, return a single fallback chunk containing the
    whole file text and minimal metadata.
    """
    filepath = Path(filepath_str)
    chunks = []

    logger.info(f"🔧 Parsing Rust file: {filepath}")

    rust_helper = Path("./tools/target/release/parse_rust_ast")
    if not rust_helper.is_file():
        logger.warning(f"Rust parser binary not found: {rust_helper}. Using fallback whole-file chunk for {filepath}")
        code = filepath.read_text(encoding="utf-8")
        fallback_meta = {
            "file": str(filepath),
            "name": filepath.name,
            "type": "file",
            "line": 1,
            "language": "rust",
        }
        return ((code, tuple(fallback_meta.items())),)

    try:
        code = filepath.read_text(encoding="utf-8")
        lines = code.splitlines(keepends=True)

        proc = subprocess.run(
            [str(rust_helper), str(filepath)],
            capture_output=True,
            text=True,
            check=True,
            timeout=20
        )
        decls = json.loads(proc.stdout)
        logger.info(f"  AST helper found {len(decls)} declarations")

        for i, d in enumerate(decls):
            start_line = max(0, d.get("start_line", 1) - 1)
            end_line = d.get("end_line", start_line + 1)
            # guard the end_line to be at least start_line+1
            if end_line <= start_line:
                end_line = start_line + 1

            # Make slice safe even if end_line > len(lines)
            chunk_code = "".join(lines[start_line:end_line])

            metadata = {
                "file": str(filepath),
                "name": d.get("name", "") or filepath.name,
                "type": d.get("type_", "") or "unknown",
                "line": start_line + 1,
                "language": "rust",
            }

            # copy other fields defensively
            if d.get("visibility"):
                metadata["visibility"] = d.get("visibility")
            if d.get("is_async"):
                metadata["is_async"] = True
            if d.get("is_unsafe"):
                metadata["is_unsafe"] = True
            if d.get("generics"):
                metadata["generics"] = d.get("generics")
            if d.get("traits"):
                metadata["implements_traits"] = d.get("traits")
            if d.get("fields"):
                metadata["fields"] = d.get("fields")
            if d.get("methods"):
                metadata["methods"] = d.get("methods")
            if d.get("return_type"):
                metadata["return_type"] = d.get("return_type")
            if d.get("parameters"):
                metadata["parameters"] = d.get("parameters")

            logger.info(f"    Declaration {i+1}: {metadata.get('type', 'unknown')} {metadata.get('name', 'unnamed')}")

            chunk_text = f"File: {filepath}\nType: {metadata.get('type')}\nName: {metadata.get('name')}\n"
            if metadata.get("visibility"):
                chunk_text += f"Visibility: {metadata.get('visibility')}\n"
            if metadata.get("is_async"):
                chunk_text += "Async: yes\n"
            if metadata.get("is_unsafe"):
                chunk_text += "Unsafe: yes\n"
            if metadata.get("generics"):
                chunk_text += f"Generics: {', '.join(metadata.get('generics', []))}\n"
            if metadata.get("implements_traits"):
                chunk_text += f"Implements: {', '.join(metadata.get('implements_traits', []))}\n"
            if metadata.get("fields"):
                chunk_text += f"Fields: {', '.join(metadata.get('fields', []))}\n"

            chunk_text += f"Code:\n{chunk_code}"

            chunks.append({
                "text": chunk_text,
                "metadata": metadata
            })

    except Exception as e:
        logger.warning(f"Rust AST helper failed for {filepath}: {e}")

    # If parser produced chunks, return them in the (text, metadata_tuple) format expected by parse_rust_file()
    if chunks:
        logger.info(f"✅ Successfully parsed {len(chunks)} chunks from {filepath}")
        return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)

    # Fallback: emit a whole-file chunk to keep counts consistent
    logger.warning(f"❌ No chunks created from Rust AST for {filepath}. Emitting fallback whole-file chunk.")
    try:
        full_code = filepath.read_text(encoding="utf-8")
    except Exception:
        full_code = ""
    fallback_meta = {
        "file": str(filepath),
        "name": filepath.name,
        "type": "file",
        "line": 1,
        "language": "rust",
    }
    return ((full_code, tuple(fallback_meta.items())),)


@lru_cache(maxsize=1000)
def parse_shell_file_cached(filepath_str: str) -> tuple:
    """Cached shell script parsing. Returns tuple for hashability."""
    filepath = Path(filepath_str)
    try:
        content = filepath.read_text(encoding="utf-8")
        chunks = []
        lines = content.splitlines(keepends=True)

        patterns = [
            r'^(\w+)\s*\(\s*\)\s*\{',
            r'^function\s+(\w+)\s*\{',
            r'^function\s+(\w+)\s*\(\s*\)\s*\{',
            r'^def\s+(\w+)\s*\(\s*\)\s*\{',
        ]

        for i, line in enumerate(lines):
            line_text = line.strip()

            for pattern in patterns:
                match = re.match(pattern, line_text)
                if match:
                    func_name = match.group(1)
                    start_line = i
                    end_line = find_shell_function_end(lines, i)
                    func_code = "".join(lines[start_line:end_line + 1])

                    # Extract leading comments
                    leading_comments = extract_shell_comments(lines, start_line)

                    chunk_text = f"File: {filepath}\nType: function\nName: {func_name}\n"
                    if leading_comments:
                        chunk_text += f"Comments:\n{leading_comments}\n\n"
                    chunk_text += f"Code:\n{func_code}"

                    chunks.append({
                        "text": chunk_text,
                        "metadata": {
                            "file": str(filepath),
                            "name": func_name,
                            "type": "function",
                            "line": start_line + 1,
                            "language": "shell",
                            "has_braces": "{" in line_text
                        }
                    })
                    break  # Move to next line after finding a match

        # If no functions found, fall back to basic chunking
        return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks) if chunks else tuple()

    except Exception as e:
        logger.warning(f"parse_shell_file failed {filepath}: {e}")
        return tuple()

def parse_shell_file(filepath: Path) -> List[Dict]:
    """Enhanced shell script parsing with caching."""
    cached_result = parse_shell_file_cached(str(filepath))
    if not cached_result:
        return chunk_text_file(filepath, "shell")
    return [{"text": text, "metadata": dict(meta)} for text, meta in cached_result]

def extract_shell_comments(lines: List[str], function_line_index: int, max_lines: int = 10) -> str:
    """Extract leading comments above a shell function."""
    comments = []
    idx = function_line_index - 1

    while idx >= 0 and len(comments) < max_lines:
        line = lines[idx].strip()
        if line.startswith('#'):
            comments.insert(0, line[1:].strip())  # Remove the # and trim
            idx -= 1
        elif line == '':
            # Allow one blank line between comments
            if comments:
                comments.insert(0, "")
                idx -= 1
            else:
                break
        else:
            break

    # Clean up: remove leading/trailing empty lines
    while comments and comments[0] == "":
        comments.pop(0)
    while comments and comments[-1] == "":
        comments.pop(-1)

    return "\n".join(comments)

def find_shell_function_end(lines: List[str], start_line_index: int) -> int:
    """
    Find the end of a shell function by tracking brace nesting.

    Args:
        lines: List of file lines
        start_line_index: Starting line index (0-based) of the function

    Returns:
        Line index (0-based) where the function ends
    """
    brace_count = 0
    in_function = False

    for i in range(start_line_index, len(lines)):
        line = lines[i]

        # Count opening and closing braces
        for char in line:
            if char == '{':
                brace_count += 1
                in_function = True
            elif char == '}':
                brace_count -= 1

            # If we've returned to brace_count 0 and we were in a function, this is the end
            if in_function and brace_count == 0:
                return i

        # Special case: shell functions without braces (single line)
        if not in_function and i > start_line_index:
            # Look for function end patterns
            stripped = line.strip()
            if (stripped.startswith(('function ', 'def ')) or
                re.match(r'^\w+\(\s*\)', stripped) or
                stripped.endswith(')') and '()' not in line):
                # New function starting, so previous one ended
                return i - 1

    # If we never found the end, return the last line
    return len(lines) - 1

# -----------------------------
# Detect language helper
# -----------------------------
def detect_language(filepath: Path) -> str:
    ext = filepath.suffix.lower()
    lang_map = {
        ".py": "python", ".go": "go", ".rs": "rust", ".ts": "typescript", ".tsx": "typescript",
        ".js": "javascript", ".jsx": "javascript", ".svelte": "svelte", ".java": "java",
        ".sql": "sql", ".pgsql": "sql", ".sh": "shell"
    }
    return lang_map.get(ext, "unknown")

# -----------------------------
# Indexing codebase
# -----------------------------
def index_codebase(codebase_path: Path) -> Tuple[List[str], List[Dict[str, Any]]]:
    ragignore = load_ragignore(codebase_path)
    all_texts = []
    all_metadata = []
    logger.info(f"Indexing codebase at {codebase_path}")
    for filepath in codebase_path.rglob("*"):
        if not filepath.is_file():
            continue
        rel = filepath.relative_to(codebase_path)
        if ragignore.match_file(str(rel)):
            continue
        if filepath.suffix.lower() not in SUPPORTED_EXTENSIONS:
            continue
        language = detect_language(filepath)
        if language == "python":
            chunks = parse_python_file(filepath)
        elif language == "go":
            chunks = parse_go_file(filepath)
        elif language == "java":
            chunks = parse_java_file(filepath)
        elif language == "svelte":
            chunks = parse_svelte_file(filepath)
        elif language == "sql":
            chunks = parse_sql_file(filepath)
        elif language == "shell":
            chunks = parse_shell_file(filepath)#
        elif language == "rust":
            chunks = parse_rust_file(filepath)
        else:
            chunks = chunk_text_file(filepath, language)

        for c in chunks:
            all_texts.append(c["text"])
            all_metadata.append(clean_metadata_for_chroma(c["metadata"]))
    logger.info(f"Indexed {len(all_texts)} chunks")

    return all_texts, all_metadata

def clean_metadata_for_chroma(metadata: Dict[str, Any]) -> Dict[str, Any]:
    """Clean metadata to only include str, int, float, bool values (no None)."""
    cleaned = {}
    for key, value in metadata.items():
        if value is None:
            cleaned[key] = ""  # Convert None to empty string
        elif isinstance(value, (str, int, float, bool)):
            cleaned[key] = value
        elif isinstance(value, list):
            # Convert lists to comma-separated strings
            cleaned[key] = ",".join(str(v) for v in value if v is not None)
        else:
            # Convert other types to strings
            cleaned[key] = str(value)
    return cleaned

def build_indexes():
    global vectorstore, bm25, bm25_corpus, chunks_metadata, embeddings, index_build_time
    logger.info("Building indexes with contextual weighting and batch embedding...")
    start_time = time.time()

    # -------------------------------------------------
    # Step 1: Parse codebase and build enhanced graph
    # -------------------------------------------------
    texts, metadatas = index_codebase(CODEBASE_PATH)
    if len(texts) != len(metadatas):
        logger.error(
            "Indexing mismatch: %d texts vs %d metadata entries. Aborting index build.",
            len(texts),
            len(metadatas),
        )
        raise ValueError(f"Indexing produced {len(texts)} texts but {len(metadatas)} metadata entries")

    logger.info("=== Sample metadata inspection ===")
    for i, m in enumerate(metadatas[:5]):
        logger.info(f"Chunk {i}: type={m.get('type')}, name={m.get('name')}")
        logger.info(f"  imports={m.get('imports')}")
        logger.info(f"  calls={m.get('calls')}")
        logger.info(f"  parameters={m.get('parameters')}")

    # Initialize graph
    graph = LocalGraph(root_path=str(CODEBASE_PATH))
    graph.clear()

    # Track all created nodes for second-pass relationships
    created_nodes = {}

    # Track language-specific features for cross-language relationships
    component_registry = {}
    interface_registry = {}
    trait_registry = {}
    class_registry = {}
    package_registry = {}

    logger.info("=== Phase 1: Creating nodes ===")

    # -------------------------------------------------
    # PHASE 1: Create all nodes with rich metadata
    # -------------------------------------------------
    for m in metadatas:
        fpath = m.get("file")
        lang = m.get("language", "unknown")
        node_type = m.get("type", "unknown")
        name = m.get("name", "")

        # Skip empty/invalid nodes
        if not name or node_type == "unknown" or not fpath:
            continue

        # Normalize metadata
        if "implements_traits" in m and m["implements_traits"]:
            m["implements"] = m["implements_traits"]

        # Create file node
        file_node_id = f"file::{fpath}"
        if file_node_id not in graph.graph.nodes:
            graph.add_node(
                file_node_id,
                type="File",
                name=fpath,
                path=fpath,
                lang=lang,
            )
            created_nodes[file_node_id] = {
                "type": "File",
                "name": fpath,
                "path": fpath,
                "lang": lang,
            }

        # Create node ID with language prefix
        node_id = f"{lang}::{node_type}::{fpath}::{name}"

        # Base attributes for all nodes
        base_attrs = {
            "name": name,
            "file": fpath,
            "line": m.get("line"),
            "lang": lang
        }

        # Add docstring/documentation if available
        if m.get("docstring"):
            base_attrs["docstring"] = m.get("docstring")

        # Store for second-pass processing
        created_nodes[node_id] = base_attrs.copy()
        created_nodes[node_id]["type"] = node_type
        created_nodes[node_id]["metadata"] = m

        # ===== FUNCTIONS & METHODS =====
        if node_type in ["function", "method", "async_function", "constructor"]:
            if node_type == "method" and "class" in m:
                class_name = m['class']
                node_id = f"{lang}::{node_type}::{fpath}::{class_name}.{name}"

                created_nodes[node_id] = base_attrs.copy()
                created_nodes[node_id]["type"] = node_type
                created_nodes[node_id]["class_name"] = class_name
                created_nodes[node_id]["metadata"] = m

                graph.add_node(
                    node_id,
                    **base_attrs,
                    type=node_type,
                    class_name=class_name,
                    return_type=m.get("return_type"),
                    parameters=m.get("parameters", []),
                    visibility=m.get("visibility"),
                    is_static=m.get("is_static", False),
                    is_abstract=m.get("is_abstract", False),
                    is_async=m.get("is_async", False),
                    is_generator=m.get("is_generator", False),
                    decorators=m.get("decorators", []),
                    annotations=m.get("annotations", [])
                )

                class_node_id = f"{lang}::class::{fpath}::{class_name}"
                graph.add_edge(class_node_id, node_id, "contains")
            else:
                graph.add_node(
                    node_id,
                    **base_attrs,
                    type=node_type,
                    return_type=m.get("return_type"),
                    parameters=m.get("parameters", []),
                    visibility=m.get("visibility"),
                    is_async=m.get("is_async", False),
                    is_generator=m.get("is_generator", False),
                    is_exported=m.get("is_exported", False),
                    is_property=m.get("is_property", False),
                    decorators=m.get("decorators", []),
                    annotations=m.get("annotations", []),
                    generics=m.get("generics", [])
                )

                graph.add_edge(file_node_id, node_id, "contains")

        # ===== CLASSES =====
        elif node_type == "class":
            graph.add_node(
                node_id,
                **base_attrs,
                type=node_type,
                extends=m.get("extends"),
                implements=m.get("implements", []),
                methods=m.get("methods", []),
                properties=m.get("properties", []),
                fields=m.get("fields", []),
                static_methods=m.get("static_methods", []),
                class_methods=m.get("class_methods", []),
                class_variables=m.get("class_variables", []),
                visibility=m.get("visibility"),
                is_abstract=m.get("is_abstract", False),
                is_final=m.get("is_final", False),
                decorators=m.get("decorators", []),
                annotations=m.get("annotations", []),
                generics=m.get("generics", []),
                metaclass=m.get("metaclass")
            )

            graph.add_edge(file_node_id, node_id, "contains")

            class_registry[name] = {
                "node_id": node_id,
                "lang": lang,
                "file": fpath,
                "metadata": m
            }

        # ===== STRUCTS (Rust/Go) =====
        elif node_type == "struct":
            graph.add_node(
                node_id,
                **base_attrs,
                type=node_type,
                fields=m.get("fields", []),
                visibility=m.get("visibility"),
                generics=m.get("generics", []),
                derives=m.get("derives", [])
            )

            graph.add_edge(file_node_id, node_id, "contains")

            class_registry[name] = {
                "node_id": node_id,
                "lang": lang,
                "file": fpath,
                "metadata": m
            }

        # ===== INTERFACES (TypeScript/Java/Go) =====
        elif node_type == "interface":
            graph.add_node(
                node_id,
                **base_attrs,
                type=node_type,
                extends=m.get("extends", []),
                methods=m.get("methods", []),
                properties=m.get("properties", []),
                interface_methods=m.get("interface_methods", []),
                visibility=m.get("visibility"),
                generics=m.get("generics", [])
            )

            graph.add_edge(file_node_id, node_id, "contains")

            interface_registry[name] = {
                "node_id": node_id,
                "lang": lang,
                "file": fpath,
                "metadata": m
            }

        # ===== TRAITS (Rust) =====
        elif node_type == "trait":
            graph.add_node(
                node_id,
                **base_attrs,
                type=node_type,
                methods=m.get("methods", []),
                visibility=m.get("visibility"),
                generics=m.get("generics", [])
            )

            graph.add_edge(file_node_id, node_id, "contains")

            trait_registry[name] = {
                "node_id": node_id,
                "lang": lang,
                "file": fpath,
                "metadata": m
            }

        # ===== RUST IMPL BLOCKS =====
        elif node_type == "impl":
            graph.add_node(
                node_id,
                **base_attrs,
                type=node_type,
                traits=m.get("traits", []),
                methods=m.get("methods", []),
                generics=m.get("generics", []),
                is_unsafe=m.get("is_unsafe", False)
            )

            graph.add_edge(file_node_id, node_id, "contains")

        # ===== ENUMS =====
        elif node_type == "enum":
            graph.add_node(
                node_id,
                **base_attrs,
                type=node_type,
                variants=m.get("variants", []),
                members=m.get("members", []),
                fields=m.get("fields", []),
                visibility=m.get("visibility")
            )

            graph.add_edge(file_node_id, node_id, "contains")

        # ===== TYPE ALIASES =====
        elif node_type == "type":
            graph.add_node(
                node_id,
                **base_attrs,
                type=node_type,
                type_definition=m.get("type_definition"),
                generics=m.get("generics", [])
            )

            graph.add_edge(file_node_id, node_id, "contains")

        # ===== MODULES (Rust/Go/Python packages) =====
        elif node_type == "module":
            graph.add_node(
                node_id,
                **base_attrs,
                type=node_type,
                visibility=m.get("visibility")
            )
            graph.add_edge(file_node_id, node_id, "contains")

            # Create symbolic module node for cross-references
            module_symbol_id = f"module_symbol::{lang}::{name}"
            if module_symbol_id not in graph.graph.nodes:
                graph.add_node(
                    module_symbol_id,
                    type=name,
                    name=name,
                    file=fpath,
                    lang=lang
                )
                created_nodes[module_symbol_id] = {
                    "type": name,
                    "name": name,
                    "file": fpath,
                    "lang": lang
                }

            graph.add_edge(node_id, module_symbol_id, "defines")

            package_registry[name] = {
                "node_id": node_id,
                "symbol_id": module_symbol_id,
                "lang": lang,
                "file": fpath
            }

        # ===== SVELTE COMPONENTS =====
        elif node_type == "markup" and lang == "svelte":
            component_name = m.get("component_name", name)
            component_node_id = f"svelte::component::{fpath}::{component_name}"

            graph.add_node(
                component_node_id,
                **base_attrs,
                type="component",
                component_name=component_name,
                props=m.get("component_props", []),
                slots=m.get("slots", []),
                has_default_slot=m.get("has_default_slot", False),
                events=m.get("event_handlers", []),
                used_components=m.get("used_components", []),
                stores=m.get("stores", []),
                template_stores=m.get("template_stores", []),
                use_directives=m.get("use_directives", []),
                transitions=m.get("transitions", []),
                bind_directives=m.get("bind_directives", [])
            )

            graph.add_edge(file_node_id, component_node_id, "contains")

            component_registry[component_name] = {
                "node_id": component_node_id,
                "file": fpath,
                "metadata": m
            }

            node_id = component_node_id

        # ===== SVELTE SCRIPTS =====
        elif node_type in ["script", "script_module"] and lang in ["typescript", "javascript"]:
            pass

        # ===== VARIABLES & CONSTANTS =====
        elif node_type in ["variable", "constant", "prop", "reactive", "var", "const"]:
            graph.add_node(
                node_id,
                **base_attrs,
                type=node_type,
                var_type=m.get("var_type"),
                var_kind=m.get("var_kind"),
                is_exported=m.get("is_exported", False),
                is_reactive=m.get("is_reactive", False),
                is_store=m.get("is_store", False),
                value=m.get("value")
            )

            graph.add_edge(file_node_id, node_id, "contains")

        # ===== STYLES (CSS/SCSS/LESS) =====
        elif node_type == "style":
            graph.add_node(
                node_id,
                **base_attrs,
                type=node_type,
                css_classes=m.get("css_classes", []),
                css_ids=m.get("css_ids", []),
                css_variables=m.get("css_variables", []),
                css_selectors=m.get("css_selectors", []),
                is_scoped=m.get("is_scoped", True),
                component_name=m.get("component_name")
            )

            graph.add_edge(file_node_id, node_id, "contains")

        # ===== IMPORTS =====
        elif node_type == "import":
            import_path = m.get("import_path") or m.get("full_name") or m.get("parameters", [name])[0] if m.get("parameters") else name

            if isinstance(import_path, list):
                import_path = import_path[0] if import_path else name

            import_node_id = f"{lang}::import::{fpath}::{import_path}"

            graph.add_node(
                import_node_id,
                **base_attrs,
                type=node_type,
                import_path=import_path,
                imported_names=m.get("imported_names", [])
            )
            graph.add_edge(file_node_id, import_node_id, "contains")

            # Create symbolic nodes - DON'T pass created_nodes, collect new nodes separately
            new_nodes = _create_import_symbolic_nodes(graph, import_node_id, file_node_id, import_path, lang)
            # Add new nodes after iteration completes
            for new_node_id, new_node_data in new_nodes.items():
                if new_node_id not in created_nodes:
                    created_nodes[new_node_id] = new_node_data

            node_id = import_node_id

        # ===== FIELDS (Java standalone fields) =====
        elif node_type == "field":
            graph.add_node(
                node_id,
                **base_attrs,
                type=node_type,
                field_type=m.get("field_type"),
                visibility=m.get("visibility"),
                is_static=m.get("is_static", False),
                is_final=m.get("is_final", False),
                is_volatile=m.get("is_volatile", False),
                annotations=m.get("annotations", [])
            )
            graph.add_edge(file_node_id, node_id, "contains")

        # ===== ANNOTATIONS (Java) =====
        elif node_type == "annotation":
            graph.add_node(
                node_id,
                **base_attrs,
                type=node_type,
                visibility=m.get("visibility")
            )
            graph.add_edge(file_node_id, node_id, "contains")

        # ===== GENERIC FALLBACK =====
        else:
            graph.add_node(node_id, **base_attrs, type=node_type)
            graph.add_edge(file_node_id, node_id, "contains")

    logger.info(f"Phase 1 complete: Created {len(created_nodes)} nodes")
    logger.info(f"  Classes: {len(class_registry)}, Interfaces: {len(interface_registry)}")
    logger.info(f"  Traits: {len(trait_registry)}, Components: {len(component_registry)}")
    logger.info(f"  Packages/Modules: {len(package_registry)}")

    # -------------------------------------------------
    # PHASE 2: Build relationships
    # -------------------------------------------------
    logger.info("=== Phase 2: Building relationships ===")

    # CRITICAL FIX: Create a snapshot of the items to iterate over
    # This prevents "dictionary changed size during iteration" errors
    nodes_snapshot = list(created_nodes.items())

    for node_id, node_data in nodes_snapshot:
        m = node_data.get("metadata", {})
        if not m:
            continue

        node_type = node_data.get("type")
        fpath = node_data.get("file")
        lang = node_data.get("lang")
        name = node_data.get("name")

        # ===== IMPORTS & DEPENDENCIES =====
        imports_raw = m.get("imports", [])
        if node_type == "import" and "parameters" in m:
            imports_raw = m["parameters"]

        if isinstance(imports_raw, str):
            imports = [imp.strip() for imp in imports_raw.split(",") if imp.strip()]
        else:
            imports = imports_raw if isinstance(imports_raw, list) else []

        if imports:
            file_node_id = f"file::{fpath}"
            # Collect new nodes instead of modifying created_nodes during iteration
            new_nodes = _add_import_edges_enhanced(graph, file_node_id, node_id, imports, lang)
            for new_node_id, new_node_data in new_nodes.items():
                if new_node_id not in created_nodes:
                    created_nodes[new_node_id] = new_node_data

        # ===== FUNCTION CALLS =====
        calls_raw = m.get("calls", [])
        if isinstance(calls_raw, str):
            calls = [c.strip() for c in calls_raw.split(",") if c.strip()]
        else:
            calls = calls_raw if isinstance(calls_raw, list) else []

        if calls:
            graph._add_call_edges(node_id, calls, fpath)

        # ===== CLASS/STRUCT INHERITANCE =====
        if node_type in ["class", "struct"]:
            extends = m.get("extends")
            if extends:
                if isinstance(extends, list):
                    extends_list = extends
                else:
                    extends_list = [extends]

                for parent_name in extends_list:
                    for class_name, class_info in class_registry.items():
                        if class_name == parent_name:
                            graph.add_edge(node_id, class_info["node_id"], "extends")

            implements = m.get("implements", [])
            if isinstance(implements, str):
                implements = [implements]

            for interface_name in implements:
                for iface_name, iface_info in interface_registry.items():
                    if iface_name == interface_name:
                        graph.add_edge(node_id, iface_info["node_id"], "implements")

                for trait_name, trait_info in trait_registry.items():
                    if trait_name == interface_name:
                        graph.add_edge(node_id, trait_info["node_id"], "implements")

        # ===== INTERFACE INHERITANCE =====
        if node_type == "interface":
            extends_list = m.get("extends", [])
            if isinstance(extends_list, str):
                extends_list = [extends_list]

            for parent_interface in extends_list:
                for iface_name, iface_info in interface_registry.items():
                    if iface_name == parent_interface:
                        graph.add_edge(node_id, iface_info["node_id"], "extends")

        # ===== RUST IMPL BLOCKS =====
        if node_type == "impl":
            traits = m.get("traits", [])
            if isinstance(traits, str):
                traits = [traits]

            for trait_name in traits:
                for t_name, t_info in trait_registry.items():
                    if t_name == trait_name:
                        graph.add_edge(node_id, t_info["node_id"], "implements")

                impl_name = m.get("name", "")
                for class_name, class_info in class_registry.items():
                    if class_name == impl_name or class_name in impl_name:
                        graph.add_edge(node_id, class_info["node_id"], "implements_for")

        # ===== GO METHOD RECEIVERS =====
        if lang == "go" and node_type in ["method", "function"]:
            receiver = m.get("receiver")
            receiver_base = m.get("receiver_base_type")

            if receiver_base:
                for struct_name, struct_info in class_registry.items():
                    if struct_name == receiver_base and struct_info["lang"] == "go":
                        graph.add_edge(struct_info["node_id"], node_id, "has_method")

        # ===== METHOD RETURN TYPES & PARAMETERS =====
        if node_type in ["function", "method", "async_function", "constructor"]:
            return_type = m.get("return_type")
            if return_type and return_type not in ['void', '()', 'Result', 'None', '', 'any', 'unknown']:
                clean_return = _clean_type_string(return_type)

                for type_name, type_info in {**class_registry, **interface_registry, **trait_registry}.items():
                    if type_name == clean_return:
                        graph.add_edge(node_id, type_info["node_id"], "returns")

            parameters = m.get("parameters", [])
            if isinstance(parameters, list):
                for param in parameters:
                    param_type = _extract_param_type(param)
                    if param_type:
                        for type_name, type_info in {**class_registry, **interface_registry, **trait_registry}.items():
                            if type_name == param_type:
                                graph.add_edge(node_id, type_info["node_id"], "uses_parameter")

        # ===== SVELTE COMPONENT RELATIONSHIPS =====
        if node_type == "component" or (node_type == "markup" and lang == "svelte"):
            used_components = m.get("used_components", [])
            for comp_name in used_components:
                if comp_name in component_registry:
                    comp_info = component_registry[comp_name]
                    graph.add_edge(node_id, comp_info["node_id"], "uses_component")

            props = m.get("component_props", [])
            for prop_name in props:
                for iface_name, iface_info in interface_registry.items():
                    if f"{m.get('component_name')}Props" == iface_name:
                        graph.add_edge(node_id, iface_info["node_id"], "uses_props_type")

        # ===== FIELD TYPE RELATIONSHIPS =====
        if node_type in ["struct", "class"]:
            fields = m.get("fields", [])
            for field in fields:
                field_type = _extract_field_type(field)
                if field_type:
                    for type_name, type_info in {**class_registry, **interface_registry, **trait_registry}.items():
                        if type_name == field_type:
                            graph.add_edge(node_id, type_info["node_id"], "has_field_type")

    logger.info("Phase 2 complete: Built intra-language relationships")

    # -------------------------------------------------
    # PHASE 3: Advanced relationship building
    # -------------------------------------------------
    logger.info("=== Phase 3: Advanced relationships ===")

    type_connections = build_type_relationships(graph, created_nodes)
    logger.info(f"Added {type_connections} type relationships")

    bevy_connections = build_bevy_relationships(graph, created_nodes)
    logger.info(f"Added {bevy_connections} Bevy-specific relationships")

    cross_lang_connections = build_cross_language_relationships(
        graph, created_nodes, component_registry, interface_registry, class_registry
    )
    logger.info(f"Added {cross_lang_connections} cross-language relationships")

    graph.save()
    graph.to_json()
    graph.to_toon()
    logger.info(f"✅ Graph built with {graph.graph.number_of_nodes()} nodes and {graph.graph.number_of_edges()} edges.")

    # -------------------------------------------------
    # PHASE 4: Embedding/index pipeline
    # -------------------------------------------------
    logger.info("=== Phase 4: Building embeddings ===")

    os.environ["OLLAMA_HOST"] = OLLAMA_BASE_URL
    embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL, base_url=OLLAMA_BASE_URL)

    logger.info(f"Generating embeddings for {len(texts)} documents in batches of {EMBEDDING_BATCH_SIZE}...")
    base_embeddings = batch_embed_documents(texts)

    logger.info("Applying contextual weights...")
    weighted_embeddings = apply_contextual_weights_to_embeddings(base_embeddings, metadatas)

    logger.info("Creating vector store (Chroma) with weighted embeddings...")

    class WeightedEmbeddingFunction:
        def __init__(self, texts, weighted_embs):
            self.lookup = defaultdict(deque)
            for t, emb in zip(texts, weighted_embs):
                self.lookup[t].append(emb)

        def embed_documents(self, docs):
            result = []
            for d in docs:
                if not self.lookup[d]:
                    raise ValueError(f"No embedding left for text {repr(d[:80])}")
                result.append(self.lookup[d].popleft())
            return result

        def embed_query(self, q):
            return embeddings.embed_query(q)

    weighted_emb_func = WeightedEmbeddingFunction(texts=texts, weighted_embs=weighted_embeddings)

    vectorstore = Chroma.from_texts(
        texts=texts,
        embedding=weighted_emb_func,
        persist_directory=str(VECTOR_DB_PATH),
        metadatas=metadatas
    )

    logger.info("Building BM25 index...")
    tokenized = [t.lower().split() for t in texts]
    bm25 = BM25Okapi(tokenized)
    bm25_corpus = texts
    chunks_metadata = metadatas

    BM25_INDEX_PATH.write_text(
        json.dumps({"corpus": texts, "metadata": metadatas}, indent=2),
        encoding="utf-8"
    )

    index_build_time = time.time()
    elapsed = index_build_time - start_time
    logger.info(f"✅ Index build complete. Total time: {elapsed:.1f}s")


# -------------------------------------------------
# Helper Functions
# -------------------------------------------------

def _add_import_edges_enhanced(graph, file_node_id, importing_node_id, imports, lang):
    """Enhanced import edge creation. Returns dict of new nodes to add."""
    new_nodes = {}

    for imp_path in imports:
        if not imp_path or not isinstance(imp_path, str):
            continue

        # Normalize path separators based on language
        if lang == "python":
            normalized_path = imp_path.replace('.', '::')
        elif lang in ["javascript", "typescript", "svelte"]:
            if imp_path.startswith('.') and len(imp_path) < 3:
                continue
            normalized_path = imp_path.replace('/', '::').replace('@', '')
        elif lang == "rust":
            normalized_path = imp_path
        elif lang == "go":
            normalized_path = imp_path.replace('/', '::')
        elif lang == "java":
            normalized_path = imp_path.replace('.', '::')
        else:
            normalized_path = imp_path.replace('.', '::').replace('/', '::')

        parts = [p.strip() for p in normalized_path.split('::') if p.strip()]
        if not parts:
            continue

        # Create package/crate/module hierarchy
        package_name = parts[0]
        package_node_id = f"package::{lang}::{package_name}"

        if package_node_id not in graph.graph.nodes:
            graph.add_node(
                package_node_id,
                type=package_name,
                name=package_name,
                file=imp_path,
                lang=lang,
                is_external=True
            )
            new_nodes[package_node_id] = {
                "type": package_name,
                "name": package_name,
                "file": imp_path,
                "lang": lang
            }

        if not graph.graph.has_edge(file_node_id, package_node_id):
            graph.add_edge(file_node_id, package_node_id, "imports")

        # Build hierarchical module structure
        current_path = package_name
        parent_node = package_node_id

        for i, part in enumerate(parts[1:], 1):
            current_path = f"{current_path}::{part}"

            if i == len(parts) - 1:
                node_type_name = part
            else:
                node_type_name = part

            module_node_id = f"symbol::{lang}::{current_path}"

            if module_node_id not in graph.graph.nodes:
                graph.add_node(
                    module_node_id,
                    type=node_type_name,
                    name=current_path,
                    file=imp_path,
                    lang=lang,
                    is_symbol=True
                )
                new_nodes[module_node_id] = {
                    "type": node_type_name,
                    "name": current_path,
                    "file": imp_path,
                    "lang": lang
                }

            if not graph.graph.has_edge(parent_node, module_node_id):
                graph.add_edge(parent_node, module_node_id, "contains")

            parent_node = module_node_id

        if parent_node != package_node_id:
            if not graph.graph.has_edge(importing_node_id, parent_node):
                graph.add_edge(importing_node_id, parent_node, "imports")

    return new_nodes

def _clean_type_string(type_str: str) -> str:
    """Clean up type strings for matching across all languages."""
    if not type_str:
        return ""

    type_str = type_str.strip()

    # Remove common Rust wrappers
    type_str = re.sub(r'Option<(.+?)>', r'\1', type_str)
    type_str = re.sub(r'Result<(.+?)(?:,|>).*', r'\1', type_str)
    type_str = re.sub(r'Vec<(.+?)>', r'\1', type_str)
    type_str = re.sub(r'Box<(.+?)>', r'\1', type_str)
    type_str = re.sub(r'Arc<(.+?)>', r'\1', type_str)
    type_str = re.sub(r'Rc<(.+?)>', r'\1', type_str)

    # Remove TypeScript/JavaScript wrappers
    type_str = re.sub(r'Promise<(.+?)>', r'\1', type_str)
    type_str = re.sub(r'Array<(.+?)>', r'\1', type_str)
    type_str = re.sub(r'Observable<(.+?)>', r'\1', type_str)

    # Remove Java wrappers
    type_str = re.sub(r'Optional<(.+?)>', r'\1', type_str)
    type_str = re.sub(r'List<(.+?)>', r'\1', type_str)
    type_str = re.sub(r'Set<(.+?)>', r'\1', type_str)
    type_str = re.sub(r'Map<.+?,\s*(.+?)>', r'\1', type_str)

    # Remove Go slices and pointers
    type_str = re.sub(r'\[\](.+)', r'\1', type_str)
    type_str = re.sub(r'\*(.+)', r'\1', type_str)

    # Remove Rust references and mutability
    type_str = type_str.replace('&', '').replace('mut ', '').strip()

    # Remove generic parameters for simpler matching
    type_str = re.sub(r'<.*?>', '', type_str).strip()

    # Remove array brackets
    type_str = re.sub(r'\[.*?\]', '', type_str).strip()

    return type_str

def _extract_param_type(param: str) -> str:
    """Extract type from parameter string for all languages."""
    if not param or not isinstance(param, str):
        return ""

    if ':' in param:
        parts = param.split(':')
        if len(parts) < 2:
            return ""
        param_type = parts[-1].strip()
    elif ' ' in param:
        parts = param.strip().split()
        if parts[0] and parts[0][0].isupper() or parts[0] in ['int', 'string', 'bool', 'float', 'double', 'long']:
            param_type = parts[0]
        else:
            param_type = parts[-1] if parts else ""
    else:
        return ""

    param_type = param_type.split('=')[0].strip()
    param_type = param_type.split('`')[0].strip()

    return _clean_type_string(param_type)


def _extract_field_type(field: str) -> str:
    """Extract type from field string for all languages."""
    if not field or not isinstance(field, str):
        return ""

    if field.startswith('embedded:'):
        return _clean_type_string(field[9:])

    if ':' in field:
        parts = field.split(':')
        if len(parts) < 2:
            return ""
        field_type = parts[-1].strip()
    elif ' ' in field:
        parts = field.strip().split()
        if len(parts) < 2:
            return ""
        if parts[0] and parts[0][0].isupper() or parts[0] in ['int', 'string', 'bool', 'float', 'double', 'long']:
            field_type = parts[0]
        else:
            field_type = parts[1] if len(parts) > 1 else parts[0]
    else:
        return ""

    field_type = re.sub(r'`.*?`', '', field_type).strip()
    field_type = re.sub(r'".*?"', '', field_type).strip()

    return _clean_type_string(field_type)


def build_cross_language_relationships(graph, created_nodes, component_registry,
                                       interface_registry, class_registry):
    """Build relationships across different programming languages."""
    connections = 0

    for ts_name, ts_info in interface_registry.items():
        if ts_info["lang"] in ["typescript", "javascript"]:
            for rust_name, rust_info in class_registry.items():
                if rust_info["lang"] == "rust" and ts_name == rust_name:
                    if not graph.graph.has_edge(ts_info["node_id"], rust_info["node_id"]):
                        graph.add_edge(ts_info["node_id"], rust_info["node_id"], "mirrors")
                        connections += 1

    for java_name, java_info in class_registry.items():
        if java_info["lang"] == "java":
            for go_name, go_info in class_registry.items():
                if go_info["lang"] == "go" and java_name == go_name:
                    if not graph.graph.has_edge(java_info["node_id"], go_info["node_id"]):
                        graph.add_edge(java_info["node_id"], go_info["node_id"], "mirrors")
                        connections += 1

    for comp_name, comp_info in component_registry.items():
        props_name = f"{comp_name}Props"
        for iface_name, iface_info in interface_registry.items():
            if iface_name == props_name or iface_name == comp_name:
                if not graph.graph.has_edge(comp_info["node_id"], iface_info["node_id"]):
                    graph.add_edge(comp_info["node_id"], iface_info["node_id"], "uses_type")
                    connections += 1

        for class_name, class_info in class_registry.items():
            if class_name == comp_name or class_name == props_name:
                if not graph.graph.has_edge(comp_info["node_id"], class_info["node_id"]):
                    graph.add_edge(comp_info["node_id"], class_info["node_id"], "uses_type")
                    connections += 1

    for py_name, py_info in class_registry.items():
        if py_info["lang"] == "python":
            for ts_name, ts_info in interface_registry.items():
                if ts_info["lang"] == "typescript" and py_name == ts_name:
                    if not graph.graph.has_edge(py_info["node_id"], ts_info["node_id"]):
                        graph.add_edge(py_info["node_id"], ts_info["node_id"], "api_type")
                        connections += 1

    return connections


def build_type_relationships(graph, created_nodes):
    """Build enhanced type relationships within language boundaries."""
    connections = 0

    type_map = {}
    for node_id, node_data in created_nodes.items():
        node_type = node_data.get("type")
        name = node_data.get("name")
        lang = node_data.get("lang")

        if node_type in ["class", "struct", "interface", "trait", "enum", "type"] and name:
            key = f"{lang}::{name}"
            if key not in type_map:
                type_map[key] = []
            type_map[key].append(node_id)

    for node_id, node_data in created_nodes.items():
        m = node_data.get("metadata", {})
        if not m:
            continue

        lang = node_data.get("lang")
        node_type = node_data.get("type")

        # Field types
        fields = m.get("fields", [])
        if isinstance(fields, list):
            for field in fields:
                field_type = _extract_field_type(field)
                if field_type:
                    type_key = f"{lang}::{field_type}"
                    if type_key in type_map:
                        for target_id in type_map[type_key]:
                            if not graph.graph.has_edge(node_id, target_id):
                                graph.add_edge(node_id, target_id, "has_field_of_type")
                                connections += 1

        # Return types
        return_type = m.get("return_type")
        if return_type:
            clean_return = _clean_type_string(return_type)
            if clean_return:
                type_key = f"{lang}::{clean_return}"
                if type_key in type_map:
                    for target_id in type_map[type_key]:
                        if not graph.graph.has_edge(node_id, target_id):
                            graph.add_edge(node_id, target_id, "returns_type")
                            connections += 1

        # Parameter types
        parameters = m.get("parameters", [])
        if isinstance(parameters, list):
            for param in parameters:
                param_type = _extract_param_type(param)
                if param_type:
                    type_key = f"{lang}::{param_type}"
                    if type_key in type_map:
                        for target_id in type_map[type_key]:
                            if not graph.graph.has_edge(node_id, target_id):
                                graph.add_edge(node_id, target_id, "accepts_type")
                                connections += 1

        # Generic type parameters
        generics = m.get("generics", [])
        if isinstance(generics, list):
            for generic in generics:
                generic_parts = generic.split(':')
                if len(generic_parts) > 1:
                    constraint = generic_parts[1].strip()
                    type_key = f"{lang}::{constraint}"
                    if type_key in type_map:
                        for target_id in type_map[type_key]:
                            if not graph.graph.has_edge(node_id, target_id):
                                graph.add_edge(node_id, target_id, "constrained_by")
                                connections += 1

    return connections


def build_bevy_relationships(graph, created_nodes):
    """Build Bevy-specific relationships (ECS patterns)."""
    connections = 0

    systems = []
    components = []
    resources = []
    bundles = []

    for node_id, node_data in created_nodes.items():
        m = node_data.get("metadata", {})
        if not m:
            continue

        lang = node_data.get("lang")
        if lang != "rust":
            continue

        node_type = node_data.get("type")
        name = node_data.get("name")

        if node_type in ["function", "method"]:
            params = m.get("parameters", [])
            has_query = any("Query" in str(p) for p in params)
            has_commands = any("Commands" in str(p) for p in params)
            has_res = any("Res" in str(p) or "ResMut" in str(p) for p in params)

            if has_query or has_commands or has_res:
                systems.append(node_id)

        if node_type == "struct":
            derives = m.get("derives", [])
            if "Component" in str(derives):
                components.append(node_id)

            if "Resource" in str(derives):
                resources.append(node_id)

            if "Bundle" in str(derives):
                bundles.append(node_id)

    for system_id in systems:
        system_data = created_nodes.get(system_id, {})
        m = system_data.get("metadata", {})
        params = m.get("parameters", [])

        for param in params:
            param_str = str(param)

            query_match = re.findall(r'Query<[^>]*?([A-Z]\w+)', param_str)
            for comp_name in query_match:
                for comp_id in components:
                    comp_data = created_nodes.get(comp_id, {})
                    if comp_data.get("name") == comp_name:
                        if not graph.graph.has_edge(system_id, comp_id):
                            graph.add_edge(system_id, comp_id, "queries_component")
                            connections += 1

            res_match = re.findall(r'Res(?:Mut)?<([A-Z]\w+)', param_str)
            for res_name in res_match:
                for res_id in resources:
                    res_data = created_nodes.get(res_id, {})
                    if res_data.get("name") == res_name:
                        if not graph.graph.has_edge(system_id, res_id):
                            graph.add_edge(system_id, res_id, "uses_resource")
                            connections += 1

    for bundle_id in bundles:
        bundle_data = created_nodes.get(bundle_id, {})
        m = bundle_data.get("metadata", {})
        fields = m.get("fields", [])

        for field in fields:
            field_type = _extract_field_type(field)
            for comp_id in components:
                comp_data = created_nodes.get(comp_id, {})
                if comp_data.get("name") == field_type:
                    if not graph.graph.has_edge(bundle_id, comp_id):
                        graph.add_edge(bundle_id, comp_id, "bundles_component")
                        connections += 1

    return connections

def build_bevy_relationships(graph, created_nodes):
    """Build Bevy-specific relationships (ECS patterns)."""
    connections = 0

    # Track Bevy systems, components, and resources
    systems = []
    components = []
    resources = []
    bundles = []

    for node_id, node_data in created_nodes.items():
        m = node_data.get("metadata", {})
        if not m:
            continue

        lang = node_data.get("lang")
        if lang != "rust":
            continue

        node_type = node_data.get("type")
        name = node_data.get("name")

        # Detect Bevy systems (functions with Query parameters)
        if node_type in ["function", "method"]:
            params = m.get("parameters", [])
            has_query = any("Query" in str(p) for p in params)
            has_commands = any("Commands" in str(p) for p in params)
            has_res = any("Res" in str(p) or "ResMut" in str(p) for p in params)

            if has_query or has_commands or has_res:
                systems.append(node_id)

        # Detect Bevy components (structs with Component derive)
        if node_type == "struct":
            derives = m.get("derives", [])
            if "Component" in str(derives):
                components.append(node_id)

        # Detect Bevy resources (structs with Resource derive)
        if node_type == "struct":
            derives = m.get("derives", [])
            if "Resource" in str(derives):
                resources.append(node_id)

        # Detect Bevy bundles (structs with Bundle derive)
        if node_type == "struct":
            derives = m.get("derives", [])
            if "Bundle" in str(derives):
                bundles.append(node_id)

    # Connect systems to the components/resources they use
    for system_id in systems:
        system_data = created_nodes.get(system_id, {})
        m = system_data.get("metadata", {})
        params = m.get("parameters", [])

        for param in params:
            param_str = str(param)

            # Extract component types from Query<Component>
            query_match = re.findall(r'Query<[^>]*?([A-Z]\w+)', param_str)
            for comp_name in query_match:
                for comp_id in components:
                    comp_data = created_nodes.get(comp_id, {})
                    if comp_data.get("name") == comp_name:
                        if not graph.graph.has_edge(system_id, comp_id):
                            graph.add_edge(system_id, comp_id, "queries_component")
                            connections += 1

            # Extract resource types from Res<Resource>
            res_match = re.findall(r'Res(?:Mut)?<([A-Z]\w+)', param_str)
            for res_name in res_match:
                for res_id in resources:
                    res_data = created_nodes.get(res_id, {})
                    if res_data.get("name") == res_name:
                        if not graph.graph.has_edge(system_id, res_id):
                            graph.add_edge(system_id, res_id, "uses_resource")
                            connections += 1

    # Connect bundles to their component fields
    for bundle_id in bundles:
        bundle_data = created_nodes.get(bundle_id, {})
        m = bundle_data.get("metadata", {})
        fields = m.get("fields", [])

        for field in fields:
            field_type = _extract_field_type(field)
            for comp_id in components:
                comp_data = created_nodes.get(comp_id, {})
                if comp_data.get("name") == field_type:
                    if not graph.graph.has_edge(bundle_id, comp_id):
                        graph.add_edge(bundle_id, comp_id, "bundles_component")
                        connections += 1

    return connections

def _create_import_symbolic_nodes(graph, import_node_id, file_node_id, import_path, lang):
    """Create rich symbolic nodes for imports. Returns dict of new nodes to add."""
    new_nodes = {}

    if not import_path:
        return new_nodes

    # Normalize path separators based on language
    if lang == "python":
        parts = import_path.replace('.', '::').split('::')
    elif lang in ["javascript", "typescript"]:
        if import_path.startswith('.'):
            return new_nodes
        parts = import_path.replace('/', '::').replace('@', '').split('::')
    elif lang == "rust":
        parts = import_path.split('::')
    elif lang == "go":
        parts = import_path.split('/')
    elif lang == "java":
        parts = import_path.split('.')
    else:
        parts = import_path.replace('.', '::').replace('/', '::').split('::')

    parts = [p.strip() for p in parts if p.strip()]
    if not parts:
        return new_nodes

    # Create package/crate node
    package_name = parts[0]
    package_node_id = f"package::{lang}::{package_name}"

    if package_node_id not in graph.graph.nodes:
        graph.add_node(
            package_node_id,
            type=package_name,
            name=package_name,
            file=import_path,
            lang=lang,
            is_external=True
        )
        new_nodes[package_node_id] = {
            "type": package_name,
            "name": package_name,
            "file": import_path,
            "lang": lang
        }

    if not graph.graph.has_edge(file_node_id, package_node_id):
        graph.add_edge(file_node_id, package_node_id, "imports")

    # Create intermediate module/namespace nodes
    current_path = package_name
    parent_node = package_node_id

    for part in parts[1:]:
        current_path = f"{current_path}::{part}"
        module_node_id = f"symbol::{lang}::{current_path}"

        if module_node_id not in graph.graph.nodes:
            graph.add_node(
                module_node_id,
                type=part,
                name=current_path,
                file=import_path,
                lang=lang,
                is_symbol=True
            )
            new_nodes[module_node_id] = {
                "type": part,
                "name": current_path,
                "file": import_path,
                "lang": lang
            }

        if not graph.graph.has_edge(parent_node, module_node_id):
            graph.add_edge(parent_node, module_node_id, "contains")

        parent_node = module_node_id

    if parent_node != package_node_id and not graph.graph.has_edge(import_node_id, parent_node):
        graph.add_edge(import_node_id, parent_node, "imports")
    return new_nodes

def _extract_parameter_type(param):
    """Extract type from parameter string, handling multiple formats"""
    if not param:
        return None

    # Handle "type: name" format (TypeScript/Python)
    if ':' in param:
        parts = param.split(':')
        if len(parts) >= 2:
            return parts[1].strip()

    # Handle "Type name" format (Rust/Java/C++)
    if ' ' in param:
        parts = param.split()
        if len(parts) >= 2:
            return parts[0].strip()

    # Handle complex types with generics
    if '<' in param and '>' in param:
        # Extract the base type before generics
        base_type = param.split('<')[0].strip()
        return base_type

    return param.strip()

def _clean_type_name(type_name):
    """Clean type names by removing common modifiers"""
    if not type_name:
        return ""

    # Remove common modifiers
    modifiers = ['&', 'mut ', 'const ', 'static ', 'pub ', 'private ', 'protected ']
    cleaned = type_name
    for mod in modifiers:
        cleaned = cleaned.replace(mod, '')

    # Remove trailing references/pointers
    cleaned = cleaned.rstrip('*&')

    return cleaned.strip()

def _is_meaningful_type(type_name):
    """Check if a type name is meaningful (not primitive, void, etc.)"""
    if not type_name:
        return False

    meaningless_types = [
        'void', 'int', 'i32', 'i64', 'f32', 'f64', 'bool', 'str', 'String',
        'char', 'u8', 'u16', 'u32', 'u64', 'usize', 'isize',
        'self', '&self', '&mut self', 'mut self', 'Self',
        'Result', 'Option', 'Vec', 'String'
    ]

    cleaned = _clean_type_name(type_name)
    return (cleaned and
            cleaned not in meaningless_types and
            not cleaned.startswith('&') and
            len(cleaned) > 1)

# -----------------------------
# Build / load indexes with contextual weighting and batch embedding
# -----------------------------
def batch_embed_documents(texts: List[str]) -> List[List[float]]:
    """Embed documents in batches for better performance."""
    all_embeddings = []
    total_batches = (len(texts) + EMBEDDING_BATCH_SIZE - 1) // EMBEDDING_BATCH_SIZE

    for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
        batch = texts[i:i + EMBEDDING_BATCH_SIZE]
        batch_num = i // EMBEDDING_BATCH_SIZE + 1
        logger.info(f"Embedding batch {batch_num}/{total_batches} ({len(batch)} docs)...")
        batch_embeddings = embeddings.embed_documents(batch)
        all_embeddings.extend(batch_embeddings)

    return all_embeddings

def load_indexes():
    global vectorstore, bm25, bm25_corpus, chunks_metadata, embeddings, index_build_time
    logger.info("Loading indexes from disk...")
    embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL, base_url=OLLAMA_BASE_URL)
    vectorstore = Chroma(persist_directory=str(VECTOR_DB_PATH), embedding_function=embeddings)
    data = json.loads(BM25_INDEX_PATH.read_text(encoding="utf-8"))
    bm25_corpus = data["corpus"]
    chunks_metadata = data["metadata"]
    bm25 = BM25Okapi([t.lower().split() for t in bm25_corpus])
    index_build_time = BM25_INDEX_PATH.stat().st_mtime
    logger.info("Indexes loaded.")

# -----------------------------
# Analytics Helper Functions
# -----------------------------
def count_by_language(metadata_list: List[Dict[str, Any]]) -> Dict[str, int]:
    """Count chunks by programming language."""
    counts = {}
    for meta in metadata_list:
        lang = meta.get("language", "unknown")
        counts[lang] = counts.get(lang, 0) + 1
    return counts

def count_by_type(metadata_list: List[Dict[str, Any]]) -> Dict[str, int]:
    """Count chunks by type (function, class, etc.)."""
    counts = {}
    for meta in metadata_list:
        typ = meta.get("type", "unknown")
        counts[typ] = counts.get(typ, 0) + 1
    return counts

def calculate_avg_chunk_size(corpus: List[str]) -> float:
    """Calculate average chunk size in characters."""
    if not corpus:
        return 0.0
    return sum(len(text) for text in corpus) / len(corpus)

# -----------------------------
# Hybrid search 80/20 + normalization
# -----------------------------
def hybrid_search(query: str, k: int = 30) -> List[Tuple[str, Dict[str, Any], float]]:
    with safe_search():
        # Vector results (Chroma returns (doc, score) where lower score might be better depending on backend)
        vector_results = vectorstore.similarity_search_with_score(query, k=50)
        # Normalize vector scores so that higher is better in [0,1]
        vec_values = [s for _, s in vector_results] if vector_results else [1.0]
        vmin, vmax = (min(vec_values), max(vec_values)) if vector_results else (0.0, 1.0)
        vector_scores = {}
        for doc, s in vector_results:
            # If Chroma returns distance-like (smaller better), invert; otherwise adapt
            # We'll normalize via (vmax - s) / (vmax - vmin + eps)
            eps = 1e-12
            denom = (vmax - vmin) + eps
            norm = (vmax - s) / denom
            vector_scores[doc.page_content] = norm

        # BM25
        tokenized_query = query.lower().split()
        bm25_scores = bm25.get_scores(tokenized_query)
        top_idxs = sorted(range(len(bm25_scores)), key=lambda i: bm25_scores[i], reverse=True)[:50]
        max_b = max((bm25_scores[i] for i in top_idxs), default=1.0)
        bm25_doc_scores = {bm25_corpus[i]: (bm25_scores[i] / max_b if max_b > 0 else 0.0) for i in top_idxs}

        # Fusion 80/20
        all_docs = set(vector_scores.keys()) | set(bm25_doc_scores.keys())
        fused = {}
        for d in all_docs:
            v = vector_scores.get(d, 0.0)
            b = bm25_doc_scores.get(d, 0.0)
            fused[d] = VECTOR_WEIGHT * v + BM25_WEIGHT * b

        # Sort by fused score descending
        sorted_items = sorted(fused.items(), key=lambda x: x[1], reverse=True)

	    # Build accurate mapping from document text to its true metadata
        text_to_meta = {text: meta for text, meta in zip(bm25_corpus, chunks_metadata)}

        results = []
        for doc_text, score in sorted_items[:k]:
            meta = text_to_meta.get(doc_text, {})
            results.append((doc_text, meta, score))
        return results

# -----------------------------
# Enhanced Reranker with Batch Processing
# -----------------------------
def rerank_with_ollama_enhanced(query: str, candidates: List[str], top_k: int = 5) -> List[Tuple[str, float]]:
    """Enhanced Qwen3 Reranker with batch processing and better scoring."""
    if not candidates:
        return []

    # Use domain-specific instruction for code search
    instruction = "Given a technical code search query, retrieve relevant code implementations, function definitions, or examples that directly address the query requirements"

    batch_scores = []

    # Process in smaller batches to avoid timeouts
    batch_size = 5
    for i in range(0, len(candidates), batch_size):
        batch = candidates[i:i + batch_size]
        batch_prompts = []

        for chunk in batch:
            system_prompt = '<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n'

            user_prompt = (
                f'<|im_start|>user\n'
                f'<Instruct>: {instruction}\n'
                f'<Query>: {query}\n'
                f'<Document>: {chunk[:8000]}\n'  # Limit document length
                f'<|im_end|>\n'
                f'<|im_start|>assistant\n<think>\n\n</think>\n\n'
            )

            batch_prompts.append(system_prompt + user_prompt)

        # Score batch
        for prompt, chunk in zip(batch_prompts, batch):
            try:
                resp = requests.post(
                    f"{OLLAMA_BASE_URL}/api/generate",
                    json={
                        "model": RERANKER_MODEL_OLLAMA,
                        "prompt": prompt,
                        "max_tokens": 3,
                        "temperature": 0.0,
                        "stream": False
                    },
                    timeout=20
                )

                j = resp.json()
                response_text = j.get("response", "").strip().lower()

                # Enhanced scoring with confidence levels
                score = parse_reranker_response(response_text)
                batch_scores.append((chunk, score))

            except requests.exceptions.Timeout:
                logger.warning("Reranker request timeout, assigning default score")
                batch_scores.append((chunk, 0.0))
            except Exception as e:
                logger.warning(f"Rerank call failed: {e}")
                batch_scores.append((chunk, 0.0))

    # Sort by score descending and return top_k
    batch_scores.sort(key=lambda x: x[1], reverse=True)
    return batch_scores[:top_k]


def parse_reranker_response(response_text: str) -> float:
    """Parse Qwen3 Reranker response and convert to confidence score."""
    response_lower = response_text.strip().lower()

    # Exact matches from the model
    if response_lower == "yes":
        return 1.0
    elif response_lower == "no":
        return 0.0

    # Handle variations and partial matches
    yes_indicators = ["yes", "relevant", "correct", "matches", "appropriate", "suitable"]
    no_indicators = ["no", "irrelevant", "incorrect", "unrelated", "inappropriate"]

    yes_count = sum(1 for indicator in yes_indicators if indicator in response_lower)
    no_count = sum(1 for indicator in no_indicators if indicator in response_lower)

    if yes_count > no_count:
        return 0.8  # Likely relevant but not confident
    elif no_count > yes_count:
        return 0.2  # Likely irrelevant but not confident
    else:
        # Ambiguous response
        return 0.5

# -----------------------------
# Comment-aware reference search (ripgrep)
# -----------------------------
def get_context_with_comments(file_path: Path, match_line: int, max_context_lines: int = 5) -> str:
    try:
        lines = file_path.read_text(encoding="utf-8").splitlines()
    except Exception:
        return ""
    idx = match_line - 1
    context_lines = []
    look = idx - 1
    collected = 0
    while look >= 0 and collected < max_context_lines:
        line = lines[look]
        stripped = line.strip()
        if stripped.startswith(("#", "//", "--")) or stripped == "":
            context_lines.insert(0, line)
            look -= 1
            collected += 1
        else:
            break
    context_lines.append(lines[idx] if 0 <= idx < len(lines) else "")
    return "\n".join(context_lines)

def search_symbol(symbol: str, root_dir: Path, top_k: int = 20) -> List[Dict[str, Any]]:
    matches = []
    rg_cmd = [
        "rg",
        "-n",
        "-i",
        "-C", "3",
        "--no-ignore",
        "--no-heading",
        "--color", "never",
        "--glob", "!node_modules",
        "--glob", "!__pycache__",
        "--glob", "!.git",
    ]
    for ext in SUPPORTED_EXTENSIONS:
        rg_cmd.extend(["--glob", f"*{ext}"])
    rg_cmd.extend(["--", symbol, "."])

    try:
        proc = subprocess.run(
            rg_cmd,
            cwd=root_dir,
            capture_output=True,
            text=True,
            timeout=30
        )
        if proc.returncode not in (0, 1):
            logger.warning(f"ripgrep returned code {proc.returncode}: {proc.stderr}")
            return []
        for line in proc.stdout.splitlines():
            if not line.strip() or line.startswith("--"):
                continue
            parts = line.split(":", 2)
            if len(parts) < 3:
                continue
            file_rel, lineno_str, snippet = parts
            try:
                lineno = int(lineno_str)
            except ValueError:
                continue
            matches.append({
                "file": file_rel,
                "line": lineno,
                "context": snippet.strip()
            })
    except subprocess.TimeoutExpired:
        logger.warning("ripgrep timed out")
    except Exception as e:
        logger.warning(f"ripgrep failed: {e}")
    return matches


def _extract_imports(source: str) -> list[str]:
    """Return a list of fully‑qualified imports found in `source`."""
    return [m.group(1).strip() for m in re.finditer(r'^\s*use\s+([^;]+);', source, re.MULTILINE)]

def _extract_calls(source: str) -> list[str]:
    """Return a list of called function names (bare name only)."""
    return [m.group(1).split("::")[-1] for m in re.finditer(r'([A-Za-z_][A-Za-z0-9_:]*)\s*\(', source, re.MULTILINE)]

# -----------------------------
# Additional tool: read file lines with context
# -----------------------------
def read_file_lines(path: str, start: int = 1, end: Optional[int] = None) -> str:
    """
    Read file lines with intelligent context inclusion.
    - Includes requested lines
    - Expands to include full function/class definition if within one
    - Includes leading comments (docstrings, inline comments)
    - Shows function signature, return type, decorators
    """
    # Handle both absolute and relative paths
    if path.startswith('/'):
        p = Path(path)
    else:
        p = CODEBASE_PATH / Path(path)

    # Security: Validate path is within codebase
    try:
        resolved = p.resolve()
        codebase_resolved = CODEBASE_PATH.resolve()
        if not resolved.is_relative_to(codebase_resolved):
            return f"Error: Path '{path}' is outside the codebase directory"
    except (ValueError, OSError) as e:
        return f"Error: Invalid path '{path}': {e}"

    try:
        content = p.read_text(encoding="utf-8")
        lines = content.splitlines()

        if end is None:
            end = start + 50  # reasonable default

        # Clamp to valid range
        start_i = max(1, start)
        end_i = min(len(lines), end)

        # Detect language
        language = detect_language(p)

        # Build context-aware line range
        context_start, context_end, context_info = find_context_boundaries(
            lines, start_i, end_i, language
        )

        # Extract the context lines
        context_lines = lines[context_start - 1:context_end]

        # Build output with metadata
        output = []
        output.append(f"File: {path}")
        output.append(f"Language: {language}")
        output.append(f"Requested: lines {start_i}-{end_i}")
        output.append(f"Showing: lines {context_start}-{context_end} (with context)")

        if context_info:
            output.append(f"Context: {context_info}")

        output.append("\n" + "="*60)

        # Add line numbers
        for i, line in enumerate(context_lines, start=context_start):
            # Highlight the originally requested range
            marker = ">>> " if start_i <= i <= end_i else "    "
            output.append(f"{marker}{i:4d} | {line}")

        output.append("="*60)

        return "\n".join(output)

    except Exception as e:
        return f"Failed to read {path}: {e}"


def find_context_boundaries(lines: List[str], start: int, end: int, language: str) -> Tuple[int, int, str]:
    """
    Find intelligent context boundaries around the requested line range.
    Returns: (context_start_line, context_end_line, description)
    """
    # Convert to 0-based indexing for processing
    start_idx = start - 1
    end_idx = end - 1

    context_info = []

    # Expand upward to include leading comments and function/class headers
    context_start = start

    # Step 1: Include leading comments above start
    comment_start = find_leading_comments_start(lines, start_idx, language)
    if comment_start < start_idx:
        context_start = comment_start + 1
        context_info.append("leading comments")

    # Step 2: Check if we're inside a function/class and include its signature
    func_start, func_end, func_name, func_type = find_enclosing_function_or_class(
        lines, start_idx, end_idx, language
    )

    if func_start is not None:
        context_start = min(context_start, func_start + 1)
        if func_end is not None and func_end > end_idx:
            context_end = func_end + 1
        else:
            context_end = end

        if func_name:
            context_info.append(f"{func_type} '{func_name}'")
    else:
        context_end = end

    # Step 3: Include decorators (Python) or annotations (Java)
    if language == "python":
        decorator_start = find_decorators_start(lines, context_start - 1)
        if decorator_start < context_start - 1:
            context_start = decorator_start + 1
            context_info.append("decorators")

    info_str = " + ".join(context_info) if context_info else "no additional context"

    return context_start, context_end, info_str


def find_leading_comments_start(lines: List[str], line_idx: int, language: str) -> int:
    """Find the start of leading comments above the given line."""
    comment_patterns = {
        "python": ("#", '"""', "'''"),
        "javascript": ("//", "/*"),
        "typescript": ("//", "/*"),
        "go": ("//", "/*"),
        "java": ("//", "/*"),
        "sql": ("--", "/*"),
        "rust": ("//", "/*"),
    }

    patterns = comment_patterns.get(language, ("#", "//", "/*", "--"))

    start_idx = line_idx
    i = line_idx - 1

    # Track if we're in a block comment
    in_block = False

    while i >= 0:
        line = lines[i].strip()

        # Empty lines are OK if we already have comments
        if not line:
            if i < line_idx - 1:  # Allow one blank line between comments
                i -= 1
                continue
            else:
                break

        # Check for block comment end (we're going backwards)
        if "*/" in line:
            in_block = True
            start_idx = i
            i -= 1
            continue

        # Check for block comment start
        if in_block and ("/*" in line or "/**" in line):
            start_idx = i
            in_block = False
            i -= 1
            continue

        # Check for single-line comments
        is_comment = any(line.startswith(p) for p in patterns)

        if is_comment or in_block:
            start_idx = i
            i -= 1
        else:
            break

    return start_idx


def find_decorators_start(lines: List[str], line_idx: int) -> int:
    """Find Python decorators above the given line."""
    i = line_idx - 1
    while i >= 0:
        line = lines[i].strip()
        if line.startswith("@"):
            i -= 1
        elif not line:  # Allow blank lines
            i -= 1
        else:
            return i + 1
    return 0


def find_enclosing_function_or_class(lines: List[str], start_idx: int, end_idx: int, language: str) -> Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]:
    """
    Find the function or class that encloses the given line range.
    Returns: (start_line_idx, end_line_idx, name, type) or (None, None, None, None)
    """
    if language == "python":
        return find_python_function_or_class(lines, start_idx, end_idx)
    elif language in ("javascript", "typescript"):
        return find_js_function(lines, start_idx, end_idx)
    elif language == "go":
        return find_go_function(lines, start_idx, end_idx)
    elif language == "java":
        return find_java_method_or_class(lines, start_idx, end_idx)
    elif language == "sql":
        return find_sql_function_or_block(lines, start_idx, end_idx)

    return None, None, None, None


def find_python_function_or_class(lines: List[str], start_idx: int, end_idx: int) -> Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]:
    """Find enclosing Python function or class."""
    # Look backwards for def/class with proper indentation
    target_indent = None

    for i in range(start_idx, -1, -1):
        line = lines[i]
        stripped = line.lstrip()

        if stripped.startswith(("def ", "class ", "async def ")):
            indent = len(line) - len(stripped)

            # Check if this could be our enclosing scope
            if target_indent is None or indent < target_indent:
                # Extract name
                match = re.match(r'(?:async\s+)?(?:def|class)\s+(\w+)', stripped)
                if match:
                    name = match.group(1)
                    func_type = "class" if stripped.startswith("class") else "function"

                    # Find the end by looking for next same-or-lower indent non-empty line
                    end_i = find_python_block_end(lines, i, indent)

                    return i, end_i, name, func_type

    return None, None, None, None


def find_python_block_end(lines: List[str], start_idx: int, base_indent: int) -> Optional[int]:
    """Find the end of a Python block based on indentation."""
    for i in range(start_idx + 1, len(lines)):
        line = lines[i]
        stripped = line.lstrip()

        if not stripped:  # Skip empty lines
            continue

        indent = len(line) - len(stripped)

        if indent <= base_indent and stripped and not stripped.startswith("#"):
            return i - 1

    return len(lines) - 1


def find_js_function(lines: List[str], start_idx: int, end_idx: int) -> Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]:
    """Find enclosing JavaScript/TypeScript function."""
    for i in range(start_idx, -1, -1):
        line = lines[i].strip()

        # Match various JS function patterns
        patterns = [
            r'function\s+(\w+)',
            r'(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s+)?\(',
            r'(\w+)\s*\([^)]*\)\s*{',  # Arrow functions
            r'async\s+function\s+(\w+)',
        ]

        for pattern in patterns:
            match = re.search(pattern, line)
            if match:
                name = match.group(1)
                end_i = find_brace_block_end(lines, i)
                return i, end_i, name, "function"

    return None, None, None, None


def find_go_function(lines: List[str], start_idx: int, end_idx: int) -> Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]:
    """Find enclosing Go function."""
    for i in range(start_idx, -1, -1):
        line = lines[i].strip()

        # Match Go function: func (receiver) name(params) returnType {
        match = re.match(r'func\s+(?:\([^)]+\)\s+)?(\w+)', line)
        if match:
            name = match.group(1)
            end_i = find_brace_block_end(lines, i)
            return i, end_i, name, "function"

    return None, None, None, None


def find_java_method_or_class(lines: List[str], start_idx: int, end_idx: int) -> Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]:
    """Find enclosing Java method or class."""
    for i in range(start_idx, -1, -1):
        line = lines[i].strip()

        # Match class
        class_match = re.match(r'(?:public|private|protected)?\s*(?:static)?\s*class\s+(\w+)', line)
        if class_match:
            name = class_match.group(1)
            end_i = find_brace_block_end(lines, i)
            return i, end_i, name, "class"

        # Match method
        method_match = re.match(r'(?:public|private|protected)?\s*(?:static)?\s*(?:\w+(?:<[^>]+>)?)\s+(\w+)\s*\(', line)
        if method_match:
            name = method_match.group(1)
            end_i = find_brace_block_end(lines, i)
            return i, end_i, name, "method"

    return None, None, None, None


def find_sql_function_or_block(lines: List[str], start_idx: int, end_idx: int) -> Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]:
    """Find enclosing SQL function or procedure."""
    for i in range(start_idx, -1, -1):
        line = lines[i].strip().upper()

        if line.startswith(("CREATE FUNCTION", "CREATE OR REPLACE FUNCTION", "CREATE PROCEDURE")):
            # Extract name
            match = re.search(r'(?:FUNCTION|PROCEDURE)\s+(\w+)', line, re.IGNORECASE)
            if match:
                name = match.group(1)
                # SQL functions typically end with $$ or END
                end_i = find_sql_function_end(lines, i)
                func_type = "procedure" if "PROCEDURE" in line else "function"
                return i, end_i, name, func_type

    return None, None, None, None


def find_sql_function_end(lines: List[str], start_idx: int) -> Optional[int]:
    """Find end of SQL function (looks for $$ or END;)."""
    for i in range(start_idx + 1, len(lines)):
        line = lines[i].strip().upper()
        if "$$" in line or line.startswith("END;") or line == "END":
            return i
    return len(lines) - 1


def find_brace_block_end(lines: List[str], start_idx: int) -> Optional[int]:
    """Find the end of a brace-delimited block {}."""
    brace_count = 0
    started = False

    for i in range(start_idx, len(lines)):
        line = lines[i]

        for char in line:
            if char == '{':
                brace_count += 1
                started = True
            elif char == '}':
                brace_count -= 1

                if started and brace_count == 0:
                    return i

    return len(lines) - 1

# -----------------------------
# MCP Tools
# -----------------------------
@mcp.tool()
def health_check() -> str:
    """What it does - Quick JSON status of the RAG server (ready, indexed path, chunk count, age, graph stats).
When to use - Before any other query, to confirm the index is current and see codebase structure.
When not to use - After you already know the server is healthy; it adds no value.
Example - health_check() → { "status":"ready","total_chunks":11234,"graph_nodes":4567,"graph_edges":12345,"index_age_hours":4.7 }"""
    try:
        with _startup_lock:
            # Check if we have a git repo in working_repo
            repo_info = "No repository loaded"
            if (CODEBASE_PATH / ".git").exists():
                try:
                    branch = subprocess.check_output(
                        ["git", "-C", str(CODEBASE_PATH), "rev-parse", "--abbrev-ref", "HEAD"],
                        text=True
                    ).strip()
                    repo_info = f"Loaded: {branch}"
                except:
                    repo_info = "Git repository (no branch info)"

            # Load graph to check if it exists - FIX: Convert Path to string
            graph = LocalGraph(root_path=str(CODEBASE_PATH))  # Convert Path to str
            graph_loaded = graph.load()

            status = {
                "status": "ready" if vectorstore is not None else "no_index",
                "working_repository": str(CODEBASE_PATH),
                "repo_status": repo_info,
                "ollama_url": OLLAMA_BASE_URL,
                "config": {
                    "embedding_model": EMBEDDING_MODEL,
                    "vector_weight": VECTOR_WEIGHT,
                    "bm25_weight": BM25_WEIGHT,
                    "rerank_enabled": ENABLE_RERANK,
                    "batch_size": EMBEDDING_BATCH_SIZE,
                    "graph_enabled": True
                },
                "Tools": {
                    "NOTE": "THESE TOOLS ARE RESTRICTED BY .gitignore AS WELL AS .ragignore",
                    "semantic_rag_search": "Hybrid RAG search with graph context (cross-file deps, inheritance, calls)",
                    "find_code_references": "Find exact symbol references across codebase",
                    "read_file_lines_tool": "Read specific file lines with syntax highlighting",
                    "find_path": "Find files by glob patterns",
                    "grep": "Search file contents with regex",
                    "list_directory": "List directory contents"
                }
            }

            if bm25_corpus:
                status["statistics"] = {
                    "total_chunks": len(bm25_corpus),
                    "by_language": count_by_language(chunks_metadata),
                    "by_type": count_by_type(chunks_metadata),
                    "avg_chunk_size_chars": round(calculate_avg_chunk_size(bm25_corpus), 1),
                    "index_age_hours": round((time.time() - index_build_time) / 3600, 1) if index_build_time > 0 else None
                }

            # Add graph statistics - FIX: Use graph.graph.nodes instead of graph.nodes
            if graph_loaded:
                graph_stats = _calculate_graph_stats(graph)
                status["graph"] = graph_stats
            else:
                status["graph"] = {
                    "status": "not_loaded",
                    "nodes": 0,
                    "edges": 0,
                    "message": "Graph will be built on next index rebuild"
                }

            return json.dumps(status, indent=2)
    except Exception as e:
        return json.dumps({"status": "error", "message": str(e)}, indent=2)

def _calculate_graph_stats(graph: LocalGraph) -> dict:
    """Calculate comprehensive graph statistics."""
    # FIX: Use graph.graph.nodes and graph.graph.edges instead of graph.nodes/graph.edges
    if graph.graph.number_of_nodes() == 0:
        return {"status": "empty", "nodes": 0, "edges": 0}

    # Basic counts
    nodes = graph.graph.number_of_nodes()
    edges = graph.graph.number_of_edges()

    # Count by node type
    node_types = {}
    for _, data in graph.graph.nodes(data=True):  # FIX: graph.graph.nodes
        node_type = data.get('type', 'unknown')
        node_types[node_type] = node_types.get(node_type, 0) + 1

    # Count by edge type
    edge_types = {}
    for _, _, data in graph.graph.edges(data=True):  # FIX: graph.graph.edges
        edge_type = data.get('type', 'unknown')
        edge_types[edge_type] = edge_types.get(edge_type, 0) + 1

    # Language distribution
    languages = {}
    for _, data in graph.graph.nodes(data=True):  # FIX: graph.graph.nodes
        lang = data.get('lang', 'unknown')  # Changed from 'language' to 'lang' to match our new attribute
        languages[lang] = languages.get(lang, 0) + 1

    # File statistics
    files = [n for n, d in graph.graph.nodes(data=True) if d.get('type') == 'File']  # FIX: graph.graph.nodes

    # Relationship density
    avg_edges_per_node = edges / nodes if nodes > 0 else 0

    # Most connected nodes (hubs)
    degree_centrality = dict(graph.graph.degree())  # FIX: graph.graph.degree()
    top_hubs = sorted(degree_centrality.items(), key=lambda x: x[1], reverse=True)[:5]
    top_hubs_info = []
    for node_id, degree in top_hubs:
        node_data = graph.graph.nodes[node_id]  # FIX: graph.graph.nodes
        top_hubs_info.append({
            "name": node_data.get('name', node_id),
            "type": node_data.get('type', 'unknown'),
            "file": node_data.get('file', ''),
            "connections": degree
        })

    return {
        "status": "loaded",
        "nodes": nodes,
        "edges": edges,
        "node_types": node_types,
        "edge_types": edge_types,
        "languages": languages,
        "files": len(files),
        "avg_edges_per_node": round(avg_edges_per_node, 2),
        "top_connected_nodes": top_hubs_info,
        "relationship_density": "high" if avg_edges_per_node > 2.0 else "medium" if avg_edges_per_node > 1.0 else "low"
    }

@mcp.tool()
def semantic_rag_search(query: str, top_k: int = 5, rerank: bool = True) -> str:
    """Hybrid RAG search: returns ranked code snippets + file/line + docstrings + graph context (cross-file deps, inheritance, callers/callees, type refs).

Use for: Open questions ("how does auth work?"), unknown patterns, architecture exploration. Natural language queries.
Don't use for: Single symbol/exact phrase (use find_code_references or grep then read_file_lines_tool or read), exploring known entity relationships.

rerank=False: fast, 65% accuracy | rerank=True: slower, 95% accuracy

Output includes: code content, docs, cross_file:[calls/extends/uses@filepath], used_by:[entity@filepath], file:[same-file entities]

Example: semantic_rag_search("user authentication", 5, True) → 5 results with code + "cross_file:calls:ValidateToken@auth/jwt.go|used_by:setupAuth@server/routes.go" """
    if vectorstore is None or bm25_corpus is None:
        return "❌ Index not built. Please call init_repo() or rebuild_index() first."

    graph = LocalGraph(root_path=str(CODEBASE_PATH))
    graph_loaded = graph.load()

    if not graph_loaded:
        logger.warning("Graph not loaded - search will proceed without graph context")

    try:
        hy = hybrid_search(query, k=max(top_k, RERANK_TOP_N))

        if ENABLE_RERANK and rerank:
            try:
                candidates = [t for t, m, s in hy]
                rr = rerank_with_ollama_enhanced(query, candidates[:RERANK_TOP_N], top_k=top_k)

                # Build enhanced results with graph context
                results = []
                for i, (chunk_text, score) in enumerate(rr, 1):
                    meta = next((m for t, m, s in hy if t == chunk_text), {})
                    results.append(_build_enhanced_result(chunk_text, meta, score, graph))

                return _format_enhanced_results(results, query, "search")

            except Exception as e:
                logger.warning(f"Rerank step failed: {e}")
                # Fall through to non-reranked results

        # Non-reranked results with graph context
        results = []
        for i, (text, meta, score) in enumerate(hy[:top_k], 1):
            results.append(_build_enhanced_result(text, meta, score, graph))

        return _format_enhanced_results(results, query, "search")

    except RuntimeError as e:
        return f"Error: {e}"
    except Exception as e:
        logger.exception("Search failed")
        return f"Search failed: {e}"


def _build_enhanced_result(chunk_text: str, meta: dict, score: float, graph: Optional[LocalGraph] = None) -> dict:
    """Build enhanced result with graph context and important metadata."""
    # Extract key metadata
    file_path = meta.get('file', '')
    line_num = meta.get('line', '')
    entity_type = meta.get('type', 'code')
    language = meta.get('language', '')
    entity_name = meta.get('name', '')

    # Extract clean content and docstring
    clean_content = _extract_clean_content(chunk_text)
    docstring = _extract_docstring(chunk_text)

    # Get graph context for this result
    graph_context = _get_result_graph_context(meta, graph)

    return {
        'file': file_path,
        'line': line_num,
        'type': entity_type,
        'language': language,
        'name': entity_name,
        'score': score,
        'content': clean_content,
        'docstring': docstring,
        'graph_context': graph_context
    }


def _get_result_graph_context(meta: dict, graph: Optional[LocalGraph] = None) -> dict:
    """Ultra-compact graph context using LocalGraph helper methods."""
    if graph is None:
        return {}

    file_path = meta.get('file', '')
    entity_name = meta.get('name', '')
    entity_type = meta.get('type', '')
    language = meta.get('language', '')

    # Build node ID using graph's method
    node_attrs = {"lang": language, "type": entity_type, "file": file_path, "name": entity_name}
    node_id = graph.add_node(node_attrs, return_id=True)

    if node_id not in graph.graph:
        return {}

    # Get all relationships in one call
    relationships = graph.get_node_relationships(node_id)
    context = {}

    if relationships.get('cross_file'):
        context['cross_file'] = relationships['cross_file']
    if relationships.get('same_file'):
        context['same_file'] = relationships['same_file']

    # Build reverse dependencies from incoming relationships
    if relationships.get('incoming'):
        reverse_deps = []
        node_file = graph.graph.nodes[node_id].get('file', '')
        for pred_id, edge_type, pred_data in relationships['incoming']:
            pred_file = pred_data.get('file', '')
            pred_name = pred_data.get('name', '')
            pred_type = pred_data.get('type', '')

            if pred_file and pred_file != node_file:
                reverse_deps.append(f"{edge_type}:{pred_type}:{pred_name}@{pred_file}")
            else:
                reverse_deps.append(f"{edge_type}:{pred_type}:{pred_name}")

        if reverse_deps:
            context['used_by'] = reverse_deps

    # Get file entities
    file_entities = graph.get_file_entities(file_path)
    if file_entities:
        context['file_contents'] = file_entities

    return context

def _format_enhanced_results(results: List[Dict], query: str, result_type: str) -> str:
    """Comprehensive format with ALL information preserved - optimized for LLM analysis."""
    if not results:
        return f"# {result_type.title()}: {query}\nNo results found.\n"

    lines = [
        f"# {result_type.title()}: {query}",
        f"results_count={len(results)}"
    ]

    for i, result in enumerate(results, 1):
        # Extract ALL metadata
        file_path = result.get('file', '')
        line_num = result.get('line', '')
        entity_type = result.get('type', 'code')
        entity_name = result.get('name', '')
        language = result.get('language', '')
        score = result.get('score', 0.0)
        content = result.get('content', '')  # NO TRUNCATION - let LLM handle it
        docstring = result.get('docstring', '')

        graph_context = result.get('graph_context', {})

        # Build comprehensive result block
        lines.append(f"\n--- RESULT {i} ---")
        lines.append(f"file:{file_path}:{line_num}")
        lines.append(f"entity:{entity_type}:{entity_name}")
        lines.append(f"language:{language}")
        lines.append(f"score:{score:.3f}")

        # Content (FULL, no truncation)
        if content:
            lines.append(f"content: {content}")

        # Docstring if available
        if docstring:
            lines.append(f"doc: {docstring}")

        # Graph context - PRESERVE EVERYTHING
        if graph_context:
            lines.append("graph_context:")

            # Cross-file dependencies (highest value)
            if graph_context.get('cross_file'):
                lines.append(f"  cross_file: {' | '.join(graph_context['cross_file'])}")

            # Reverse dependencies
            if graph_context.get('used_by'):
                lines.append(f"  used_by: {' | '.join(graph_context['used_by'])}")

            # Same-file relationships
            if graph_context.get('same_file'):
                lines.append(f"  same_file_rels: {' | '.join(graph_context['same_file'])}")

            # All relationships (comprehensive)
            if graph_context.get('all_rels'):
                lines.append(f"  all_relationships: {' | '.join(graph_context['all_rels'][:15])}")  # Reasonable limit

            # File contents
            if graph_context.get('file_contents'):
                lines.append(f"  file_entities: {' | '.join(graph_context['file_contents'])}")

    return "\n".join(lines)


def _extract_clean_content(chunk_text: str) -> str:
    """Extract clean content WITHOUT truncation - preserve full context."""
    lines = chunk_text.split('\n')
    content_lines = []

    # Skip only the pure metadata header lines
    skip_metadata = True
    for line in lines:
        stripped = line.strip()
        if skip_metadata:
            # Only skip actual metadata labels, not content
            if any(stripped.startswith(prefix) for prefix in ['File:', 'Type:', 'Name:', 'Language:', 'Doc:']):
                continue
            # Stop skipping when we hit actual code/content
            if stripped and not stripped.startswith(('File:', 'Type:', 'Name:', 'Language:', 'Doc:')):
                skip_metadata = False
                content_lines.append(stripped)
        else:
            content_lines.append(stripped)

    # Join with spaces but preserve ALL content
    content = ' '.join(content_lines)
    content = re.sub(r'\s+', ' ', content)  # Only normalize whitespace
    return content.strip()


def _build_enhanced_result(chunk_text: str, meta: dict, score: float, graph: Optional[LocalGraph] = None) -> dict:
    """Build comprehensive result with ALL context preserved."""
    # Extract key metadata
    file_path = meta.get('file', '')
    line_num = meta.get('line', '')
    entity_type = meta.get('type', 'code')
    language = meta.get('language', '')
    entity_name = meta.get('name', '')

    # Extract FULL content and docstring
    clean_content = _extract_clean_content(chunk_text)
    docstring = _extract_docstring(chunk_text)

    # Get COMPREHENSIVE graph context
    graph_context = _get_result_graph_context(meta, graph)

    return {
        'file': file_path,
        'line': line_num,
        'type': entity_type,
        'language': language,
        'name': entity_name,
        'score': score,
        'content': clean_content,  # NO TRUNCATION
        'docstring': docstring,
        'graph_context': graph_context
    }

def _extract_docstring(chunk_text: str) -> str:
    """Extract docstring or comments from chunk text."""
    lines = chunk_text.split('\n')
    doc_lines = []

    # Look for docstring patterns
    in_docstring = False
    for line in lines:
        stripped = line.strip()

        # Java/Go style comments
        if stripped.startswith('//') or stripped.startswith('/*') or stripped.startswith('*'):
            doc_lines.append(stripped)
        # Python style docstrings
        elif '"""' in line or "'''" in line:
            if not in_docstring:
                in_docstring = True
            else:
                in_docstring = False
                break
        elif in_docstring:
            doc_lines.append(stripped)
        # Specific doc patterns in the metadata section
        elif stripped.startswith('Doc:') and len(stripped) > 4:
            doc_content = stripped[4:].strip()
            if doc_content and doc_content not in ['None', '""']:
                doc_lines.append(doc_content)

    # Clean up docstring
    if doc_lines:
        docstring = ' '.join(doc_lines)
        docstring = re.sub(r'\s+', ' ', docstring)
        return docstring.strip()[:300]  # Reasonable limit

    return ""


def _format_enhanced_results(results: List[Dict], query: str, result_type: str) -> str:
    """Comprehensive format with ALL information preserved - optimized for LLM analysis."""
    if not results:
        return f"# {result_type.title()}: {query}\nNo results found.\n"

    lines = [
        f"# {result_type.title()}: {query}",
        f"results_count={len(results)}"
    ]

    for i, result in enumerate(results, 1):
        # Extract ALL metadata
        file_path = result.get('file', '')
        line_num = result.get('line', '')
        entity_type = result.get('type', 'code')
        entity_name = result.get('name', '')
        language = result.get('language', '')
        score = result.get('score', 0.0)
        content = result.get('content', '')  # NO TRUNCATION - let LLM handle it
        docstring = result.get('docstring', '')

        graph_context = result.get('graph_context', {})

        # Build comprehensive result block
        lines.append(f"\n--- RESULT {i} ---")
        lines.append(f"file:{file_path}:{line_num}")
        lines.append(f"entity:{entity_type}:{entity_name}")
        lines.append(f"language:{language}")
        lines.append(f"score:{score:.3f}")

        # Content (FULL, no truncation)
        if content:
            lines.append(f"content: {content}")

        # Docstring if available
        if docstring:
            lines.append(f"doc: {docstring}")

        # Graph context - PRESERVE EVERYTHING
        if graph_context:
            lines.append("graph_context:")

            # Cross-file dependencies (highest value)
            if graph_context.get('cross_file'):
                lines.append(f"  cross_file: {' | '.join(graph_context['cross_file'])}")

            # Reverse dependencies
            if graph_context.get('used_by'):
                lines.append(f"  used_by: {' | '.join(graph_context['used_by'])}")

            # Same-file relationships
            if graph_context.get('same_file'):
                lines.append(f"  same_file_rels: {' | '.join(graph_context['same_file'])}")

            # All relationships (comprehensive)
            if graph_context.get('all_rels'):
                lines.append(f"  all_relationships: {' | '.join(graph_context['all_rels'][:15])}")  # Reasonable limit

            # File contents
            if graph_context.get('file_contents'):
                lines.append(f"  file_entities: {' | '.join(graph_context['file_contents'])}")

    return "\n".join(lines)


def _extract_clean_content(chunk_text: str) -> str:
    """Extract clean content WITHOUT truncation - preserve full context."""
    lines = chunk_text.split('\n')
    content_lines = []

    # Skip only the pure metadata header lines
    skip_metadata = True
    for line in lines:
        stripped = line.strip()
        if skip_metadata:
            # Only skip actual metadata labels, not content
            if any(stripped.startswith(prefix) for prefix in ['File:', 'Type:', 'Name:', 'Language:', 'Doc:']):
                continue
            # Stop skipping when we hit actual code/content
            if stripped and not stripped.startswith(('File:', 'Type:', 'Name:', 'Language:', 'Doc:')):
                skip_metadata = False
                content_lines.append(stripped)
        else:
            content_lines.append(stripped)

    # Join with spaces but preserve ALL content
    content = ' '.join(content_lines)
    content = re.sub(r'\s+', ' ', content)  # Only normalize whitespace
    return content.strip()

@mcp.tool()
def find_code_references(symbol: str, top_k: int = 20) -> str:
    """What it does - Fast, unranked lookup of all file/line occurrences of a symbol.
When to use - After identifying a symbol in a search hit, or when you need every use of a class/method.
When not to use - When you need context-rich code; use read_file_lines_tool after getting the locations.
Example - find_code_references("AuthService") → list of references."""
    matches = search_symbol(symbol, CODEBASE_PATH)

    if not matches:
        return f"# References: {symbol}\nNo references found.\n"

    # Build clean TOON format directly
    results = []
    for i, m in enumerate(matches[:top_k], 1):
        results.append({
            'file': m['file'],
            'line': m['line'],
            'context': m['context']
        })

    return _format_reference_results(results, symbol)


def _format_reference_results(results: List[Dict], symbol: str) -> str:
    """Format reference results in clean TOON format."""
    if not results:
        return f"# References: {symbol}\nNo references found.\n"

    lines = [f"# References: {symbol}", f"references[{len(results)}]{{file,line,context}}:"]

    for result in results:
        row = [
            result.get('file', ''),
            result.get('line', ''),
            result.get('context', '')[:300]  # Reasonable limit
        ]

        # Escape fields if needed
        escaped_row = []
        for field in row:
            field_str = str(field)
            if any(char in field_str for char in [',', '"', '\n', '\r']):
                escaped = field_str.replace('"', '\\"')
                escaped_row.append(f'"{escaped}"')
            else:
                escaped_row.append(field_str)

        lines.append("  " + ",".join(escaped_row))

    return "\n".join(lines)

@mcp.tool()
def read_file_lines_tool(path: str, start: int = 1, end: Optional[int] = None) -> str:
    """What it does - Reads a requested line range and auto-expands to include surrounding context (function body, docstring, comments).
When to use - To inspect a specific hit from find_code_references or a snippet from semantic_rag_search.
When not to use - For browsing the entire file or unrelated sections; use semantic_rag_search first to pinpoint relevant parts.
Example - read_file_lines_tool("src/auth/login.py", start=45, end=60) → full function with context.
Note: This tool treats / as '/home/popertots/Crussell/' so adjust paths accordingly"""
    return EnhancedToon.file_content_results(read_file_lines(path, start, end), path)

@mcp.tool()
def init_repo(git_url: str) -> str:
    """What it does - Initialise the RAG stack with a brand-new Git repository."""
    import shutil
    from pathlib import Path
    import subprocess

    # ----- 1. Always wipe and recreate working_repo ---------------------
    if CODEBASE_PATH.exists():
        shutil.rmtree(CODEBASE_PATH)
        logger.info(f"Wiped existing working repo: {CODEBASE_PATH}")

    CODEBASE_PATH.mkdir(exist_ok=True)

    # ----- 2. Clone fresh -----------------------------------------------
    try:
        subprocess.run(
            ["git", "clone", git_url, str(CODEBASE_PATH)],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        logger.info(f"Cloned {git_url} to {CODEBASE_PATH}")
    except subprocess.CalledProcessError as e:
        return f"❌ Failed to clone {git_url}: {e.stderr.strip()}"

    # ----- 3. Reset ALL global state ------------------------------------
    global vectorstore, bm25, bm25_corpus, chunks_metadata, index_build_time

    vectorstore = None
    bm25 = None
    bm25_corpus = None
    chunks_metadata = None
    index_build_time = 0

    # ----- 4. Remove any stale indices ----------------------------------
    for artefact in ("chroma_db", "bm25_index.json", "embeddings", ".mcp_cache"):
        artefact_path = CODEBASE_PATH / artefact
        if artefact_path.exists():
            try:
                if artefact_path.is_dir():
                    shutil.rmtree(artefact_path)
                else:
                    artefact_path.unlink()
                logger.info(f"Removed stale artefact: {artefact_path}")
            except Exception as e:
                logger.warning(f"Could not delete {artefact_path}: {e}")

    # ----- 5. Re-build the index ----------------------------------------
    try:
        build_indexes()
        logger.info("Index rebuilt successfully")
    except Exception as e:
        logger.exception("Index rebuild failed")
        return f"❌ Index rebuild failed: {e}"

    # ----- 6. Report success --------------------------------------------
    try:
        branch_name = subprocess.check_output(
            ["git", "-C", str(CODEBASE_PATH), "rev-parse", "--abbrev-ref", "HEAD"],
            text=True,
        ).strip()
        git_hash = subprocess.check_output(
            ["git", "-C", str(CODEBASE_PATH), "rev-parse", "HEAD"],
            text=True,
        ).strip()

        # Verify the index was actually built
        if vectorstore is None or bm25_corpus is None:
            return f"⚠️ Repository cloned but index failed to build properly"

        return f"✅ Initialized repository at {CODEBASE_PATH}\n   Branch: {branch_name}\n   Commit: {git_hash[:8]}\n   Indexed chunks: {len(bm25_corpus)}"

    except Exception as e:
        return f"⚠️ Repository cloned but status check failed: {e}"

@mcp.tool()
def rebuild_index() -> str:
    """What it does - Re-creates all embeddings, BM25, and metadata indexes after major code changes.
When to use - When the codebase has been pulled or refactored and you suspect stale search results, and only after explicit instruction to do so.
When not to use - On every query; it's expensive and unnecessary if the index is already up-to-date.
Example - rebuild_index() → '✅ Index rebuilt (13.2 s); 11,234 chunks'"""
    with _startup_lock:
        try:
            # Clear LRU caches
            parse_python_file_cached.cache_clear()
            parse_sql_file_cached.cache_clear()
            parse_go_file_cached.cache_clear()
            parse_java_file_cached.cache_clear()
            parse_svelte_file_cached.cache_clear()
            parse_shell_file_cached.cache_clear()

            LocalGraph(root_path=CODEBASE_PATH).clear()

            t0 = time.time()
            build_indexes()
            dt = time.time() - t0
            return f"✅ Index rebuilt successfully! (took {dt:.1f}s)\n\nStatistics:\n" + json.dumps({
                "total_chunks": len(bm25_corpus),
                "by_language": count_by_language(chunks_metadata),
                "by_type": count_by_type(chunks_metadata)
            }, indent=2)
        except Exception as e:
            logger.exception("Index rebuild failed")
            return f"❌ Rebuild failed: {e}"

# -----------------------------
# Startup
# -----------------------------

def signal_handler(sig, frame):
    """Handle graceful shutdown on Ctrl+C."""
    logger.info("\n=== Shutdown signal received ===")

    # Clear any LRU caches to release memory
    try:
        parse_python_file_cached.cache_clear()
        parse_sql_file_cached.cache_clear()
        parse_go_file_cached.cache_clear()
        parse_java_file_cached.cache_clear()
        parse_svelte_file_cached.cache_clear()
        parse_shell_file_cached.cache_clear()
        logger.info("✓ Caches cleared")
    except Exception as e:
        logger.warning(f"Cache clearing error: {e}")

    # Note: Chroma auto-persists, no explicit close needed
    logger.info("✓ Indexes are already persisted to disk")
    logger.info("Goodbye!\n")

    # Force exit immediately to avoid thread hang
    os._exit(0)

if __name__ == "__main__":
    # Register signal handlers
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    try:
        logger.info("="*60)
        logger.info("🚀 MCP RAG Server is ready!")
        logger.info("Press Ctrl+C to stop")

        # Just run in stdio mode by default
        mcp.run(transport='stdio')

    except KeyboardInterrupt:
        signal_handler(signal.SIGINT, None)
    except Exception as e:
        logger.exception("Fatal error during startup")
        os._exit(1)