main.py

"""
main.py
-------
FastAPI backend for the Codebase Oracle system.
This is the HTTP layer — thin wrapper around inference.py.

Endpoints:
    POST /index          — ingest + embed a codebase from given path
    POST /query          — run a query (macro / micro / cross_module)
    GET  /status         — check if a codebase is indexed and ready
    GET  /tree           — return parsed codebase tree for UI sidebar
    GET  /health         — simple health check

Run:
    uvicorn main:app --reload --port 8000

Depends on:
    - inference.py
    - embedder.py
    - call_graph.py
    - ast_parser.py
    - vector_store.py
    - fastapi, uvicorn, pydantic, python-dotenv
"""

import os
from contextlib import asynccontextmanager
from dotenv import load_dotenv

import tempfile
import zipfile
import shutil
from fastapi import FastAPI, HTTPException, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field
from rich.console import Console

from inference.inference import get_engine, InferenceRequest
from ingest.embed import embed_codebase
from store.call_graph import build_and_save, get_call_graph, CALL_GRAPH_PATH
from ingest.parse_ast import parse_codebase
from store.vector_store import get_vector_store

load_dotenv()

console = Console()

# ── App Lifespan ──────────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize shared resources on startup."""
    console.rule("[bold cyan]Codebase Oracle — Starting[/bold cyan]")
    # Pre-warm the inference engine (loads embedding model once)
    get_engine()
    console.print("[green]✔[/green] Server ready.\n")
    yield
    console.print("[dim]Server shutting down.[/dim]")


# ── App ───────────────────────────────────────────────────────────────────────

app = FastAPI(
    title="Codebase Oracle",
    description="AI-powered monolithic codebase comprehension system.",
    version="1.0.0",
    lifespan=lifespan,
)

# Allow UI (served from same origin or localhost dev)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Serve UI static files
UI_DIR = os.path.join(os.path.dirname(__file__), "ui")
if os.path.exists(UI_DIR):
    app.mount("/ui", StaticFiles(directory=UI_DIR), name="ui")
    app.mount("/static", StaticFiles(directory=os.path.join(UI_DIR, "static")), name="static")


# ── Pydantic Request / Response Models ────────────────────────────────────────

class IndexRequest(BaseModel):
    """Request body for POST /index"""
    root_path: str = Field(
        ...,
        description="Absolute path to the monolithic codebase root directory.",
        example="/home/user/projects/my-django-app"
    )


class QueryRequest(BaseModel):
    """Request body for POST /query"""
    query_type: str = Field(
        ...,
        description="One of: 'macro', 'micro', 'cross_module'",
        example="micro"
    )
    query: str = Field(
        ...,
        description="Natural language developer query.",
        example="What does process_payment do and how do I use it?"
    )
    subtype: str = Field(
        default="",
        description="Macro subtype: 'overall_architecture' | 'module_responsibility' | 'data_flow'",
        example="overall_architecture"
    )
    function_name: str = Field(
        default="",
        description="Target function/method name for micro and cross_module queries.",
        example="process_payment"
    )
    class_name: str = Field(
        default="",
        description="Target class name if function is a method.",
        example="PaymentProcessor"
    )
    module_name: str = Field(
        default="",
        description="Target module name for macro module_responsibility queries.",
        example="payments"
    )
    followup: bool = Field(
        default=False,
        description="True if this is a follow-up to a previous response."
    )
    previous_response: str = Field(
        default="",
        description="Previous LLM response for follow-up context."
    )


class IndexResponse(BaseModel):
    success:       bool
    message:       str
    class_chunks:  int = 0
    function_chunks: int = 0
    total_chunks:  int = 0
    graph_nodes:   int = 0
    graph_edges:   int = 0


class QueryResponse(BaseModel):
    success:  bool
    content:  str
    error:    str = ""
    metadata: dict = {}


class StatusResponse(BaseModel):
    indexed:          bool
    class_chunks:     int
    function_chunks:  int
    total_chunks:     int
    graph_loaded:     bool
    graph_nodes:      int


class TreeNode(BaseModel):
    name:     str
    type:     str           # "module" | "file" | "class" | "function"
    children: list["TreeNode"] = []

TreeNode.model_rebuild()


class TreeResponse(BaseModel):
    success: bool
    tree:    list[TreeNode] = []
    error:   str = ""


