api.py

"""
CodeLens – FastAPI REST service.

Fixes vs original:
- CORS restricted to configured origins (not wildcard *)
- /api/init no longer mutates global os.environ (race condition fix)
- /api/upload cleans up temp directory after indexing
- Config centralised in config.py
- Health check reports actual DB + API key status
"""
import os
import shutil
import tempfile
from pathlib import Path
from contextlib import asynccontextmanager

import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException, UploadFile, File, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, List

from config import DB_PATH, REPO_PATH, API_ALLOWED_ORIGINS
from tools import (
    codebase_search, read_file, get_directory_tree,
    grep_search, get_file_outline,
)
from create_db import main as create_database

load_dotenv()


# ── Lifespan ──────────────────────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    print("CodeLens API starting...")
    yield
    print("CodeLens API shutting down.")


# ── App ───────────────────────────────────────────────────────────────────
app = FastAPI(
    title="CodeLens API",
    description="AI-powered code intelligence — semantic search, grep, file ops.",
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=API_ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)


# ── Pydantic models ───────────────────────────────────────────────────────
class QueryRequest(BaseModel):
    query: str
    context_size: Optional[int] = 5

class SearchRequest(BaseModel):
    query: str
    path: Optional[str] = "."
    is_regex: Optional[bool] = False

class FileRequest(BaseModel):
    file_path: str

class TreeRequest(BaseModel):
    directory: Optional[str] = "."
    max_depth: Optional[int] = 2

class InitRequest(BaseModel):
    repo_path: str
    db_path: Optional[str] = None
    force: Optional[bool] = False

class SearchResult(BaseModel):
    source: str
    content: str

class QueryResponse(BaseModel):
    success: bool
    results: List[SearchResult]
    message: Optional[str] = None


# ── Endpoints ─────────────────────────────────────────────────────────────
@app.get("/")
async def root():
    return {
        "service": "CodeLens API",
        "version": "1.0.0",
        "docs": "/docs",
        "endpoints": ["/health", "/api/ask", "/api/search", "/api/tree",
                      "/api/outline", "/api/read", "/api/init", "/api/upload"],
    }


@app.get("/health")
async def health_check():
    """Returns database and API key status."""
    return {
        "status":               "ok",
        "database_initialized": os.path.exists(DB_PATH),
        "groq_api_configured":  bool(os.getenv("GROQ_API_KEY")),
        "db_path":              DB_PATH,
    }


@app.post("/api/ask", response_model=QueryResponse)
async def ask_question(request: QueryRequest):
    """
    Semantic search over the indexed codebase.

    Example: {"query": "How does authentication work?"}
    """
    result = codebase_search.invoke({"query": request.query})

    if result.startswith("Error"):
        raise HTTPException(status_code=503, detail=result)

    sections = result.split("--- Source:")
    results  = []
    for section in sections[1:]:
        if not section.strip():
            continue
        lines   = section.strip().split("\n", 1)
        source  = lines[0].strip().replace("---", "").strip()
        content = lines[1] if len(lines) > 1 else ""
        results.append(SearchResult(source=source, content=content))

    return QueryResponse(
        success=True,
        results=results[: request.context_size],
        message=f"Found {len(results)} matches",
    )


@app.post("/api/search")
async def search_code(request: SearchRequest):
    """
    Grep search — exact string or regex across files.

    Example: {"query": "def main", "path": "./src"}
    """
    result  = grep_search.invoke({"query": request.query, "path": request.path, "is_regex": request.is_regex})
    if "No matches found" in result:
        return {"success": True, "matches": [], "message": "No matches found"}

    matches = [l for l in result.split("\n") if l.strip() and not l.startswith("Matches:")]
    return {"success": True, "matches": matches, "count": len(matches)}


@app.post("/api/tree")
async def get_tree(request: TreeRequest):
    """Get directory tree. Example: {"directory": "./src", "max_depth": 3}"""
    result = get_directory_tree.invoke({"directory": request.directory, "max_depth": request.max_depth})
    return {"success": True, "tree": result}


@app.post("/api/outline")
async def get_outline(request: FileRequest):
    """Get class/function outline of a file."""
    if not os.path.exists(request.file_path):
        raise HTTPException(status_code=404, detail="File not found")

    result = get_file_outline.invoke({"file_path": request.file_path})
    items  = []
    for line in result.split("\n"):
        if line.strip() and line.startswith("L"):
            parts = line.split(": ", 1)
            if len(parts) == 2:
                items.append({"line": parts[0], "definition": parts[1]})

    return {"success": True, "file": request.file_path, "outline": items}


@app.post("/api/read")
async def read_file_content(request: FileRequest):
    """Read a file's content."""
    if not os.path.exists(request.file_path):
        raise HTTPException(status_code=404, detail="File not found")

    content = read_file.invoke({"file_path": request.file_path})
    return {"success": True, "file": request.file_path, "content": content}


@app.post("/api/init")
async def initialize_database(request: InitRequest, background_tasks: BackgroundTasks):
    """
    Index a repository into the vector database (runs in background).

    Example: {"repo_path": "/path/to/repo"}
    """
    if not os.path.exists(request.repo_path):
        raise HTTPException(status_code=404, detail="Repository path not found")

    db_path = request.db_path or DB_PATH

    background_tasks.add_task(
        create_database,
        repo_path=request.repo_path,
        db_path=db_path,
        force=request.force or False,
    )

    return {
        "success": True,
        "message": "Indexing started in background. Poll /health to check db status.",
        "repo_path": request.repo_path,
        "db_path": db_path,
    }


@app.post("/api/upload")
async def upload_repository(file: UploadFile = File(...)):
    """
    Upload a zip of a repository, index it, and return results.
    Temp files are cleaned up automatically.
    """
    temp_dir = tempfile.mkdtemp()
    try:
        zip_path    = os.path.join(temp_dir, file.filename or "repo.zip")
        extract_dir = os.path.join(temp_dir, "repo")

        with open(zip_path, "wb") as buf:
            shutil.copyfileobj(file.file, buf)

        shutil.unpack_archive(zip_path, extract_dir)
        create_database(repo_path=extract_dir, db_path=DB_PATH, force=True)

        return {
            "success": True,
            "message": "Repository uploaded and indexed successfully.",
        }
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))
    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)


if __name__ == "__main__":
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=False)