import os
import signal
import json
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

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

# -----------------------------
# Configuration
# -----------------------------
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
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(os.environ.get("CODEBASE_PATH", "."))  # change as needed
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.2,    # Class definitions
    "python_function": 1.15, # Function definitions

    # Go weights
    "go_type": 1.2,         # Type definitions
    "go_function": 1.15,    # Function definitions

    # Java weights
    "java_class": 1.2,      # Class definitions
    "java_method": 1.15,    # Method definitions

    # 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()

    # 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
    elif language == "python":
        if chunk_type == "classdef":
            return CONTEXTUAL_WEIGHTS.get("python_class", 1.0)
        elif chunk_type in ("functiondef", "asyncfunctiondef"):
            return CONTEXTUAL_WEIGHTS.get("python_function", 1.0)

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

    # Java-specific weights
    elif language == "java":
        if chunk_type == "class":
            return CONTEXTUAL_WEIGHTS.get("java_class", 1.0)
        elif chunk_type == "method":
            return CONTEXTUAL_WEIGHTS.get("java_method", 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 = []

    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. Returns tuple for hashability."""
    filepath = Path(filepath_str)
    chunks = []
    try:
        code = filepath.read_text(encoding="utf-8")
        tree = ast.parse(code)
        lines = code.splitlines(keepends=True)
        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.ClassDef, ast.AsyncFunctionDef)):
                docstring = ast.get_docstring(node) or ""
                start_line = node.lineno - 1
                end_line = getattr(node, "end_lineno", start_line + 1)
                chunk_code = "".join(lines[start_line:end_line])
                chunks.append({
                    "text": f"File: {filepath}\nType: {type(node).__name__}\nName: {getattr(node,'name', '<anon>')}\nDocstring: {docstring}\n\nCode:\n{chunk_code}",
                    "metadata": {
                        "file": str(filepath),
                        "name": getattr(node, "name", ""),
                        "type": type(node).__name__,
                        "line": node.lineno,
                        "language": "python",
                    }
                })
    except Exception as e:
        logger.warning(f"parse_python_file: failed {filepath}: {e}")
    return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)

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)
    """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:
        # find position of this statement in the original SQL (using running search_pos)
        start_idx = sql_code.find(stmt, search_pos)
        if start_idx == -1:
            # fallback: try from beginning
            start_idx = sql_code.find(stmt)
        if start_idx == -1:
            # cannot locate, approximate line number as last line
            start_line = len(lines)
        else:
            start_line = sql_code[:start_idx].count("\n") + 1
            search_pos = start_idx + len(stmt)

        # include leading comments above this statement
        leading_comments = include_leading_comments(lines, start_line)

        # Trim and normalize statement
        stmt_clean = stmt.strip()
        stmt_text_for_embed = (leading_comments + "\n\n" + stmt_clean) if leading_comments else stmt_clean

        # Try to parse with sqlglot to extract tables / functions / ctes metadata
        meta: Dict[str, Any] = {
            "file": str(filepath),
            "name": None,
            "type": None,
            "line": start_line,
            "language": "sql"
        }
        try:
            parsed = sqlglot.parse_one(stmt_clean, read="postgres")
            # stmt type
            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"

            # table names (if any)
            tables = [t.this for t in parsed.find_all(sqlglot.exp.Table)]
            meta["tables"] = [t for t in tables if isinstance(t, str)]
            # ctes
            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"] = ctes

            # functions called in the statement (names)
            funcs = []
            for f in parsed.find_all(sqlglot.exp.Func):
                try:
                    name = f.name
                    if name:
                        funcs.append(name)
                except Exception:
                    pass
            meta["functions"] = funcs

            # choose a sensible name for the metadata
            if meta.get("tables"):
                meta["name"] = meta["tables"][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\nComments:\n{leading_comments}\n\nSQL:\n{stmt_clean}" if leading_comments else f"File: {filepath}\nType: {meta.get('type')}\nName: {meta.get('name')}\n\nSQL:\n{stmt_clean}"

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

            continue
        except Exception:
            # sqlglot failed on this statement — fall back to heuristics (still keep comments)
            pass

        # Fallback: simple heuristics if sqlglot fails
        # Determine an approximate type by the first token
        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\nComments:\n{leading_comments}\n\nSQL:\n{stmt_clean}" if leading_comments else f"File: {filepath}\nType: {meta['type']}\nName: {meta.get('name')}\n\nSQL:\n{stmt_clean}"

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

    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 AST parsing. Returns tuple for hashability."""
    filepath = Path(filepath_str)
    chunks = []
    helper = Path("./tools/parse_go_ast")

    if helper.exists():
        try:
            proc = subprocess.run(
                [str(helper), str(filepath)],
                capture_output=True,
                text=True,
                check=True,
                timeout=20
            )
            decls = json.loads(proc.stdout)
            lines = filepath.read_text(encoding="utf-8").splitlines(keepends=True)

            for d in decls:
                # Use lowercase keys (matching the Go JSON output)
                start = max(0, d.get("start_line", 1) - 1)
                end = d.get("end_line", start + 1)
                name = d.get("name", "")
                typ = d.get("type", "")
                doc_comment = d.get("doc_comment", "")
                receiver = d.get("receiver", "")

                # Skip package declarations
                if typ == "package":
                    continue

                # Build display type
                display_type = typ
                if receiver:
                    display_type = f"method ({receiver})"

                chunk_code = "".join(lines[start:end])

                chunk_text = f"File: {filepath}\nType: {display_type}\nName: {name}\n"
                if doc_comment:
                    chunk_text += f"Doc:\n{doc_comment}\n\n"
                chunk_text += f"Code:\n{chunk_code}"

                chunks.append({
                    "text": chunk_text,
                    "metadata": {
                        "file": str(filepath),
                        "name": name,
                        "type": typ,
                        "receiver": receiver,
                        "line": start + 1,
                        "language": "go",
                    }
                })

            if chunks:  # If helper worked, return its results
                return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)

        except subprocess.CalledProcessError as e:
            logger.warning(f"parse_go_ast failed for {filepath}: {e.stderr}; falling back to regex")
        except json.JSONDecodeError as e:
            logger.warning(f"parse_go_ast output invalid JSON for {filepath}: {e}; falling back to regex")
        except Exception as e:
            logger.warning(f"parse_go_file helper failed for {filepath}: {e}; falling back to regex")

    # Fallback: Use regex-based parsing (same as before)
    try:
        content = filepath.read_text(encoding="utf-8")
        lines = content.splitlines(keepends=True)

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

        for match in func_pattern.finditer(content):
            func_name = match.group(1)
            start_pos = match.start()
            start_line = content[:start_pos].count('\n')
            end_line = find_brace_block_end_go(lines, start_line)
            func_code = "".join(lines[start_line:end_line + 1])
            comments = extract_go_comments(lines, start_line)

            chunk_text = f"File: {filepath}\nType: function\nName: {func_name}\n"
            if comments:
                chunk_text += f"Comments:\n{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": "go",
                }
            })

        # Structs
        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_go(lines, start_line)
            struct_code = "".join(lines[start_line:end_line + 1])
            comments = extract_go_comments(lines, start_line)

            chunk_text = f"File: {filepath}\nType: struct\nName: {struct_name}\n"
            if comments:
                chunk_text += f"Comments:\n{comments}\n\n"
            chunk_text += f"Code:\n{struct_code}"

            chunks.append({
                "text": chunk_text,
                "metadata": {
                    "file": str(filepath),
                    "name": struct_name,
                    "type": "struct",
                    "line": start_line + 1,
                    "language": "go",
                }
            })

        # Interfaces
        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_go(lines, start_line)
            interface_code = "".join(lines[start_line:end_line + 1])
            comments = extract_go_comments(lines, start_line)

            chunk_text = f"File: {filepath}\nType: interface\nName: {interface_name}\n"
            if comments:
                chunk_text += f"Comments:\n{comments}\n\n"
            chunk_text += f"Code:\n{interface_code}"

            chunks.append({
                "text": chunk_text,
                "metadata": {
                    "file": str(filepath),
                    "name": interface_name,
                    "type": "interface",
                    "line": start_line + 1,
                    "language": "go",
                }
            })

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

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


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. Returns tuple for hashability."""
    filepath = Path(filepath_str)
    chunks = []
    try:
        code = filepath.read_text(encoding="utf-8")
        lines = code.splitlines(keepends=True)
        tree = javalang.parse.parse(code)

        def find_brace_block_end(start_idx: int) -> int:
            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
            return len(lines) - 1

        for path, node in tree:
            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 + 1])
                doc = get_context_with_comments(filepath, start_line + 1)
                chunks.append({
                    "text": f"File: {filepath}\nType: Class\nName: {node.name}\nDoc:\n{doc}\nCode:\n{chunk_text}",
                    "metadata": {
                        "file": str(filepath),
                        "name": node.name,
                        "type": "class",
                        "line": start_line + 1,
                        "language": "java"
                    }
                })

            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 + 1])
                doc = get_context_with_comments(filepath, start_line + 1)
                # Find enclosing class name
                class_name = None
                for p in reversed(path):
                    if isinstance(p, javalang.tree.ClassDeclaration):
                        class_name = p.name
                        break
                chunks.append({
                    "text": f"File: {filepath}\nType: Method\nClass: {class_name or 'unknown'}\nName: {node.name}\nDoc:\n{doc}\nCode:\n{chunk_text}",
                    "metadata": {
                        "file": str(filepath),
                        "name": node.name,
                        "type": "method",
                        "class": class_name,
                        "line": start_line + 1,
                        "language": "java"
                    }
                })

            elif isinstance(node, javalang.tree.FieldDeclaration):
                # Fields may have multiple variable declarators
                if not hasattr(node, 'position') or not node.position:
                    continue
                start_line = node.position.line - 1
                # Field is usually one line, but include next if annotation
                end_line = start_line
                chunk_text = lines[start_line]
                doc = get_context_with_comments(filepath, start_line + 1)
                for declarator in node.declarators:
                    chunks.append({
                        "text": f"File: {filepath}\nType: Field\nName: {declarator.name}\nDoc:\n{doc}\nCode:\n{chunk_text}",
                        "metadata": {
                            "file": str(filepath),
                            "name": declarator.name,
                            "type": "field",
                            "line": start_line + 1,
                            "language": "java"
                        }
                    })

            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 + 1])
                doc = get_context_with_comments(filepath, start_line + 1)
                class_name = None
                for p in reversed(path):
                    if isinstance(p, javalang.tree.ClassDeclaration):
                        class_name = p.name
                        break
                chunks.append({
                    "text": f"File: {filepath}\nType: Constructor\nClass: {class_name or 'unknown'}\nDoc:\n{doc}\nCode:\n{chunk_text}",
                    "metadata": {
                        "file": str(filepath),
                        "name": f"{class_name} constructor",
                        "type": "constructor",
                        "class": class_name,
                        "line": start_line + 1,
                        "language": "java"
                    }
                })

            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 + 1])
                doc = get_context_with_comments(filepath, start_line + 1)
                chunks.append({
                    "text": f"File: {filepath}\nType: Enum\nName: {node.name}\nDoc:\n{doc}\nCode:\n{chunk_text}",
                    "metadata": {
                        "file": str(filepath),
                        "name": node.name,
                        "type": "enum",
                        "line": start_line + 1,
                        "language": "java"
                    }
                })

    except javalang.parser.JavaSyntaxError as e:
        logger.warning(f"Failed to parse Java file {filepath}: {e}")
    except Exception as e:
        logger.warning(f"Unexpected error parsing Java file {filepath}: {e}")
    return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)