# ── Endpoints ─────────────────────────────────────────────────────────────────

@app.post("/upload-index", response_model=IndexResponse)
async def upload_index(file: UploadFile = File(...)):
    """
    Accept a ZIP file, extract it to a temp directory, and index it.
    Allows deployment without requiring local filesystem access.
    """
    if not file.filename.endswith(".zip"):
        raise HTTPException(status_code=400, detail="Only .zip files are accepted.")

    tmp_dir = tempfile.mkdtemp()

    try:
        zip_path = os.path.join(tmp_dir, file.filename)
        with open(zip_path, "wb") as f:
            shutil.copyfileobj(file.file, f)

        with zipfile.ZipFile(zip_path, "r") as zf:
            zf.extractall(tmp_dir)

        os.remove(zip_path)

        # Find the extracted root — skip __MACOSX and similar artifacts
        candidates = [
            os.path.join(tmp_dir, d)
            for d in os.listdir(tmp_dir)
            if os.path.isdir(os.path.join(tmp_dir, d)) and not d.startswith("__")
        ]
        root = candidates[0] if candidates else tmp_dir

        console.rule(f"[bold cyan]Indexing ZIP: {file.filename}[/bold cyan]")

        embed_codebase(root)

        graph = build_and_save(root)
        graph_stats = graph.stats()

        store = get_vector_store()
        vstats = store.stats()

        console.print("[bold green]✔ ZIP Indexing complete.[/bold green]\n")

        return IndexResponse(
            success=True,
            message=f"ZIP indexed successfully: {file.filename}",
            class_chunks=vstats["class_chunks"],
            function_chunks=vstats["function_chunks"],
            total_chunks=vstats["total"],
            graph_nodes=graph_stats["total_nodes"],
            graph_edges=graph_stats["total_edges"],
        )

    except zipfile.BadZipFile:
        raise HTTPException(status_code=400, detail="Invalid or corrupted ZIP file.")

    except Exception as e:
        console.print(f"[red]❌ ZIP indexing failed: {e}[/red]")
        raise HTTPException(status_code=500, detail=f"ZIP indexing failed: {str(e)}")

    finally:
        shutil.rmtree(tmp_dir, ignore_errors=True)

@app.get("/health")
async def health():
    """Simple health check."""
    return {"status": "ok", "service": "Codebase Oracle"}


@app.get("/", response_class=FileResponse)
async def serve_ui():
    """Serve the UI index.html at root."""
    ui_path = os.path.join(UI_DIR, "index.html")
    if not os.path.exists(ui_path):
        raise HTTPException(
            status_code=404,
            detail="UI not found. Place index.html in the ui/ directory."
        )
    return FileResponse(ui_path)


@app.post("/index", response_model=IndexResponse)
async def index_codebase(req: IndexRequest):
    """
    Ingest, parse, embed, and index a monolithic codebase.
    Builds both ChromaDB vector index and call_graph.json.

    This is the first endpoint to call before any queries.
    """
    root = req.root_path.strip()

    if not os.path.exists(root):
        raise HTTPException(
            status_code=400,
            detail=f"Path does not exist: {root}"
        )

    if not os.path.isdir(root):
        raise HTTPException(
            status_code=400,
            detail=f"Path is not a directory: {root}"
        )

    try:
        console.rule(f"[bold cyan]Indexing: {root}[/bold cyan]")

        # Step 1 — Embed codebase into ChromaDB
        embed_codebase(root)

        # Step 2 — Build and save call graph
        graph = build_and_save(root)
        graph_stats = graph.stats()

        # Step 3 — Fetch vector store stats
        store = get_vector_store()
        vstats = store.stats()

        console.print("[bold green]✔ Indexing complete.[/bold green]\n")

        return IndexResponse(
            success=True,
            message=f"Codebase indexed successfully: {root}",
            class_chunks=vstats["class_chunks"],
            function_chunks=vstats["function_chunks"],
            total_chunks=vstats["total"],
            graph_nodes=graph_stats["total_nodes"],
            graph_edges=graph_stats["total_edges"],
        )

    except Exception as e:
        console.print(f"[red]❌ Indexing failed: {e}[/red]")
        raise HTTPException(status_code=500, detail=f"Indexing failed: {str(e)}")


