api.py

import os
import re
from datetime import datetime, timezone
from pathlib import Path
from threading import Lock

import fitz
import uvicorn
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from pdf_to_markdown import convert_pdf_to_markdown
from src import mongo_store
from src.ingestion import IngestionStats, ingest_sources, remove_document

app = FastAPI(title="Agentic RAG API")

_origins = os.environ.get(
    "CORS_ORIGINS", "http://localhost:3000,http://127.0.0.1:3000"
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=[origin.strip() for origin in _origins.split(",") if origin.strip()]
    or ["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

DATA_DIR = Path(__file__).resolve().parent.parent / "data"
PDF_DIR = DATA_DIR / "pdf"
MD_DIR = DATA_DIR / "md"
PDF_DIR.mkdir(parents=True, exist_ok=True)
MD_DIR.mkdir(parents=True, exist_ok=True)

MAX_BYTES = 25 * 1024 * 1024
_INVALID_NAME = re.compile(r'[<>:"/\\|?*\x00-\x1f]')
_INGESTION_LOCK = Lock()


class ChatRequest(BaseModel):
    user_id: str
    query: str
    session_id: str | None = None


class ChatResponse(BaseModel):
    user_id: str
    query: str
    answer: str


class DocumentItem(BaseModel):
    id: str
    filename: str
    status: str
    pages: int
    chars: int
    chunks: int
    sizeBytes: int
    hasMarkdown: bool
    vectorized: bool
    createdAt: str


class IngestResponse(BaseModel):
    document: DocumentItem
    parentsRegistered: int
    chunksUpserted: int
    collectionCount: int
    warnings: list[str] = Field(default_factory=list)


class DocumentsResponse(BaseModel):
    documents: list[DocumentItem]


class DeleteResponse(BaseModel):
    id: str
    deleted: bool


def _safe_pdf_name(filename: str) -> str:
    base = Path(filename or "").name.strip()
    if not base.lower().endswith(".pdf"):
        raise HTTPException(status_code=415, detail="Only PDF files are accepted.")
    stem = _INVALID_NAME.sub("_", base[:-4]).strip()
    if not stem:
        stem = "document"
    return f"{stem}.pdf"


def _pdf_page_count(pdf_path: Path) -> int:
    try:
        with fitz.open(str(pdf_path)) as doc:
            return doc.page_count
    except Exception:
        return 0


def _doc_from_pdf_with_stats(pdf_path: Path, stats: IngestionStats) -> DocumentItem:
    md_path = MD_DIR / f"{pdf_path.stem}.md"
    has_md = md_path.exists()
    chars = 0
    if has_md:
        try:
            chars = len(md_path.read_text(encoding="utf-8", errors="ignore"))
        except OSError:
            pass

    doc_stats = next((d for d in stats.documents if d.doc_id == pdf_path.stem), None)
    chunks = doc_stats.chunks if doc_stats else 0
    vectorized = has_md and chunks > 0 and stats.upserted > 0

    return DocumentItem(
        id=pdf_path.stem,
        filename=pdf_path.name,
        status="ready" if vectorized else "error",
        pages=_pdf_page_count(pdf_path),
        chars=chars,
        chunks=chunks,
        sizeBytes=pdf_path.stat().st_size,
        hasMarkdown=has_md,
        vectorized=vectorized,
        createdAt=datetime.fromtimestamp(pdf_path.stat().st_mtime, tz=timezone.utc).isoformat(),
    )


def _normalize_created(value) -> str:
    if isinstance(value, datetime):
        dt = value if value.tzinfo else value.replace(tzinfo=timezone.utc)
        return dt.astimezone(timezone.utc).isoformat()
    if not value:
        return datetime.now(timezone.utc).isoformat()
    text = str(value)
    for fmt in (
        "%Y-%m-%dT%H:%M:%S.%f",
        "%Y-%m-%dT%H:%M:%S",
        "%Y-%m-%d %H:%M:%S.%f",
        "%Y-%m-%d %H:%M:%S",
    ):
        try:
            return datetime.strptime(text, fmt).replace(tzinfo=timezone.utc).isoformat()
        except ValueError:
            continue
    return text


def _doc_item_from_mongo(rec: dict) -> DocumentItem:
    doc_id = str(rec.get("doc_id") or "")
    source_file = str(rec.get("source_file") or (f"{doc_id}.pdf" if doc_id else ""))
    chunks = int(rec.get("chunk_count") or rec.get("chunks") or 0)

    pdf_path = PDF_DIR / source_file if source_file else None
    md_path = MD_DIR / f"{doc_id}.md" if doc_id else None
    has_md = bool(rec.get("has_markdown")) or bool(md_path and md_path.exists())

    pages = int(rec.get("pages") or 0)
    if not pages and pdf_path and pdf_path.exists():
        pages = _pdf_page_count(pdf_path)

    chars = int(rec.get("chars") or 0)
    if not chars and md_path and md_path.exists():
        try:
            chars = len(md_path.read_text(encoding="utf-8", errors="ignore"))
        except OSError:
            pass

    size_bytes = int(rec.get("size_bytes") or 0)
    if not size_bytes and pdf_path and pdf_path.exists():
        try:
            size_bytes = pdf_path.stat().st_size
        except OSError:
            pass

    vectorized = bool(rec["vectorized"]) if "vectorized" in rec else chunks > 0
    status = str(rec.get("status") or ("ready" if vectorized and chunks > 0 else "error"))

    return DocumentItem(
        id=doc_id or source_file,
        filename=source_file or f"{doc_id}.pdf",
        status=status,
        pages=pages,
        chars=chars,
        chunks=chunks,
        sizeBytes=size_bytes,
        hasMarkdown=has_md,
        vectorized=vectorized,
        createdAt=_normalize_created(rec.get("created_at")),
    )


def _reset_search_cache() -> None:
    try:
        from src.tools.vector_search import reset_vector_search_cache
        reset_vector_search_cache()
    except Exception:
        pass


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------

@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/chat", response_model=ChatResponse)
def chat(request: ChatRequest):
    from src.core import final_answer
    answer = final_answer(
        user_id=request.user_id,
        query=request.query,
        session_id=request.session_id,
    )
    return ChatResponse(user_id=request.user_id, query=request.query, answer=answer)


@app.post("/ingest", response_model=IngestResponse, status_code=201)
def ingest(file: UploadFile = File(...)):
    name = _safe_pdf_name(file.filename or "")
    data = file.file.read()
    if len(data) == 0:
        raise HTTPException(status_code=400, detail="Empty file.")
    if len(data) > MAX_BYTES:
        raise HTTPException(status_code=413, detail="File exceeds the 25 MB limit.")

    with _INGESTION_LOCK:
        pdf_path = PDF_DIR / name
        pdf_path.write_bytes(data)

        md_path = MD_DIR / f"{pdf_path.stem}.md"
        conv = convert_pdf_to_markdown(str(pdf_path), str(md_path), overwrite=True)
        if not conv.success:
            raise HTTPException(
                status_code=500,
                detail=conv.error or "PDF to Markdown conversion failed.",
            )

        try:
            stats = ingest_sources(md_path, incremental=True)
            _reset_search_cache()
        except Exception as exc:
            raise HTTPException(
                status_code=500,
                detail=f"Markdown was saved, but chunk/embed/Qdrant failed: {exc}",
            ) from exc

        document = _doc_from_pdf_with_stats(pdf_path, stats)
        doc_stats = next((d for d in stats.documents if d.doc_id == pdf_path.stem), None)
        warnings = list(doc_stats.warnings) if doc_stats else []

        mongo_record = {
            "doc_id": pdf_path.stem,
            "filename": pdf_path.name,
            "source_file": pdf_path.name,
            "status": document.status,
            "pages": document.pages,
            "chars": document.chars,
            "chunk_count": document.chunks,
            "size_bytes": document.sizeBytes,
            "has_markdown": document.hasMarkdown,
            "vectorized": document.vectorized,
        }

    try:
        mongo_store.upsert_document(mongo_record)
    except mongo_store.MongoUnavailableError as exc:
        raise HTTPException(
            status_code=502,
            detail=f"Đã lưu file và cập nhật Qdrant, nhưng ghi MongoDB thất bại: {exc}",
        ) from exc

    return IngestResponse(
        document=document,
        parentsRegistered=stats.parent_documents,
        chunksUpserted=stats.upserted,
        collectionCount=stats.collection_count,
        warnings=warnings,
    )


@app.get("/documents", response_model=DocumentsResponse)
def list_documents():
    try:
        records = mongo_store.list_documents()
    except mongo_store.MongoUnavailableError as exc:
        raise HTTPException(
            status_code=503, detail=f"Không kết nối được MongoDB: {exc}"
        ) from exc
    return DocumentsResponse(documents=[_doc_item_from_mongo(rec) for rec in records])


@app.delete("/documents/{doc_id}", response_model=DeleteResponse)
def delete_document(doc_id: str):
    try:
        mongo_rec = mongo_store.get_document(doc_id)
    except mongo_store.MongoUnavailableError as exc:
        raise HTTPException(
            status_code=503, detail=f"Không kết nối được MongoDB: {exc}"
        ) from exc

    if mongo_rec and mongo_rec.get("source_file"):
        file_name = Path(str(mongo_rec["source_file"])).name
    else:
        file_name = Path(doc_id).name
    stem = file_name[:-4] if file_name.lower().endswith(".pdf") else file_name
    pdf_path = PDF_DIR / f"{stem}.pdf"
    md_path = MD_DIR / f"{stem}.md"

    file_existed = pdf_path.exists() or md_path.exists()
    if mongo_rec is None and not file_existed:
        raise HTTPException(status_code=404, detail="Document not found.")

    with _INGESTION_LOCK:
        pdf_path.unlink(missing_ok=True)
        md_path.unlink(missing_ok=True)
        try:
            remove_document(doc_id)  # deletes Qdrant points + MongoDB doc + chunks
            _reset_search_cache()
        except Exception as exc:
            raise HTTPException(
                status_code=500,
                detail=f"Xóa file OK, nhưng xóa Qdrant/MongoDB thất bại: {exc}",
            ) from exc

    return DeleteResponse(id=doc_id, deleted=True)


if __name__ == "__main__":
    uvicorn.run("src.api:app", host="127.0.0.1", port=8080, reload=True)