def parse_java_file(filepath: Path) -> List[Dict]:
    """Parse Java source file with caching."""
    file_hash = get_file_hash(filepath)
    cached_result = parse_java_file_cached(str(filepath), file_hash)
    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. Returns tuple for hashability."""
    filepath = Path(filepath_str)
    chunks = []
    try:
        text = filepath.read_text(encoding="utf-8")
        lines = text.splitlines(keepends=True)
        script_matches = list(re.finditer(r"<script(?:\s+[^>]*)?>(.*?)</script>", text, flags=re.DOTALL))
        idx = 0
        for m in script_matches:
            script_content = m.group(1)
            script_start_line = text[:m.start(1)].count("\n") + 1
            ts_helper = Path("./tools/parse_ts.js")
            if ts_helper.exists():
                tmp = filepath.parent / f".tmp_{filepath.name}_{idx}.ts"
                try:
                    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)
                    for d in decls:
                        start_rel = max(0, d.get("start_line", 1) - 1)
                        end_rel = d.get("end_line", 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
                        chunk_text = f"File: {filepath}\nType: {d.get('type')}\nName: {d.get('name')}\nDoc: {doc}\nCode:\n{chunk_code}"
                        chunks.append({
                            "text": chunk_text,
                            "metadata": {
                                "file": str(filepath),
                                "name": d.get("name"),
                                "type": d.get("type"),
                                "line": abs_start,
                                "language": "typescript",
                                "block_type": "script"
                            }
                        })
                except Exception as e:
                    logger.warning(f"TS parser failed for {filepath}: {e}; falling back to heuristic")
                    chunks.append(_make_script_chunk(filepath, script_content, script_start_line, "script"))
                finally:
                    tmp.unlink(missing_ok=True)
            else:
                chunks.append(_make_script_chunk(filepath, script_content, script_start_line, "script"))
            idx += 1

        # Style blocks
        for m in re.finditer(r"<style(?:\s+[^>]*)?>(.*?)</style>", text, flags=re.DOTALL):
            style_content = m.group(1)
            style_start_line = text[:m.start(1)].count("\n") + 1
            chunks.append({
                "text": f"File: {filepath}\nType: style\n{style_content}",
                "metadata": {
                    "file": str(filepath),
                    "name": f"style_block_{idx}",
                    "type": "style",
                    "line": style_start_line,
                    "language": "css",
                    "block_type": "style"
                }
            })
            idx += 1

        # Markup
        markup = re.sub(r"<script[^>]*>.*?</script>", "", text, flags=re.DOTALL)
        markup = re.sub(r"<style[^>]*>.*?</style>", "", markup, flags=re.DOTALL).strip()
        if markup:
            chunks.append({
                "text": f"File: {filepath}\nType: markup\n{markup}",
                "metadata": {
                    "file": str(filepath),
                    "name": "markup",
                    "type": "markup",
                    "line": 1,
                    "language": "svelte",
                    "block_type": "markup"
                }
            })
    except Exception as e:
        logger.warning(f"parse_svelte_file failed {filepath}: {e}")
        return tuple()
    return tuple((c["text"], tuple(c["metadata"].items())) for c in chunks)

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]


def _make_script_chunk(filepath: Path, content: str, start_line: int, block_type: str = "script") -> Dict:
    return {
        "text": f"File: {filepath}\nType: {block_type}\n{content}",
        "metadata": {
            "file": str(filepath),
            "name": f"{block_type}_chunk",
            "type": block_type,
            "line": start_line,
            "language": "typescript" if block_type == "script" else "css",
            "block_type": block_type
        }
    }

# -----------------------------
# 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

@lru_cache(maxsize=1000)
def parse_shell_file_cached(filepath_str: str, file_hash: 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."""
    file_hash = get_file_hash(filepath)
    cached_result = parse_shell_file_cached(str(filepath), file_hash)
    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)
        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

