pdf_loader.py

"""
pdf_loader.py — PDF parsing and chunking pipeline.

Takes a PDF file path, extracts text page-by-page using PyMuPDF,
then splits into overlapping chunks suitable for embedding.
"""

import json
import re
from pathlib import Path
from typing import List, Dict, Any

import fitz  # PyMuPDF

from config import CHUNK_SIZE, CHUNK_OVERLAP, CHUNKS_JSON_PATH, DOCUMENTS_DIR


# ── Chunking ─────────────────────────────────────────────────────────────────

def _split_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> List[str]:
    """
    Manual recursive character text splitter with overlap.
    Tries to break on paragraphs → sentences → words before hard-cutting.
    """
    if len(text) <= chunk_size:
        return [text] if text.strip() else []

    chunks: List[str] = []
    start = 0

    while start < len(text):
        end = start + chunk_size

        if end >= len(text):
            chunk = text[start:]
            if chunk.strip():
                chunks.append(chunk.strip())
            break

        # Try to find a clean break point (paragraph > newline > space)
        slice_str = text[start:end]
        break_point = None

        for sep in ["\n\n", "\n", ". ", " "]:
            idx = slice_str.rfind(sep)
            if idx != -1 and idx > chunk_size // 2:
                break_point = idx + len(sep)
                break

        if break_point is None:
            break_point = chunk_size

        chunk = text[start : start + break_point].strip()
        if chunk:
            chunks.append(chunk)

        # Advance with overlap
        start += break_point - overlap

    return chunks


# ── PDF extraction ────────────────────────────────────────────────────────────

def extract_text_from_pdf(pdf_path: str | Path) -> List[Dict[str, Any]]:
    """
    Extract text from each page of a PDF.

    Returns a list of dicts:
        { "page": int, "text": str }
    """
    pdf_path = Path(pdf_path)
    if not pdf_path.exists():
        raise FileNotFoundError(f"PDF not found: {pdf_path}")

    pages = []
    with fitz.open(str(pdf_path)) as doc:
        for page_num, page in enumerate(doc, start=1):
            text = page.get_text("text")
            # Collapse excessive whitespace but keep paragraph breaks
            text = re.sub(r"[ \t]{2,}", " ", text)
            text = re.sub(r"\n{3,}", "\n\n", text)
            if text.strip():
                pages.append({"page": page_num, "text": text.strip()})

    return pages


def load_and_chunk_pdf(pdf_path: str | Path) -> List[Dict[str, Any]]:
    """
    Full pipeline: extract text → chunk → return chunk metadata list.

    Each chunk dict contains:
        {
            "source": str,         # filename
            "page": int,           # page number the chunk came from
            "chunk_index": int,    # sequential index across all chunks
            "text": str            # chunk text
        }
    """
    pdf_path = Path(pdf_path)
    pages = extract_text_from_pdf(pdf_path)

    all_chunks: List[Dict[str, Any]] = []
    chunk_index = 0

    for page_info in pages:
        page_chunks = _split_text(page_info["text"])
        for chunk_text in page_chunks:
            all_chunks.append({
                "source": pdf_path.name,
                "page": page_info["page"],
                "chunk_index": chunk_index,
                "text": chunk_text,
            })
            chunk_index += 1

    return all_chunks


# ── Persistence ───────────────────────────────────────────────────────────────

def load_chunks_from_cache() -> List[Dict[str, Any]]:
    """Load previously saved chunks from disk."""
    if CHUNKS_JSON_PATH.exists():
        try:
            with open(CHUNKS_JSON_PATH, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            return []
    return []


def save_chunks_to_cache(chunks: List[Dict[str, Any]]):
    """Persist chunks to cache/chunks.json."""
    CHUNKS_JSON_PATH.parent.mkdir(parents=True, exist_ok=True)
    with open(CHUNKS_JSON_PATH, "w", encoding="utf-8") as f:
        json.dump(chunks, f, indent=2, ensure_ascii=False)


def ingest_pdf_file(pdf_path: str | Path, profile_id: str = "") -> List[Dict[str, Any]]:
    """
    Convenience wrapper: parse, chunk, merge with existing cache,
    deduplicate by source+page+chunk_index, and save.

    Returns the full updated chunk list.
    """
    new_chunks = load_and_chunk_pdf(pdf_path)
    if profile_id:
        for chunk in new_chunks:
            chunk["profile_id"] = profile_id
    existing = load_chunks_from_cache()

    # Remove stale chunks for this profile + source before re-adding
    source_name = Path(pdf_path).name
    existing = [
        c for c in existing
        if not (c.get("source") == source_name and c.get("profile_id", "") == profile_id)
    ]

    merged = existing + new_chunks
    save_chunks_to_cache(merged)
    return merged


def get_indexed_sources(profile_id: str = "") -> List[str]:
    """Return list of document file names in the chunk cache for a profile."""
    chunks = load_chunks_from_cache()
    if profile_id:
        chunks = [c for c in chunks if c.get("profile_id") == profile_id]
    return sorted({c["source"] for c in chunks})


def remove_source(source_name: str):
    """Remove all chunks from a given source file and save."""
    chunks = load_chunks_from_cache()
    chunks = [c for c in chunks if c.get("source") != source_name]
    save_chunks_to_cache(chunks)