ingest/embed.py

"""
embedder.py
-----------
Converts parsed codebase (FileInfo objects) into hybrid chunks and
stores them in ChromaDB across two collections:

    - class_chunks    : one chunk per class (for macro / cross-module queries)
    - function_chunks : one chunk per function/method (for micro queries)

Each chunk carries rich metadata so the retriever can filter precisely.

Depends on:
    - ast_parser.parse_codebase() → list[FileInfo]
    - chromadb
    - sentence-transformers (local embedding, no API needed)

Install:
    pip install chromadb sentence-transformers rich
"""

import json
import hashlib
from pathlib import Path

from pinecone import Pinecone, ServerlessSpec
from sentence_transformers import SentenceTransformer
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn
from rich.panel import Panel
from rich.table import Table
from rich import box
from config.config import PINECONE_API_KEY, PINECONE_INDEX, EMBEDDING_DIM
from ingest.parse_ast import parse_codebase, FileInfo, ClassInfo, FunctionInfo

console = Console()

# ── Constants ─────────────────────────────────────────────────────────────────

EMBEDDING_MODEL  = "all-MiniLM-L6-v2"    # fast, lightweight, good quality
CLASS_COLLECTION    = "class_chunks"
FUNCTION_COLLECTION = "function_chunks"
COMPLETE_COLLECTION = "complete_chunks"


# ── Embedding Model ───────────────────────────────────────────────────────────

def load_embedding_model() -> SentenceTransformer:
    """Load the sentence transformer embedding model."""
    with console.status("[bold cyan]Loading embedding model...[/bold cyan]"):
        model = SentenceTransformer(EMBEDDING_MODEL)
    console.print(f"[green]✔[/green] Embedding model loaded: [cyan]{EMBEDDING_MODEL}[/cyan]")
    return model


# ── ChromaDB Client ───────────────────────────────────────────────────────────

def get_pinecone_index():
    """Return a Pinecone index, creating it if it does not exist."""
    pc = Pinecone(api_key=PINECONE_API_KEY)
    existing = [i.name for i in pc.list_indexes()]
    if PINECONE_INDEX not in existing:
        pc.create_index(
            name=PINECONE_INDEX,
            dimension=EMBEDDING_DIM,
            metric="cosine",
            spec=ServerlessSpec(cloud="aws", region="us-east-1"),
        )
    return pc.Index(PINECONE_INDEX)


# ── Chunk Builders ────────────────────────────────────────────────────────────

def _make_id(text: str) -> str:
    """Generate a stable unique ID from chunk text."""
    return hashlib.md5(text.encode()).hexdigest()


def build_class_chunk(cls: ClassInfo, file_info: FileInfo) -> dict:
    """
    Build a class-level chunk document.
    Contains: class name, bases, docstring, all method signatures.
    """
    method_signatures = []
    for m in cls.methods:
        params = ", ".join(
            f"{p.name}: {p.annotation}" if p.annotation else p.name
            for p in m.parameters
        )
        ret = f" -> {m.return_type}" if m.return_type else ""
        method_signatures.append(f"  def {m.name}({params}){ret}")

    methods_block = "\n".join(method_signatures) if method_signatures else "  # no methods"
    bases_str = ", ".join(cls.bases) if cls.bases else "object"
    docstring = cls.docstring or "No docstring provided."

    text = (
        f"Class: {cls.name}\n"
        f"Inherits: {bases_str}\n"
        f"Module: {file_info.module}\n"
        f"File: {file_info.relative}\n"
        f"Docstring: {docstring}\n"
        f"Methods:\n{methods_block}"
    )

    metadata = {
        "type":       "class",
        "name":       cls.name,
        "module":     file_info.module,
        "file":       file_info.relative,
        "bases":      json.dumps(cls.bases),
        "methods":    json.dumps([m.name for m in cls.methods]),
        "lineno":     cls.lineno,
    }

    return {"id": _make_id(text), "text": text, "metadata": metadata}