# -----------------------------
# Build / load indexes with contextual weighting
# -----------------------------
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()

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

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

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

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

    # Create a wrapper that will return our pre-computed weighted embeddings
    class WeightedEmbeddingFunction:
        def __init__(self, weighted_embs):
            self.weighted_embs = weighted_embs
            self.idx = 0

        def embed_documents(self, texts):
            # Return the pre-computed weighted embeddings
            start = self.idx
            end = start + len(texts)
            result = self.weighted_embs[start:end]
            self.idx = end
            return result

        def batch_embed_documents(texts: List[str]) -> List[List[float]]:
            # Embed documents in batches
            all_embeddings = []
            for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
                batch = texts[i:i + EMBEDDING_BATCH_SIZE]
                embeddings = embeddings.embed_documents(batch)
                all_embeddings.extend(embeddings)
            return all_embeddings

        def embed_query(self, text):
            # For queries, use the base embedding model (no weighting)
            return embeddings.embed_query(text)

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

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

    # Save BM25 index and metadata
    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 with contextual weighting applied. Total time: {elapsed:.1f}s")

# -----------------------------
# 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

# -----------------------------
# 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")

    # Limit expansion to reasonable bounds (max 100 lines of context)
    max_context = 100
    if context_end - context_start > max_context:
        context_end = context_start + max_context
        context_info.append("truncated to 100 lines")

    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).
