src/chatbot/document_hub.py

import os
import re
import hashlib
import threading
from pathlib import Path
from typing import Optional

# PDF Extraction
try:
    from pypdf import PdfReader  # preferred
except ImportError:
    from PyPDF2 import PdfReader  # fallback


# DOCUMENT STORAGE 
# Stores extracted documents: { doc_id: { "filename", "chunks", "metadata" } }
documents_db = {}
doc_lock = threading.Lock()

# Chunking configuration
CHUNK_SIZE = 500        # characters per chunk
CHUNK_OVERLAP = 100     # overlap between consecutive chunks to preserve context



# Internal helpers
def _generate_doc_id(filename: str, content: str) -> str:
    """Create a stable, unique ID from filename + content hash."""
    digest = hashlib.md5((filename + content[:200]).encode()).hexdigest()[:10]
    return f"{Path(filename).stem}_{digest}"


def _chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> list[dict]:
    """
    Split *text* into overlapping fixed-size chunks, preserving the source
    page number for each chunk.

    Each chunk is a dict:  { "text": str, "page": int }

    Strategy
    --------
    1. Split the raw text on [Page N] markers produced by extract_text_from_pdf.
    2. For each page block, slide a fixed-size window (with overlap) to
       produce sub-chunks, all tagged with the page number they came from.
    """
    chunks = []

    # re.split with a capturing group keeps the delimiters in the result list:
    # ['', '1', 'page1 text...', '2', 'page2 text...', ...]
    page_blocks = re.split(r'\[Page (\d+)\]', text)

    for i in range(1, len(page_blocks) - 1, 2):
        page_num = int(page_blocks[i])
        page_text = page_blocks[i + 1].strip()

        start = 0
        while start < len(page_text):
            end = start + chunk_size
            chunk_text = page_text[start:end].strip()
            if chunk_text:
                chunks.append({"text": chunk_text, "page": page_num})
            start += chunk_size - overlap

    return chunks



# Public API
def extract_text_from_pdf(file_path: str) -> str:
    """
    Extract all text from a PDF file using pypdf (or PyPDF2 as fallback).

    Each page is prefixed with a [Page N] marker so that downstream
    chunking can recover the source page number.

    Returns the concatenated text of all pages, or an error string if
    extraction fails.
    """
    path = Path(file_path)
    if not path.exists():
        return f"Error: File not found — {file_path}"
    if path.suffix.lower() != ".pdf":
        return f"Error: Expected a .pdf file, got '{path.suffix}'."

    try:
        reader = PdfReader(str(path))
        pages_text = []
        for i, page in enumerate(reader.pages):
            page_text = page.extract_text() or ""
            if page_text.strip():
                # [Page N] marker is parsed by _chunk_text — keep format stable
                pages_text.append(f"[Page {i + 1}]\n{page_text.strip()}")

        if not pages_text:
            return "Warning: No extractable text found. The PDF may be scanned/image-based."

        return "\n\n".join(pages_text)

    except Exception as e:
        return f"Error extracting PDF: {str(e)}"


def process_document(file_path: str, filename: Optional[str] = None) -> dict:
    """
    Full ingestion pipeline: extract → chunk (with page metadata) → store.

    Parameters
    ----------
    file_path : str
        Absolute or relative path to the PDF file.
    filename : str, optional
        Display name for the document. Defaults to the file's basename.

    Returns
    -------
    dict with keys: status, doc_id, filename, num_chunks
    """
    filename = filename or Path(file_path).name

    # 1. Extract raw text (with [Page N] markers)
    raw_text = extract_text_from_pdf(file_path)
    if raw_text.startswith("Error"):
        return {"status": "error", "message": raw_text}

    # 2. Chunk into page-aware dicts
    chunks = _chunk_text(raw_text)

    # 3. Stable document ID
    doc_id = _generate_doc_id(filename, raw_text)

    # 4. Thread-safe storage
    with doc_lock:
        documents_db[doc_id] = {
            "filename": filename,
            "raw_text": raw_text,
            "chunks": chunks,          # list[dict]  {"text", "page"}
            "metadata": {
                "file_path": file_path,
                "num_pages": raw_text.count("[Page "),
                "num_chunks": len(chunks),
                "chunk_size": CHUNK_SIZE,
                "chunk_overlap": CHUNK_OVERLAP,
            },
        }

    print(f"[DocumentHub] Processed '{filename}' → {len(chunks)} chunks (doc_id={doc_id})")
    return {
        "status": "success",
        "doc_id": doc_id,
        "filename": filename,
        "num_chunks": len(chunks),
    }


def get_document(doc_id: str) -> Optional[dict]:
    """Return the stored document record, or None if not found."""
    with doc_lock:
        return documents_db.get(doc_id)


def list_documents() -> list[dict]:
    """Return a summary list of all processed documents."""
    with doc_lock:
        return [
            {
                "doc_id": doc_id,
                "filename": doc["filename"],
                "num_chunks": doc["metadata"]["num_chunks"],
            }
            for doc_id, doc in documents_db.items()
        ]


def delete_document(doc_id: str) -> bool:
    """Remove a document from the store. Returns True if it existed."""
    with doc_lock:
        if doc_id in documents_db:
            del documents_db[doc_id]
            return True
        return False