core.py

"""
Core document processing and retrieval engine.
Handles PDF extraction, OCR, chunking, indexing, and hybrid retrieval.
"""

import logging
import os
import re
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

import fitz  # PyMuPDF

from config import config

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Optional dependencies
# ---------------------------------------------------------------------------

try:
    import easyocr
    easyocr_available = True
except ImportError:
    easyocr_available = False
    logger.warning("easyocr not available. Scanned PDFs will not be OCR-processed.")

try:
    from sentence_transformers import SentenceTransformer, CrossEncoder
    import faiss
    st_available = True
except ImportError:
    st_available = False
    logger.warning("sentence-transformers / FAISS not available. Semantic retrieval disabled.")

try:
    import nltk
    # Download both punkt and punkt_tab — newer NLTK (3.9+) uses punkt_tab
    # while older versions use punkt. Download both for full compatibility.
    nltk.download("punkt",      quiet=True)
    nltk.download("punkt_tab",  quiet=True)
    nltk.download("stopwords",  quiet=True)
    from nltk.tokenize import sent_tokenize
    nltk_available = True
    logger.info("NLTK punkt and punkt_tab downloaded successfully.")
except Exception as e:
    nltk_available = False
    logger.warning(f"NLTK unavailable, falling back to regex sentence split: {e}")


def _sentence_split(text: str) -> List[str]:
    if nltk_available:
        return sent_tokenize(text)
    return re.split(r"(?<=[.!?])\s+", text)


# =============================================================================
# OCR PROCESSOR
# =============================================================================

class OCRProcessor:
    def __init__(self):
        self._reader = None
        self._initialized = False

    def _init(self) -> bool:
        if self._initialized:
            return True
        if not easyocr_available:
            return False
        try:
            self._reader = easyocr.Reader(
                config.OCR_LANGUAGES,
                gpu=config.USE_GPU,
                model_storage_directory="/tmp/easyocr_models",
                download_enabled=True,
            )
            self._initialized = True
            logger.info("OCR processor initialized.")
            return True
        except Exception as e:
            logger.error(f"OCR initialization failed: {e}")
            return False

    def extract(self, image_path: str) -> str:
        if not self._init():
            return ""
        try:
            results = self._reader.readtext(image_path, detail=0)
            return " ".join(results) if results else ""
        except Exception as e:
            logger.error(f"OCR failed on {image_path}: {e}")
            return ""


# =============================================================================
# PDF TEXT EXTRACTOR
# =============================================================================

class PDFExtractor:
    def __init__(self):
        self._ocr = OCRProcessor()

    def extract(self, pdf_path: str) -> Dict[str, Any]:
        """
        Extract text from a PDF file.
        Returns a dict with keys: text, total_pages, has_ocr_content, pages.
        """
        try:
            doc = fitz.open(pdf_path)
            pages_data = []
            full_text = ""
            has_ocr = False

            for page_num in range(len(doc)):
                page = doc[page_num]
                text = page.get_text()

                ocr_text = ""
                if len(text.strip()) < 50:
                    img_path = f"/tmp/page_{page_num}.png"
                    try:
                        pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
                        pix.save(img_path)
                        ocr_text = self._ocr.extract(img_path)
                        if ocr_text:
                            text = ocr_text
                            has_ocr = True
                    except Exception as e:
                        logger.warning(f"OCR failed for page {page_num}: {e}")
                    finally:
                        if os.path.exists(img_path):
                            os.remove(img_path)

                page_record = {
                    "page_number": page_num + 1,
                    "text": text,
                    "has_ocr": bool(ocr_text),
                    "char_count": len(text),
                }
                pages_data.append(page_record)
                full_text += f"\n[Page {page_num + 1}]\n{text}\n"

            doc.close()

            return {
                "text": full_text,
                "pages": pages_data,
                "total_pages": len(pages_data),
                "has_ocr_content": has_ocr,
                "total_chars": len(full_text),
            }

        except Exception as e:
            logger.error(f"Failed to process PDF {pdf_path}: {e}")
            raise


# =============================================================================
# TEXT CHUNKER
# =============================================================================

class TextChunker:
    def create_chunks(self, text: str, doc_name: str) -> List[Dict]:
        """Split text into overlapping sentence-based chunks."""
        try:
            sentences = _sentence_split(text)
            chunks: List[Dict] = []
            current: List[str] = []
            current_len = 0

            for sentence in sentences:
                slen = len(sentence)
                if current_len + slen > config.CHUNK_SIZE and current:
                    chunk_text = " ".join(current)
                    if len(chunk_text) >= config.MIN_CHUNK_LENGTH:
                        chunks.append({
                            "text": chunk_text,
                            "document": doc_name,
                            "chunk_id": len(chunks),
                            "length": len(chunk_text),
                        })
                    current = current[-1:] if config.CHUNK_OVERLAP > 0 else []
                    current_len = sum(len(s) for s in current)

                current.append(sentence)
                current_len += slen

            if current:
                chunk_text = " ".join(current)
                if len(chunk_text) >= config.MIN_CHUNK_LENGTH:
                    chunks.append({
                        "text": chunk_text,
                        "document": doc_name,
                        "chunk_id": len(chunks),
                        "length": len(chunk_text),
                    })

            return chunks
        except Exception as e:
            logger.error(f"Chunking failed: {e}")
            return []


# =============================================================================
# RETRIEVER
# =============================================================================