When to use - Before any other query, to confirm the index is current.
When not to use - After you already know the server is healthy; it adds no value.
Example - health_check() → { "status":"ready","total_chunks":11234,"index_age_hours":4.7 }"""
    try:
        with _startup_lock:
            status = {
                "status": "ready" if vectorstore is not None else "initializing",
                "codebase": str(CODEBASE_PATH),
                "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
                },
                "Tools": {
                    "NOTE": "THESE TOOLS ARE RESTRICTED BY .gitignore AS WELL AS .ragignore",
                    "search_codebase": search_codebase.__doc__,
                    "find_code_references": find_code_references.__doc__,
                    "read_file_lines_tool": read_file_lines_tool.__doc__
                }
            }

            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
                }

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

@mcp.tool()
def search_codebase(query: str, top_k: int = 5, rerank: bool = True) -> str:
    """What it does - Hybrid RAG search that returns semantically ranked snippets with file/line info from natural language queries.
When to use - For open-ended questions or unknown patterns (“how does auth work?”). Natural language and verbose queries ONLY.
When not to use - For a single symbol or exact phrase; use find_code_references then read instead to avoid a large result set.
Rerank false is good enough for 65% of searches and is VERY fast, rerank true is much slower but good enough for 95% of searches. Best to use rerank true if rerank false didn't quite help enough.
Example - search_codebase("user authentication python", top_k=5, rerank=True) → 5 relevant snippets."""
    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 clean TOON format directly
                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({
                        'file': meta.get('file', ''),
                        'line': meta.get('line', ''),
                        'type': meta.get('type', 'code'),
                        'score': score,
                        'content': _extract_clean_content(chunk_text)
                    })

                return _format_toon_results(results, query, "search")

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

        # Non-reranked results - build clean TOON format directly
        results = []
        for i, (text, meta, score) in enumerate(hy[:top_k], 1):
            results.append({
                'file': meta.get('file', ''),
                'line': meta.get('line', ''),
                'type': meta.get('type', 'code'),
                'score': score,
                'content': _extract_clean_content(text)
            })

        return _format_toon_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 _extract_clean_content(chunk_text: str) -> str:
    """Extract clean content from chunk text by removing duplicate metadata."""
    lines = chunk_text.split('\n')
    content_lines = []

    # Skip the first few metadata lines (File:, Type:, Name:, etc.)
    skip_metadata = True
    for line in lines:
        stripped = line.strip()
        if skip_metadata:
            if any(stripped.startswith(prefix) for prefix in ['File:', 'Type:', 'Name:', 'Language:', 'Doc:']):
                continue
            # Once we hit actual content, stop skipping
            if stripped and not any(stripped.startswith(prefix) for prefix in ['File:', 'Type:', 'Name:', 'Language:', 'Doc:', 'SQL:', 'Code:']):
                skip_metadata = False
                content_lines.append(stripped)
        else:
            content_lines.append(stripped)

    # Join and clean up
    content = ' '.join(content_lines)
    content = re.sub(r'\s+', ' ', content)  # Normalize whitespace
    return content.strip()


def _format_toon_results(results: List[Dict], query: str, result_type: str) -> str:
    """Format results in clean TOON format."""
    if not results:
        return f"# {result_type.title()}: {query}\nNo results found.\n"

    lines = [f"# {result_type.title()}: {query}", f"results[{len(results)}]{{file,line,score,content}}:"]

    for result in results:
        row = [
            result.get('file', ''),
            result.get('line', ''),
            f"{result.get('score', 0.0):.2f}",
            result.get('content', '')[:500]  # 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 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 search_codebase.
When not to use - For browsing the entire file or unrelated sections; use search_codebase 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 or update an existing one.

    The function will:
    1. Clone the repo if it does not already exist locally,
    2. Pull the latest changes if it does,
    3. Reset the global `CODEBASE_PATH` to the local clone,
    4. Remove any stale index artefacts,
    5. Re-build the vector/BM25 indexes with the same lock-protected flow used by `rebuild_index`,
    6. Return a human-readable status message.

    Next steps:
        • Call `health_check` to confirm that the server is ready.
        • Use `search_codebase` or `find_code_references` to explore the new code.
        • Use `read_file_lines_tool` to inspect any hit in detail.
    """
    import shutil
    from pathlib import Path
    import subprocess

    # ----- 1. Determine where to clone ---------------------------------
    repo_name = git_url.split('/')[-1]
    if repo_name.endswith('.git'):
        repo_name = repo_name[:-4]
    clone_dir = Path.cwd() / repo_name

    # ----- 2. Clone / pull -----------------------------------------------
    try:
        if clone_dir.exists():
            # Existing repo → pull
            subprocess.run(
                ["git", "-C", str(clone_dir), "pull"],
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
        else:
            # New repo → clone
            subprocess.run(
                ["git", "clone", git_url],
                cwd=str(Path.cwd()),
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
    except subprocess.CalledProcessError as e:
        return f"❌ Failed to clone/pull {git_url}: {e.stderr.strip()}"

    # ----- 3. Update the global CODEBASE_PATH --------------------------------
    global CODEBASE_PATH
    CODEBASE_PATH = clone_dir

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

    # ----- 5. Re-build the index (protected by the startup lock) ------------
    try:
        rebuild_index()
    except Exception as e:
        return f"❌ Index rebuild failed: {e}"

    # ----- 6. Report success -----------------------------------------------
    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()

    return f"✅ Initialized repository at {CODEBASE_PATH} on branch {branch_name} with hash {git_hash}"


@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()

            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("Use 'python serve_http.py' for HTTP server")
        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)