def build_function_chunk(func: FunctionInfo,
                          file_info: FileInfo,
                          class_name: str | None = None,
                          class_docstring: str | None = None) -> dict:
    """
    Build a function/method-level chunk document.
    Carries class context as metadata so micro queries stay grounded.
    """
    params = ", ".join(
        f"{p.name}: {p.annotation}" if p.annotation else p.name
        for p in func.parameters
    )
    ret = f" -> {func.return_type}" if func.return_type else ""
    signature = f"def {func.name}({params}){ret}"

    docstring  = func.docstring or "No docstring provided."
    calls_str  = ", ".join(func.calls[:15]) if func.calls else "none"
    class_ctx  = (
        f"Class: {class_name}\nClass purpose: {class_docstring or 'N/A'}\n"
        if class_name else "Scope: top-level function\n"
    )

    text = (
        f"{class_ctx}"
        f"Function: {func.name}\n"
        f"Module: {file_info.module}\n"
        f"File: {file_info.relative}\n"
        f"Signature: {signature}\n"
        f"Docstring: {docstring}\n"
        f"Calls: {calls_str}"
    )

    metadata = {
        "type":            "function",
        "name":            func.name,
        "module":          file_info.module,
        "file":            file_info.relative,
        "class_name":      class_name or "",
        "return_type":     func.return_type or "",
        "parameters":      json.dumps([p.name for p in func.parameters]),
        "calls":           json.dumps(func.calls[:15]),
        "is_method":       str(func.is_method),
        "lineno":          func.lineno,
    }

    return {"id": _make_id(text), "text": text, "metadata": metadata}

def build_module_chunk(file_info: FileInfo) -> dict:
    """
    Build a module-level chunk for files that contain no classes or functions.
    Captures imports and docstring as the indexable content.
    """
    imports_str = ", ".join(file_info.imports) if file_info.imports else "none"
    docstring = file_info.docstring or "No module docstring."

    text = (
        f"Module: {file_info.module}\n"
        f"File: {file_info.relative}\n"
        f"Docstring: {docstring}\n"
        f"Imports: {imports_str}\n"
        f"Note: This file contains only module-level statements."
    )

    metadata = {
        "type":       "module",
        "name":       Path(file_info.relative).stem,
        "module":     file_info.module,
        "file":       file_info.relative,
        "class_name": "",
        "return_type": "",
        "parameters": "[]",
        "calls":      "[]",
        "is_method":  "False",
        "lineno":     0,
    }

    return {"id": _make_id(text), "text": text, "metadata": metadata}

def build_complete_function_chunk(func: FunctionInfo,
                                   file_info: FileInfo,
                                   class_name: str | None = None,
                                   class_docstring: str | None = None) -> dict:
    """
    Build a complete function chunk including full source code.
    Used for edge case analysis and usage example generation.
    """
    params = ", ".join(
        f"{p.name}: {p.annotation}" if p.annotation else p.name
        for p in func.parameters
    )
    ret = f" -> {func.return_type}" if func.return_type else ""
    signature = f"def {func.name}({params}){ret}"

    docstring = func.docstring or "No docstring provided."
    calls_str = ", ".join(func.calls[:15]) if func.calls else "none"
    class_ctx = (
        f"Class: {class_name}\nClass purpose: {class_docstring or 'N/A'}\n"
        if class_name else "Scope: top-level function\n"
    )
    source_block = func.source if func.source else "Source not available."

    text = (
        f"{class_ctx}"
        f"Function: {func.name}\n"
        f"Module: {file_info.module}\n"
        f"File: {file_info.relative}\n"
        f"Signature: {signature}\n"
        f"Docstring: {docstring}\n"
        f"Calls: {calls_str}\n"
        f"Source Code:\n{source_block}"
    )

    metadata = {
        "type":        "complete_function",
        "name":        func.name,
        "module":      file_info.module,
        "file":        file_info.relative,
        "class_name":  class_name or "",
        "return_type": func.return_type or "",
        "parameters":  json.dumps([p.name for p in func.parameters]),
        "calls":       json.dumps(func.calls[:15]),
        "is_method":   str(func.is_method),
        "lineno":      func.lineno,
    }

    return {"id": _make_id(text), "text": text, "metadata": metadata}


def build_complete_class_chunk(cls: ClassInfo, file_info: FileInfo) -> dict:
    """
    Build a complete class chunk including full source code.
    Used for class-level deep queries.
    """
    bases_str = ", ".join(cls.bases) if cls.bases else "object"
    docstring = cls.docstring or "No docstring provided."
    source_block = cls.source if cls.source else "Source not available."

    text = (
        f"Class: {cls.name}\n"
        f"Inherits: {bases_str}\n"
        f"Module: {file_info.module}\n"
        f"File: {file_info.relative}\n"
        f"Docstring: {docstring}\n"
        f"Source Code:\n{source_block}"
    )

    metadata = {
        "type":    "complete_class",
        "name":    cls.name,
        "module":  file_info.module,
        "file":    file_info.relative,
        "bases":   json.dumps(cls.bases),
        "methods": json.dumps([m.name for m in cls.methods]),
        "lineno":  cls.lineno,
    }

    return {"id": _make_id(text), "text": text, "metadata": metadata}


