app/app.py

import json
import re
from pathlib import Path

import faiss
import torch
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from peft import PeftModel

from app.indexing import (
    append_chunks_to_index,
    build_chunks_from_document,
    existing_sources,
    list_indexed_documents,
    save_raw_document,
)
from app.paths import index_path, lora_path, metadata_path


# -----------------------------
# API Initialization
# -----------------------------

app = FastAPI(
    title="Enterprise Code Comment Generator",
    description="CodeT5 + LoRA + RAG for ERPNext Accounting",
    version="1.0",
)


# -----------------------------
# Request Schema
# -----------------------------


class CodeRequest(BaseModel):
    code: str


def _safe_filename(filename: str) -> str:
    name = Path(filename).name
    if not name or name.startswith("."):
        raise HTTPException(status_code=400, detail="Invalid filename")
    if not name.lower().endswith(".md"):
        raise HTTPException(status_code=400, detail="Only Markdown (.md) files are allowed")
    return name


def _safe_category(category: str) -> str:
    if not category or not re.fullmatch(r"[a-zA-Z0-9_-]+", category):
        raise HTTPException(status_code=400, detail="Invalid category")
    return category


# -----------------------------
# Global Variables for Models
# -----------------------------

tokenizer = None
model = None
embedder = None
index = None
metadata = None

BASE_MODEL = "Salesforce/codet5-small"


# -----------------------------
# Load Models (Startup Event)
# -----------------------------


@app.on_event("startup")
async def load_models() -> None:
    """
    Lazily load all heavy models and indexes when the app starts.
    Mirrors the setup logic from the RAG notebook.
    """
    global tokenizer, model, embedder, index, metadata

    print("Loading CodeT5 model...")

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    base_model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)

    model = PeftModel.from_pretrained(
        base_model,
        str(lora_path()),
        is_trainable=False,
    )
    model.eval()

    print("CodeT5 loaded")

    print("Loading SentenceBERT...")
    embedder = SentenceTransformer("all-MiniLM-L6-v2")
    print("SentenceBERT loaded")

    print("Loading FAISS index and metadata...")
    print(f"DATA_DIR={index_path().parent}")

    index = faiss.read_index(str(index_path()))
    with open(metadata_path(), "r", encoding="utf-8") as f:
        metadata = json.load(f)

    print("FAISS index loaded")


# -----------------------------
# Comment Generation
# -----------------------------


def generate_comment(code: str) -> str:
    """
    Generate a natural-language summary/comment for the given code.
    """
    prompt = "summarize the purpose of this function: " + code

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
    )

    with torch.no_grad():
        outputs = model.generate(
            **inputs.to(model.device),
            max_new_tokens=80,
            min_length=30,
            num_beams=6,
            repetition_penalty=1.25,
            length_penalty=1.35,
            no_repeat_ngram_size=3,
            early_stopping=True,
        )

    comment = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return comment.strip()


# -----------------------------
# Context Retrieval (RAG-style)
# -----------------------------


def retrieve_context(code: str, top_k: int = 5) -> str:
    """
    Retrieve the most relevant documentation chunk for the given code.

    Mirrors the logic used in the RAG notebook:
    - Extract function name from the code
    - Use the function name as the semantic query
    - Re-rank FAISS candidates using simple keyword overlap
    """
    # Extract function name
    fn_match = re.search(r"def\s+([a-zA-Z0-9_]+)", code)
    fn_name = fn_match.group(1).lower() if fn_match else ""
    keywords = fn_name.split("_") if fn_name else []

    # Encode query using the function name (more focused than raw code)
    query_embedding = embedder.encode(
        [fn_name or code],
        convert_to_numpy=True,
    )

    # Retrieve FAISS candidates
    _, I = index.search(query_embedding, k=top_k)
    candidates = [metadata[i] for i in I[0]]

    # Hybrid re-ranking: boost chunks that contain function-name keywords
    scored = []
    for chunk in candidates:
        text = chunk["text"].lower()
        keyword_score = sum(1 for k in keywords if k and k in text)
        scored.append((keyword_score, chunk))

    scored.sort(key=lambda x: x[0], reverse=True)

    best_chunk = scored[0][1] if scored else candidates[0]
    return best_chunk["text"]


# -----------------------------
# Context Summarization
# -----------------------------


def summarize_context(context: str) -> str:
    """
    Clean and lightly summarize the retrieved context, similar to the notebook:
    - Strip markdown artifacts (#, \-, backslashes)
    - Collapse newlines
    - Keep only the first couple of meaningful sentences
    """
    clean = context.replace("#", "")
    clean = clean.replace("\\-", "")
    clean = clean.replace("\\", "")
    clean = re.sub(r"\n+", " ", clean)

    sentences = re.split(r'(?<=[.!?])\s+', clean.strip())
    meaningful = [s.strip() for s in sentences if len(s.strip()) > 25]

    return " ".join(meaningful[:2]) if meaningful else clean.strip()


# -----------------------------
# API Endpoint
# -----------------------------


@app.post("/generate_comment")
def generate_comment_api(request: CodeRequest):
    code = request.code

    comment = generate_comment(code)
    raw_context = retrieve_context(code)
    context = summarize_context(raw_context)

    return {
        "comment": comment,
        "context": context,
    }


@app.get("/documents")
def list_documents():
    """
    List all documents currently indexed in the FAISS metadata.
    """
    return list_indexed_documents(metadata, index.ntotal)


@app.post("/documents")
async def add_document(
    file: UploadFile = File(..., description="Markdown file (.md only)"),
    category: str = Form(
        default="accounting",
        description="Subfolder under data/enterprise_accounting_docs/",
    ),
):
    """
    Upload a Markdown file and append its chunks to the FAISS index and metadata.

    Existing documents and vectors are preserved; duplicate sources are rejected.
    """
    if not file.filename:
        raise HTTPException(status_code=400, detail="Missing filename")

    filename = _safe_filename(file.filename)
    category = _safe_category(category)
    source = f"{category}/{filename}"

    raw = await file.read()
    try:
        content = raw.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise HTTPException(
            status_code=400,
            detail="File must be valid UTF-8 text",
        ) from exc

    if not content.strip():
        raise HTTPException(status_code=400, detail="Uploaded file is empty")

    if source in existing_sources(metadata):
        raise HTTPException(
            status_code=409,
            detail=f"Document already indexed: {source}",
        )

    new_chunks = build_chunks_from_document(source, content)
    if not new_chunks:
        raise HTTPException(
            status_code=400,
            detail="No indexable chunks produced (document empty or too short)",
        )

    save_raw_document(category, filename, content)
    vectors_added = append_chunks_to_index(
        embedder,
        index,
        metadata,
        new_chunks,
        persist=True,
    )

    return {
        "message": "Document indexed successfully",
        "source": source,
        "chunks_added": len(new_chunks),
        "vectors_added": vectors_added,
        "total_chunks": len(metadata),
        "total_vectors": index.ntotal,
    }