class HybridRetriever:
    """
    Combines TF-IDF lexical search with FAISS-backed semantic search.
    Score = LEXICAL_WEIGHT * tfidf + SEMANTIC_WEIGHT * embedding_similarity
    """

    def __init__(self):
        self._embedding_model: Optional[Any] = None
        self._tfidf: Optional[TfidfVectorizer] = None
        self._tfidf_matrix = None
        self._embeddings: Optional[np.ndarray] = None
        self._faiss_index = None
        self.chunks: List[Dict] = []

    def initialize(self):
        if st_available:
            try:
                logger.info(f"Loading embedding model: {config.EMBEDDING_MODEL}")
                self._embedding_model = SentenceTransformer(config.EMBEDDING_MODEL)
                logger.info("Embedding model loaded.")
            except Exception as e:
                logger.error(f"Failed to load embedding model: {e}")

    def index(self, chunks: List[Dict]):
        self.chunks = chunks
        texts = [c["text"] for c in chunks]

        if not texts:
            logger.error(
                "Indexing received 0 chunks. This usually means NLTK punkt data "
                "was not downloaded and sentence splitting failed. "
                "Chunking will fall back to regex splitting."
            )
            return

        # TF-IDF index
        self._tfidf = TfidfVectorizer(
            max_features=5000,
            stop_words="english",
            ngram_range=(1, 2),
            min_df=1,
            max_df=0.95,
        )
        try:
            self._tfidf_matrix = self._tfidf.fit_transform(texts)
        except ValueError as e:
            logger.error(f"TF-IDF fit failed: {e}. Retrying without stop_words filter.")
            self._tfidf = TfidfVectorizer(
                max_features=5000,
                stop_words=None,
                ngram_range=(1, 1),
                min_df=1,
            )
            self._tfidf_matrix = self._tfidf.fit_transform(texts)

        # Semantic index
        if self._embedding_model:
            self._embeddings = self._embedding_model.encode(
                texts,
                convert_to_numpy=True,
                show_progress_bar=False,
                batch_size=8,
            )
            dim = self._embeddings.shape[1]
            self._faiss_index = faiss.IndexFlatIP(dim)
            self._faiss_index.add(self._embeddings.astype("float32"))

        logger.info(f"Indexed {len(chunks)} chunks.")

    def retrieve(self, query: str, top_k: int = None) -> List[Dict]:
        top_k = top_k or config.MAX_CHUNKS_RETRIEVE
        if not self.chunks:
            return []

        n = len(self.chunks)
        lex_scores = np.zeros(n)
        sem_scores = np.zeros(n)

        if self._tfidf and self._tfidf_matrix is not None:
            q_vec = self._tfidf.transform([query])
            lex_scores = cosine_similarity(q_vec, self._tfidf_matrix)[0]

        if self._embedding_model and self._embeddings is not None:
            q_emb = self._embedding_model.encode([query], convert_to_numpy=True)
            if self._faiss_index:
                sims, idxs = self._faiss_index.search(
                    q_emb.astype("float32"), min(n, top_k)
                )
                for idx, sim in zip(idxs[0], sims[0]):
                    if idx < n:
                        sem_scores[idx] = sim
            else:
                sem_scores = cosine_similarity(q_emb, self._embeddings)[0]

        combined = config.LEXICAL_WEIGHT * lex_scores + config.SEMANTIC_WEIGHT * sem_scores
        top_idx = np.argsort(combined)[::-1][:top_k]

        results = []
        for idx in top_idx:
            if combined[idx] > 0:
                chunk = dict(self.chunks[idx])
                chunk["score"] = float(combined[idx])
                results.append(chunk)

        return results


# =============================================================================
# RERANKER
# =============================================================================

class CrossEncoderReranker:
    def __init__(self):
        self._model: Optional[Any] = None

    def initialize(self):
        if st_available:
            try:
                logger.info(f"Loading reranker: {config.RERANKER_MODEL}")
                self._model = CrossEncoder(config.RERANKER_MODEL)
                logger.info("Reranker loaded.")
            except Exception as e:
                logger.error(f"Failed to load reranker: {e}")

    def rerank(self, query: str, chunks: List[Dict]) -> List[Dict]:
        top_k = config.TOP_K_AFTER_RERANK
        if not self._model or not chunks:
            return chunks[:top_k]

        pairs = [[query, c["text"]] for c in chunks]
        try:
            scores = self._model.predict(pairs)
            for chunk, score in zip(chunks, scores):
                chunk["rerank_score"] = float(score)
            reranked = sorted(chunks, key=lambda x: x.get("rerank_score", 0), reverse=True)
            return reranked[:top_k]
        except Exception as e:
            logger.error(f"Reranking failed: {e}")
            return chunks[:top_k]


# =============================================================================
# DOCUMENT STORE (in-memory)
# =============================================================================

class DocumentStore:
    """Holds all processed document state for the current session."""

    def __init__(self):
        self.documents: List[Dict] = []
        self.chunks: List[Dict] = []
        self.processed_file_paths: List[str] = []

    def clear(self):
        self.documents.clear()
        self.chunks.clear()
        self.processed_file_paths.clear()

    def add_document(self, doc_meta: Dict, chunks: List[Dict], file_path: str):
        self.documents.append(doc_meta)
        self.chunks.extend(chunks)
        self.processed_file_paths.append(file_path)

    def is_empty(self) -> bool:
        return len(self.chunks) == 0