def build_file_chunk(file_info: FileInfo) -> dict:
    """
    Build a file-level chunk containing the entire source of a file.
    Used for file-wide queries.
    """
    try:
        source_block = Path(file_info.path).read_text(encoding="utf-8", errors="ignore")
    except Exception:
        source_block = "Source not available."

    docstring = file_info.docstring or "No module docstring."
    imports_str = ", ".join(file_info.imports) if file_info.imports else "none"

    text = (
        f"File: {file_info.relative}\n"
        f"Module: {file_info.module}\n"
        f"Docstring: {docstring}\n"
        f"Imports: {imports_str}\n"
        f"Source Code:\n{source_block}"
    )

    metadata = {
        "type":        "file",
        "name":        Path(file_info.relative).stem,
        "module":      file_info.module,
        "file":        file_info.relative,
        "class_name":  "",
        "return_type": "",
        "parameters":  "[]",
        "calls":       "[]",
        "is_method":   "False",
        "lineno":      0,
    }

    return {"id": _make_id(text), "text": text, "metadata": metadata}


# ── Embedding & Upserting ─────────────────────────────────────────────────────

def _upsert_batch(index, chunks: list[dict], model: SentenceTransformer, namespace: str) -> None:
    """Embed and upsert a list of chunks into a Pinecone namespace."""
    if not chunks:
        return

    texts     = [c["text"]     for c in chunks]
    ids       = [c["id"]       for c in chunks]
    metadatas = [c["metadata"] for c in chunks]

    embeddings = model.encode(texts, show_progress_bar=False).tolist()

    vectors = [
        {"id": vid, "values": vec, "metadata": {**meta, "text": txt}}
        for vid, vec, meta, txt in zip(ids, embeddings, metadatas, texts)
    ]

    index.upsert(vectors=vectors, namespace=namespace)


# ── Main Embed Pipeline ───────────────────────────────────────────────────────

def embed_codebase(root_path: str) -> None:
    """
    Full pipeline:
        1. Parse codebase via ast_parser
        2. Build hybrid chunks (class + function level)
        3. Embed with sentence-transformers
        4. Store in ChromaDB (two collections)

    Args:
        root_path: Absolute path to the monolithic codebase root.
    """
    console.rule("[bold cyan]Codebase Oracle — Embedder[/bold cyan]")

    # Step 1 — Parse
    console.print(f"\n[bold]📂 Root:[/bold] {root_path}\n")
    parsed_files: list[FileInfo] = parse_codebase(root_path)

    if not parsed_files:
        console.print("[yellow]⚠ No Python files parsed. Exiting.[/yellow]")
        return

    # Step 2 — Build chunks
    class_chunks:    list[dict] = []
    function_chunks: list[dict] = []

    for file_info in parsed_files:
        # Class-level chunks
        for cls in file_info.classes:
            class_chunks.append(build_class_chunk(cls, file_info))

            # Method-level chunks (carry class context)
            for method in cls.methods:
                function_chunks.append(build_function_chunk(
                    method, file_info,
                    class_name=cls.name,
                    class_docstring=cls.docstring,
                ))

        # Top-level function chunks
        for func in file_info.functions:
            function_chunks.append(build_function_chunk(func, file_info))

        # Module-level chunk for files with no classes and no functions
        if not file_info.classes and not file_info.functions:
            function_chunks.append(build_module_chunk(file_info))

    complete_chunks: list[dict] = []

    for file_info in parsed_files:
        complete_chunks.append(build_file_chunk(file_info))
        for cls in file_info.classes:
            complete_chunks.append(build_complete_class_chunk(cls, file_info))
            for method in cls.methods:
                complete_chunks.append(build_complete_function_chunk(
                    method, file_info,
                    class_name=cls.name,
                    class_docstring=cls.docstring,
                ))
        for func in file_info.functions:
            complete_chunks.append(build_complete_function_chunk(func, file_info))

    console.print(
        f"[green]✔[/green] Chunks built: "
        f"[magenta]{len(class_chunks)}[/magenta] class chunks · "
        f"[cyan]{len(function_chunks)}[/cyan] function chunks · "
        f"[yellow]{len(complete_chunks)}[/yellow] complete chunks\n"
    )

    # Step 3 — Load model
    model = load_embedding_model()

    # Step 4 — Pinecone
    index = get_pinecone_index()

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TextColumn("{task.completed}/{task.total}"),
        console=console,
    ) as progress:

        # Embed class chunks in batches of 32
        BATCH = 32
        task1 = progress.add_task(
            "[magenta]Embedding class chunks...", total=len(class_chunks)
        )
        for i in range(0, len(class_chunks), BATCH):
            batch = class_chunks[i:i + BATCH]
            _upsert_batch(index, batch, model, CLASS_COLLECTION)
            progress.advance(task1, len(batch))

        task2 = progress.add_task(
            "[cyan]Embedding function chunks...", total=len(function_chunks)
        )
        for i in range(0, len(function_chunks), BATCH):
            batch = function_chunks[i:i + BATCH]
            _upsert_batch(index, batch, model, FUNCTION_COLLECTION)
            progress.advance(task2, len(batch))

        task3 = progress.add_task(
            "[yellow]Embedding complete chunks...", total=len(complete_chunks)
        )
        for i in range(0, len(complete_chunks), BATCH):
            batch = complete_chunks[i:i + BATCH]
            _upsert_batch(index, batch, model, COMPLETE_COLLECTION)
            progress.advance(task3, len(batch))

    # Step 5 — Summary
    _render_embed_summary(root_path, class_chunks, function_chunks, complete_chunks)