@app.post("/query", response_model=QueryResponse)
async def query(req: QueryRequest):
    """
    Run a macro, micro, or cross-module query against the indexed codebase.
    Returns a markdown-formatted response string.
    """
    store = get_vector_store()
    if not store.is_indexed():
        raise HTTPException(
            status_code=400,
            detail="Codebase is not indexed yet. Call POST /index first."
        )

    engine = get_engine()

    inference_req = InferenceRequest(
        query_type=req.query_type,
        query=req.query,
        subtype=req.subtype,
        function_name=req.function_name,
        class_name=req.class_name,
        module_name=req.module_name,
        followup=req.followup,
        previous_response=req.previous_response,
    )

    resp = engine.infer(inference_req)

    return QueryResponse(
        success=resp.success,
        content=resp.content,
        error=resp.error,
        metadata=resp.metadata,
    )


@app.get("/status", response_model=StatusResponse)
async def status():
    """
    Check whether the codebase is indexed and the system is ready for queries.
    """
    store  = get_vector_store()
    vstats = store.stats()

    graph_loaded = False
    graph_nodes  = 0

    if os.path.exists(CALL_GRAPH_PATH):
        try:
            graph = get_call_graph()
            graph_loaded = graph.is_loaded()
            graph_nodes  = graph.stats()["total_nodes"]
        except Exception:
            pass

    return StatusResponse(
        indexed=store.is_indexed(),
        class_chunks=vstats["class_chunks"],
        function_chunks=vstats["function_chunks"],
        total_chunks=vstats["total"],
        graph_loaded=graph_loaded,
        graph_nodes=graph_nodes,
    )


@app.get("/tree", response_model=TreeResponse)
async def get_tree():
    """
    Return the parsed codebase structure as a nested tree.
    Used by the UI sidebar to render the codebase explorer.
    """
    store = get_vector_store()
    if not store.is_indexed():
        return TreeResponse(
            success=False,
            error="Codebase not indexed yet. Call POST /index first."
        )

    try:
        # Fetch both class and function chunks to reconstruct tree
        class_results = store.get_all("class_chunks", limit=500)
        func_results = store.get_all("function_chunks", limit=500)

        # Group by module → file → classes/functions
        modules: dict[str, dict[str, dict[str, set]]] = {}

        # --- classes ---
        for chunk in class_results:
            mod = chunk.module
            file = chunk.file
            modules.setdefault(mod, {}).setdefault(file, {"classes": set(), "functions": set()})
            modules[mod][file]["classes"].add(chunk.name)

        # --- functions (top-level only) ---
        for chunk in func_results:
            if not chunk.class_name:
                mod = chunk.module
                file = chunk.file
                modules.setdefault(mod, {}).setdefault(file, {"classes": set(), "functions": set()})
                modules[mod][file]["functions"].add(chunk.name)

        # Also fetch function chunks for top-level functions
        func_results = store.get_all("function_chunks", limit=500)
        func_by_file: dict[str, list[str]] = {}
        for chunk in func_results:
            if not chunk.class_name:  # top-level only
                func_by_file.setdefault(chunk.file, []).append(chunk.name)

                # Build tree structure

                all_files = set()
                for files in modules.values():
                    for file_path in files:
                        all_files.add(file_path)

                # Derive root directory name from common first path component
                first_parts = [f.split("/")[0] for f in all_files if "/" in f]
                root_name = first_parts[0] if first_parts else "codebase"

                root_node = TreeNode(name=root_name, type="module")

                for module_name, files in sorted(modules.items()):
                    module_node = TreeNode(name=module_name, type="module")

                    for file_path, content in sorted(files.items()):
                        file_node = TreeNode(
                            name=os.path.basename(file_path),
                            type="file"
                        )

                        for cls_name in sorted(content["classes"]):
                            file_node.children.append(
                                TreeNode(name=cls_name, type="class")
                            )

                        for fn_name in sorted(content["functions"]):
                            file_node.children.append(
                                TreeNode(name=fn_name, type="function")
                            )

                        module_node.children.append(file_node)

                    root_node.children.append(module_node)

                return TreeResponse(success=True, tree=[root_node])

    except Exception as e:
        console.print(f"[red]❌ Tree build failed: {e}[/red]")
        return TreeResponse(success=False, error=str(e))


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

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "main:app",
        port=8000,
        reload=True,
    )