def _render_embed_summary(root_path: str,
                              class_chunks: list[dict],
                              function_chunks: list[dict],
                              complete_chunks: list[dict]) -> None:
    """Render a rich summary panel after embedding."""
    table = Table(box=box.SIMPLE, show_header=False, padding=(0, 2))
    table.add_column(style="dim")
    table.add_column(style="bold white")

    table.add_row("Codebase",          root_path)
    table.add_row("Embedding model",   EMBEDDING_MODEL)
    table.add_row("Class chunks", str(len(class_chunks)))
    table.add_row("Function chunks", str(len(function_chunks)))
    table.add_row("Complete chunks", str(len(complete_chunks)))
    table.add_row("Total chunks", str(len(class_chunks) + len(function_chunks) + len(complete_chunks)))
    table.add_row("Collections", f"{CLASS_COLLECTION}, {FUNCTION_COLLECTION}, {COMPLETE_COLLECTION}")
    table.add_row("Status",            "[bold green]✔ Indexing complete[/bold green]")

    console.print(Panel(table, title="[bold cyan]Embedding Summary[/bold cyan]",
                        border_style="cyan"))
    console.print("\n[bold green]✔ Codebase indexed. Ready for queries.[/bold green]\n")


# ── Query Helper (for retriever.py later) ────────────────────────────────────

def query_chunks(query: str,
                 collection_name: str,
                 model: SentenceTransformer,
                 n_results: int = 5,
                 filters: dict | None = None) -> list[dict]:
    """
    Query a Pinecone namespace and return top-n matching chunks.

    Args:
        query:           Natural language query string.
        collection_name: Namespace — CLASS_COLLECTION, FUNCTION_COLLECTION, or COMPLETE_COLLECTION.
        model:           Loaded SentenceTransformer model.
        n_results:       Number of results to return.
        filters:         Optional Pinecone metadata filters.

    Returns:
        List of dicts with keys: text, metadata, distance.
    """
    index     = get_pinecone_index()
    embedding = model.encode([query]).tolist()[0]

    kwargs: dict = {
        "vector":          embedding,
        "top_k":           n_results,
        "namespace":       collection_name,
        "include_metadata": True,
    }
    if filters:
        kwargs["filter"] = filters

    results = index.query(**kwargs)

    output = []
    for match in results["matches"]:
        meta = dict(match["metadata"])
        text = meta.pop("text", "")
        output.append({
            "text":     text,
            "metadata": meta,
            "distance": 1 - match["score"],
        })

    return output


# ── Entry Point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import sys

    path = sys.argv[1] if len(sys.argv) > 1 else "."

    try:
        embed_codebase(path)
    except (FileNotFoundError, NotADirectoryError) as e:
        console.print(f"[red]❌ {e}[/red]")
        sys.exit(1)