document_engine.py

import hashlib
import logging
import os
import re
import time
import queue
import threading
import json
import sqlite3
from typing import List, NotRequired, Optional
import msoffcrypto  # noqa
import openpyxl  # noqa
import unstructured  # noqa
import docx  # noqa

# Concurrency imports
import concurrent.futures

from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler

from langchain_ollama import ChatOllama, OllamaEmbeddings
from pydantic import BaseModel, Field
from typing_extensions import TypedDict

from langchain_community.document_loaders import (
    PyPDFLoader,
    CSVLoader,
    TextLoader
)
from langchain_community.vectorstores.utils import filter_complex_metadata
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langgraph.graph import END, StateGraph, START

from rank_bm25 import BM25Okapi
from langchain_experimental.text_splitter import SemanticChunker
from langchain_text_splitters import RecursiveCharacterTextSplitter

from lore_utils import CHROMA_COLLECTION_NAME, CHROMA_DB_PATH, THINKING_OLLAMA_MODEL, FAST_OLLAMA_MODEL, \
    EMBEDDING_MODEL, RERANK_MODEL, SUPPORTED_EXTENSIONS, DOC_FOLDER

logger = logging.getLogger(__name__)

GLOBAL_VECTORSTORE: Optional[Chroma] = None


def _format_location(meta: dict) -> str:
    """Convert chunk metadata into a human-readable location string."""
    if "line_start" in meta:
        ls, le = meta["line_start"], meta.get("line_end", meta["line_start"])
        return f"Line {ls}" if ls == le else f"Lines {ls}–{le}"
    if meta.get("sheet") and "row" in meta:
        return f"Sheet '{meta['sheet']}' · Row {meta['row']}"
    if "row" in meta:
        return f"Row {meta['row']}"
    if "paragraph_index" in meta:
        return f"Paragraph {meta['paragraph_index']}"
    if "page" in meta:
        return f"Page {meta['page'] + 1}"
    if "start_index" in meta:
        return f"Char {meta['start_index']}"
    return "Unknown"


def _citation_location(meta: dict) -> str:
    """Location string for citations: prefer page number, then line, then other."""
    if "page" in meta:
        return f"Page {meta['page'] + 1}"
    if "line_start" in meta:
        ls, le = meta["line_start"], meta.get("line_end", meta["line_start"])
        return f"Line {ls}" if ls == le else f"Lines {ls}–{le}"
    if meta.get("sheet") and "row" in meta:
        return f"Sheet '{meta['sheet']}' · Row {meta['row']}"
    if "row" in meta:
        return f"Row {meta['row']}"
    if "paragraph_index" in meta:
        return f"Paragraph {meta['paragraph_index']}"
    if "start_index" in meta:
        return f"Char {meta['start_index']}"
    return ""


def _enrich_line_numbers(splits: list) -> list:
    """Add line_start and line_end to splits from .txt/.md files.

    The loader stores a JSON-encoded newline-offset list in
    ``_line_offsets_json`` metadata. This function reads it, computes line
    numbers via bisect, writes line_start / line_end, then removes the
    temporary key so it is not stored in ChromaDB.
    """
    import bisect
    for doc in splits:
        raw = doc.metadata.pop("_line_offsets_json", None)
        if raw is None:
            continue
        try:
            offsets = json.loads(raw)
        except Exception:
            continue
        start = doc.metadata.get("start_index", 0)
        end = start + len(doc.page_content)
        doc.metadata["line_start"] = bisect.bisect_right(offsets, start)
        doc.metadata["line_end"]   = bisect.bisect_right(offsets, max(end - 1, start))
    return splits


def _enrich_page_numbers(splits: list) -> list:
    """Add ``page`` (0-indexed) to splits from .docx files.

    The loader stores a JSON-encoded list of [char_offset, page_0indexed] pairs
    in ``_page_offsets_json``. This function uses bisect to find which page
    each chunk starts on and removes the temporary key before ChromaDB storage.
    """
    import bisect
    for doc in splits:
        raw = doc.metadata.pop("_page_offsets_json", None)
        if raw is None:
            continue
        try:
            offsets_pages = json.loads(raw)
        except Exception:
            continue
        if not offsets_pages:
            doc.metadata.setdefault("page", 0)
            continue
        offsets = [op[0] for op in offsets_pages]
        pages = [op[1] for op in offsets_pages]
        start = doc.metadata.get("start_index", 0)
        idx = bisect.bisect_right(offsets, start) - 1
        doc.metadata["page"] = pages[idx] if idx >= 0 else 0
    return splits


def _expand_doc_context(doc: Document, n_neighbors: int = 1) -> str:
    """Return expanded context for a chunk by including neighboring chunks.

    Fetches all chunks for the same source file from ChromaDB, sorts by
    start_index, and returns the text of the target chunk plus up to
    n_neighbors on each side joined by newlines.  Falls back to the
    original chunk text if the vectorstore is unavailable or the chunk
    cannot be located.
    """
    if GLOBAL_VECTORSTORE is None:
        return doc.page_content
    source = doc.metadata.get("source")
    if not source:
        return doc.page_content
    try:
        results = GLOBAL_VECTORSTORE.get(
            where={"source": source},
            include=["documents", "metadatas"],
        )
        pairs = list(zip(results.get("documents", []), results.get("metadatas", [])))
        pairs.sort(key=lambda x: x[1].get("start_index", 0))
        texts = [t for t, _ in pairs]
        metas = [m for _, m in pairs]

        target_start = doc.metadata.get("start_index", -1)
        pos = next((i for i, m in enumerate(metas) if m.get("start_index") == target_start), None)
        if pos is None:
            pos = next((i for i, t in enumerate(texts) if t == doc.page_content), None)
        if pos is None:
            return doc.page_content

        lo = max(0, pos - n_neighbors)
        hi = min(len(texts) - 1, pos + n_neighbors)
        return "\n".join(texts[lo:hi + 1])
    except Exception:
        return doc.page_content


# --- BM25 INDEX (rebuilt in-memory after each ingestion) ---

_BM25_INDEX = None          # rank_bm25.BM25Okapi instance
_BM25_DOCS: list = []       # Document list parallel to the BM25 corpus
_BM25_LOCK = threading.Lock()


def _build_bm25_index() -> None:
    """Rebuild the in-memory BM25 index from all documents currently in ChromaDB.

    Called at the end of initialize_vectorstore() and after each ingestion in
    the background worker so the index stays in sync with the vector store.
    """
    global _BM25_INDEX, _BM25_DOCS
    if GLOBAL_VECTORSTORE is None:
        return
    try:
        results = GLOBAL_VECTORSTORE.get(include=["documents", "metadatas"])
        texts = results.get("documents") or []
        metas = results.get("metadatas") or []
        if not texts:
            with _BM25_LOCK:
                _BM25_INDEX = None
                _BM25_DOCS = []
            return
        docs = [Document(page_content=t, metadata=m) for t, m in zip(texts, metas)]
        tokenized = [t.lower().split() for t in texts]
        index = BM25Okapi(tokenized)
        with _BM25_LOCK:
            _BM25_INDEX = index
            _BM25_DOCS = docs
        logger.debug(f"[BM25] Rebuilt index with {len(docs)} documents")
    except Exception as e:
        logger.warning(f"[BM25] Failed to build index: {e}")


def _bm25_search(
    query: str,
    k: int = 8,
    source_filter: Optional[str] = None,
    tag_filter: Optional[str] = None,
) -> list:
    """Return up to k Documents from the BM25 index with positive scores."""
    with _BM25_LOCK:
        if _BM25_INDEX is None or not _BM25_DOCS:
            return []
        scores = _BM25_INDEX.get_scores(query.lower().split())
        pairs = list(zip(scores, _BM25_DOCS))

    if source_filter:
        pairs = [(s, d) for s, d in pairs if d.metadata.get("source") == source_filter]
    if tag_filter:
        pairs = [(s, d) for s, d in pairs if tag_filter in d.metadata.get("tags", "")]

    pairs.sort(key=lambda x: x[0], reverse=True)
    return [d for s, d in pairs[:k] if s > 0]


def _rrf_merge(ranked_lists: list, k: int = 60) -> list:
    """Merge multiple ranked document lists using Reciprocal Rank Fusion.

    A document appearing at rank r in a list contributes 1/(k+r+1) to its
    total score. Documents appearing in multiple lists accumulate scores,
    naturally promoting consistent results.
    """
    scores: dict[int, float] = {}
    doc_map: dict[int, object] = {}
    for lst in ranked_lists:
        for rank, doc in enumerate(lst):
            key = hash(doc.page_content)
            scores[key] = scores.get(key, 0.0) + 1.0 / (k + rank + 1)
            doc_map[key] = doc
    return [doc_map[key] for key in sorted(scores, key=lambda x: scores[x], reverse=True)]


def _rerank_documents(query: str, docs: list) -> list:
    """Re-rank docs using Ollama's rerank API.

    Returns docs in descending relevance order. Falls back to the original
    order silently if RERANK_MODEL is not configured or the call fails.
    """
    if not RERANK_MODEL or not docs:
        return docs
    try:
        import json as _json
        import urllib.request as _urllib
        payload = _json.dumps({
            "model": RERANK_MODEL,
            "query": query,
            "documents": [d.page_content for d in docs],
        }).encode()
        req = _urllib.Request(
            "http://localhost:11434/api/rerank",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with _urllib.urlopen(req, timeout=30) as resp:
            data = _json.loads(resp.read())
        results = data.get("results") or []
        if not results:
            return docs
        ranked = sorted(results, key=lambda r: r.get("relevance_score", 0.0), reverse=True)
        return [docs[r["index"]] for r in ranked if r["index"] < len(docs)]
    except Exception as e:
        logger.warning(f"[RERANK] {e}")
        return docs


INGESTION_QUEUE = queue.Queue()
# Completed ingestion events: (action, file_path, success, chunk_count, error_msg)
COMPLETION_QUEUE: queue.Queue = queue.Queue()
MANIFEST_LOCK = threading.Lock()
INDEXED_FILES_PATH = os.path.join(CHROMA_DB_PATH, "indexed_files.txt")

_EMBEDDINGS: Optional[OllamaEmbeddings] = None
_OBSERVER: Optional[Observer] = None


def _get_embeddings() -> OllamaEmbeddings:
    """Returns the shared OllamaEmbeddings instance, creating it once on first call."""
    global _EMBEDDINGS
    if _EMBEDDINGS is None:
        _EMBEDDINGS = OllamaEmbeddings(model=EMBEDDING_MODEL)
    return _EMBEDDINGS


_CHAR_SPLITTER = RecursiveCharacterTextSplitter(
    chunk_size=1000, chunk_overlap=200, add_start_index=True
)


def _get_semantic_splitter() -> SemanticChunker:
    """Return a SemanticChunker that splits at topic boundaries using the shared embeddings model.

    Uses percentile-based breakpoint detection, which splits wherever the cosine
    distance between adjacent sentence embeddings exceeds the 95th-percentile
    threshold for the document.  The sentence-split regex includes newlines so
    that long paragraphs are broken into embeddable pieces before the embedding
    step, avoiding 400 errors from models with a short context window (e.g.
    mxbai-embed-large at 512 tokens).  ``min_chunk_size=200`` prevents orphan
    fragments.  ``add_start_index=True`` preserves character offsets so
    page/line enrichment continues to work.
    """
    return SemanticChunker(
        embeddings=_get_embeddings(),
        add_start_index=True,
        breakpoint_threshold_type="percentile",
        min_chunk_size=200,
        sentence_split_regex=r"(?<=[.?!])\s+|\n+",
    )


# mxbai-embed-large has a 512-token context; at ~4 chars/token that is ~2048 chars.
# We use a conservative limit so chunks are safely within bounds.
_MAX_EMBED_CHARS = 1500


def _split_documents(docs: list) -> list:
    """Split documents using SemanticChunker with a hard size ceiling.

    For each document, semantic chunking is attempted first. If it raises
    (e.g. an individual sentence still exceeds the embedding context window),
    that document falls back to character-based splitting.

    After semantic splitting, any chunk that exceeds _MAX_EMBED_CHARS is
    re-split with RecursiveCharacterTextSplitter so that no chunk sent to
    ChromaDB's add_documents() can trigger a 400 from the embedding model.
    """
    semantic = _get_semantic_splitter()
    raw: list = []
    for doc in docs:
        try:
            raw.extend(semantic.split_documents([doc]))
        except Exception as e:
            src = os.path.basename(doc.metadata.get("source", "?"))
            logger.warning(f"Semantic chunking failed for {src} ({e}); using character split")
            raw.extend(_CHAR_SPLITTER.split_documents([doc]))

    # Second pass: enforce the embedding context limit
    final: list = []
    for chunk in raw:
        if len(chunk.page_content) > _MAX_EMBED_CHARS:
            final.extend(_CHAR_SPLITTER.split_documents([chunk]))
        else:
            final.append(chunk)
    return final


def _load_docx(file_path: str) -> List[Document]:
    """Load a .docx file using python-docx with explicit page-break tracking.

    Walks every paragraph in document order, tracking page breaks via:
      - ``w:pageBreakBefore`` paragraph property
      - ``w:sectPr`` section break at the end of a paragraph
      - ``w:br w:type="page"`` run-level page break

    Returns a SINGLE Document containing all paragraph text joined by newlines.
    Page-break character offsets are stored in ``_page_offsets_json`` metadata
    (a JSON list of [char_offset, page_0indexed] pairs) so that
    ``_enrich_page_numbers`` can assign the correct page to each split chunk.
    Joining into one document lets the text splitter create chunks that span
    paragraph boundaries, preventing short headings from becoming isolated chunks.
    """
    import docx as _docx
    from docx.oxml.ns import qn

    try:
        doc = _docx.Document(file_path)
    except Exception as e:
        logger.error(f"Failed to open {file_path}: {e}")
        return []

    text_parts: list[str] = []
    # Each entry: [char_offset_into_full_text, page_0indexed]
    page_offsets: list[list[int]] = []
    page = 1
    current_pos = 0
    prev_section_break = False

    for para in doc.paragraphs:
        p = para._p
        pPr = p.find(qn('w:pPr'))

        page_break_before = prev_section_break
        prev_section_break = False

        if pPr is not None:
            pb_elem = pPr.find(qn('w:pageBreakBefore'))
            if pb_elem is not None:
                val = pb_elem.get(qn('w:val'), 'true')
                if val.lower() not in ('false', '0', 'off'):
                    page_break_before = True

            sectPr = pPr.find(qn('w:sectPr'))
            if sectPr is not None:
                type_elem = sectPr.find(qn('w:type'))
                break_type = (type_elem.get(qn('w:val'), 'nextPage')
                              if type_elem is not None else 'nextPage')
                if break_type in ('nextPage', 'evenPage', 'oddPage'):
                    prev_section_break = True

        if page_break_before:
            page += 1
            page_offsets.append([current_pos, page - 1])

        has_run_break = any(
            br.get(qn('w:type')) == 'page'
            for br in p.findall('.//' + qn('w:br'))
        )

        text = para.text
        if text.strip():
            text_parts.append(text)
            current_pos += len(text) + 1  # +1 for the joining newline

        if has_run_break:
            page += 1
            page_offsets.append([current_pos, page - 1])

    if not text_parts:
        return []

    full_text = '\n'.join(text_parts)
    return [Document(
        page_content=full_text,
        metadata={
            "source": file_path,
            "_page_offsets_json": json.dumps(page_offsets),
        },
    )]


def load_pdf(file_path: str) -> List[Document]:
    loader = PyPDFLoader(file_path)
    docs = loader.load()
    return docs


def _load_xlsx(file_path: str) -> List[Document]:
    """Load an Excel workbook row-by-row using openpyxl, preserving sheet/row metadata."""
    import openpyxl
    docs = []
    try:
        wb = openpyxl.load_workbook(file_path, read_only=True, data_only=True)
        for sheet_name in wb.sheetnames:
            ws = wb[sheet_name]
            for row in ws.iter_rows():
                row_num = row[0].row
                text = "\t".join(
                    str(cell.value) for cell in row if cell.value is not None
                ).strip()
                if text:
                    docs.append(Document(
                        page_content=text,
                        metadata={"source": file_path, "sheet": sheet_name, "row": row_num},
                    ))
        wb.close()
    except Exception as e:
        logger.error(f"Error loading xlsx {file_path}: {e}")
    return docs


def load_document_by_extension(file_path: str) -> List[Document]:
    """
    Selects the appropriate loader based on file extension.
    Top-level function for pickling support in multiprocessing.
    """
    # Safety check for watcher race conditions
    if not os.path.exists(file_path):
        return []

    ext = os.path.splitext(file_path)[1].lower()
    try:
        if ext == '.pdf':
            return load_pdf(file_path)
        elif ext == '.docx':
            return _load_docx(file_path)
        elif ext == '.xlsx':
            return _load_xlsx(file_path)
        elif ext == '.csv':
            loader = CSVLoader(file_path)
            return loader.load()
        elif ext in ['.txt', '.md']:
            loader = TextLoader(file_path, encoding='utf-8', autodetect_encoding=True)
            docs = loader.load()
            # Pre-compute newline offsets for post-split line number enrichment
            if docs:
                raw_text = docs[0].page_content
                offsets = [0]
                for line in raw_text.splitlines(keepends=True):
                    offsets.append(offsets[-1] + len(line))
                offsets_json = json.dumps(offsets)
                for doc in docs:
                    doc.metadata["_line_offsets_json"] = offsets_json
            return docs
        return []
    except Exception as e:
        logger.error(f"Error loading {file_path}: {e}")
        return []


class DocumentEventHandler(FileSystemEventHandler):
    """
    Watchdog Event Handler.
    Pushes events to a thread-safe queue to be processed by the worker.
    """

    def on_created(self, event):
        if not event.is_directory and event.src_path.lower().endswith(SUPPORTED_EXTENSIONS):
            logger.info(f"[WATCHER] File Created: {event.src_path}")
            INGESTION_QUEUE.put(("add", event.src_path))

    def on_modified(self, event):
        if not event.is_directory and event.src_path.lower().endswith(SUPPORTED_EXTENSIONS):
            logger.info(f"[WATCHER] File Modified: {event.src_path}")
            INGESTION_QUEUE.put(("update", event.src_path))

    def on_deleted(self, event):
        if not event.is_directory and event.src_path.lower().endswith(SUPPORTED_EXTENSIONS):
            logger.info(f"[WATCHER] File Deleted: {event.src_path}")
            INGESTION_QUEUE.put(("delete", event.src_path))

    def on_moved(self, event):
        if not event.is_directory:
            if event.src_path.lower().endswith(SUPPORTED_EXTENSIONS):
                logger.info(f"[WATCHER] File Moved (Delete old): {event.src_path}")
                INGESTION_QUEUE.put(("delete", event.src_path))
            if event.dest_path.lower().endswith(SUPPORTED_EXTENSIONS):
                logger.info(f"[WATCHER] File Moved (Add new): {event.dest_path}")
                INGESTION_QUEUE.put(("add", event.dest_path))


def _get_file_signature(file_path: str) -> dict:
    hasher = hashlib.md5()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            hasher.update(chunk)
    return {
        "mtime": os.path.getmtime(file_path),
        "size": os.path.getsize(file_path),
        "hash": hasher.hexdigest(),
    }


def get_duplicate_source(file_path: str) -> Optional[str]:
    """
    Returns the basename of an already-indexed file whose content is identical
    to file_path, or None if no duplicate exists.
    """
    if not os.path.exists(file_path):
        return None
    sig = _get_file_signature(file_path)
    new_hash = sig.get("hash")
    manifest = _load_index_manifest()
    for indexed_path, stored_sig in manifest.items():
        if indexed_path == file_path:
            continue
        if stored_sig.get("hash") == new_hash:
            return os.path.basename(indexed_path)
    return None


def _load_index_manifest() -> dict:
    if not os.path.exists(INDEXED_FILES_PATH):
        return {}
    try:
        with open(INDEXED_FILES_PATH, "r", encoding="utf-8") as f:
            raw = f.read().strip()
        if not raw:
            return {}
        if raw.startswith("{"):
            return json.loads(raw)
        # Backward compatibility with legacy newline list
        return {line.strip(): {} for line in raw.splitlines() if line.strip()}
    except Exception as e:
        logger.error(f"Error reading index manifest: {e}")
        return {}


def _write_index_manifest(manifest: dict) -> None:
    try:
        with open(INDEXED_FILES_PATH, "w", encoding="utf-8") as f:
            json.dump(manifest, f, indent=2)
    except Exception as e:
        logger.error(f"Error writing index manifest: {e}")


def _update_manifest(action: str, file_path: str, tags: Optional[list] = None) -> None:
    with MANIFEST_LOCK:
        manifest = _load_index_manifest()
        if action == "delete":
            manifest.pop(file_path, None)
            _write_index_manifest(manifest)
            return
        if os.path.exists(file_path):
            existing_tags = manifest.get(file_path, {}).get("tags", [])
            sig = _get_file_signature(file_path)
            sig["tags"] = tags if tags is not None else existing_tags
            manifest[file_path] = sig
            _write_index_manifest(manifest)


def get_tags(file_path: str) -> list[str]:
    """Return the tag list for a file from the manifest, or [] if not found."""
    return _load_index_manifest().get(file_path, {}).get("tags", [])


def set_tags(file_path: str, tags: list[str]) -> None:
    """Replace the tag list for a file and re-queue it for ingestion so
    ChromaDB chunk metadata is refreshed with the new tags."""
    cleaned = [t.lower().strip()[:50] for t in tags if t.strip()][:20]
    with MANIFEST_LOCK:
        manifest = _load_index_manifest()
        if file_path not in manifest:
            return
        manifest[file_path]["tags"] = cleaned
        _write_index_manifest(manifest)
    INGESTION_QUEUE.put(("update", file_path))


def ingestion_worker():
    """
    Background daemon that consumes the queue.
    Handles 'add', 'update', and 'delete' operations on the global VectorStore.
    """
    global GLOBAL_VECTORSTORE

    while True:
        try:
            # Block until an item is available; None is the shutdown sentinel
            item = INGESTION_QUEUE.get()
            if item is None:
                INGESTION_QUEUE.task_done()
                break
            action, file_path = item

            # Wait briefly to let file writes settle (debounce)
            time.sleep(1.0)

            if GLOBAL_VECTORSTORE is None:
                INGESTION_QUEUE.task_done()
                continue

            base_name = os.path.basename(file_path)

            logger.info(f"SYNC WORKER: Processing {action} for {base_name}")

            if action == "delete":
                try:
                    GLOBAL_VECTORSTORE.delete(where={"source": file_path})
                    _update_manifest("delete", file_path)
                    logger.info(f"Removed chunks for {base_name}")
                    COMPLETION_QUEUE.put(("delete", file_path, True, 0, None))
                except Exception as e:
                    logger.error(f"Error deleting {file_path}: {e}")
                    COMPLETION_QUEUE.put(("delete", file_path, False, 0, str(e)))

            elif action in ["add", "update"]:
                # For update, we delete first to avoid duplicates
                if action == "update":
                    try:
                        GLOBAL_VECTORSTORE.delete(where={"source": file_path})
                    except Exception:
                        pass  # Might not exist yet

                # Load and Ingest
                try:
                    docs = load_document_by_extension(file_path)
                    if docs:
                        splits = _split_documents(docs)
                        splits = _enrich_line_numbers(splits)
                        splits = _enrich_page_numbers(splits)
                        # Read current tags from manifest and stamp onto each split
                        manifest_snap = _load_index_manifest()
                        for split in splits:
                            src = split.metadata.get("source", "")
                            t = manifest_snap.get(src, {}).get("tags", [])
                            if t:
                                split.metadata["tags"] = ",".join(t)
                        splits = filter_complex_metadata(splits)
                        if splits:
                            GLOBAL_VECTORSTORE.add_documents(splits)
                            _update_manifest("add", file_path)
                            logger.info(f"Added {len(splits)} chunks for {os.path.basename(file_path)}")
                            COMPLETION_QUEUE.put(("add", file_path, True, len(splits), None))
                        else:
                            COMPLETION_QUEUE.put(("add", file_path, False, 0, "No chunks extracted"))
                    else:
                        COMPLETION_QUEUE.put(("add", file_path, False, 0, "No content loaded"))
                except Exception as e:
                    logger.error(f"Error ingesting {file_path}: {e}")
                    COMPLETION_QUEUE.put(("add", file_path, False, 0, str(e)))

            # Keep the BM25 index in sync after any change
            _build_bm25_index()

            INGESTION_QUEUE.task_done()

        except Exception as e:
            logger.error(f"Worker error: {e}")


# --- PART 1: OPTIMIZED INGESTION ENGINE ---

def initialize_vectorstore():
    """
    Ingests documents using Multiprocessing and returns the VectorStore.
    Also starts the background Watchdog observer.
    """
    global GLOBAL_VECTORSTORE, _OBSERVER

    embeddings = _get_embeddings()

    # Ensure directories exist
    if not os.path.exists(DOC_FOLDER):
        os.makedirs(DOC_FOLDER)
    if not os.path.exists(CHROMA_DB_PATH):
        os.makedirs(CHROMA_DB_PATH)

    # 1. Connect to VectorStore
    vectorstore = Chroma(
        persist_directory=CHROMA_DB_PATH,
        embedding_function=embeddings,
        collection_name=CHROMA_COLLECTION_NAME
    )
    GLOBAL_VECTORSTORE = vectorstore

    # 2. Identify New Files
    manifest = _load_index_manifest()
    indexed_files = set(manifest.keys())
    current_files = {}
    for root, dirs, files in os.walk(DOC_FOLDER):
        for file in files:
            if file.startswith("~$"): continue
            if file.lower().endswith(SUPPORTED_EXTENSIONS):
                path = os.path.join(root, file)
                current_files[path] = _get_file_signature(path)

    current_paths = set(current_files.keys())
    deleted_files = list(indexed_files - current_paths)
    candidate_files = list(current_paths)

    new_files = []
    updated_files = []
    for file_path in candidate_files:
        if file_path not in manifest:
            new_files.append(file_path)
            continue
        if manifest.get(file_path) != current_files[file_path]:
            updated_files.append(file_path)

    # 3. Parallel Loading & Ingestion
    if deleted_files:
        logger.info(f"Detected {len(deleted_files)} deleted document(s)")
        for file_path in deleted_files:
            try:
                vectorstore.delete(where={"source": file_path})
                _update_manifest("delete", file_path)
                logger.info(f"Removed chunks for {os.path.basename(file_path)}")
            except Exception as e:
                logger.error(f"Error deleting {file_path}: {e}")

    files_to_ingest = new_files + updated_files

    if files_to_ingest:
        if new_files:
            logger.info(f"Detected {len(new_files)} new document(s)")
        if updated_files:
            logger.info(f"Detected {len(updated_files)} updated document(s)")

        for file_path in updated_files:
            try:
                vectorstore.delete(where={"source": file_path})
            except Exception:
                pass

        docs = []

        # Optimization: Multiprocessing
        with concurrent.futures.ProcessPoolExecutor() as executor:
            future_to_file = {executor.submit(load_document_by_extension, fp): fp for fp in files_to_ingest}

            for i, future in enumerate(concurrent.futures.as_completed(future_to_file)):
                fp = future_to_file[future]
                try:
                    loaded_docs = future.result()
                    docs.extend(loaded_docs)
                    logger.info(f"[{i + 1}/{len(files_to_ingest)}] Loaded: {os.path.basename(fp)}")
                except Exception as exc:
                    logger.error(f"[{i + 1}/{len(files_to_ingest)}] Failed: {fp}: {exc}")

        if docs:
            logger.info(f"Splitting & embedding {len(docs)} document chunk(s)")
            splits = _split_documents(docs)
            splits = _enrich_line_numbers(splits)
            splits = _enrich_page_numbers(splits)
            # Stamp tags from manifest onto each split
            for split in splits:
                src = split.metadata.get("source", "")
                t = manifest.get(src, {}).get("tags", [])
                if t:
                    split.metadata["tags"] = ",".join(t)
            splits = filter_complex_metadata(splits)

            vectorstore.add_documents(splits)

            # Update index tracking
            for file_path in files_to_ingest:
                _update_manifest("add", file_path)
            logger.info("Ingestion complete")
    else:
        logger.info("No document changes to ingest")

    # 4. Start Background Sync (Watchdog)
    logger.info("Starting live document watcher")

    # Start Worker Thread
    worker_thread = threading.Thread(target=ingestion_worker, daemon=True)
    worker_thread.start()

    # Start Observer
    event_handler = DocumentEventHandler()
    _OBSERVER = Observer()
    _OBSERVER.schedule(event_handler, DOC_FOLDER, recursive=True)
    _OBSERVER.start()

    # Build initial BM25 index from everything now in ChromaDB
    _build_bm25_index()

    return vectorstore


def shutdown() -> None:
    """Stop the ingestion worker and watchdog observer cleanly."""
    INGESTION_QUEUE.put(None)  # sentinel tells the worker to exit
    if _OBSERVER is not None:
        _OBSERVER.stop()
        _OBSERVER.join()
    logger.info("Document engine shutdown complete")


def get_indexed_files() -> dict:
    """Returns the index manifest. Keys are file paths, values are {mtime, size, hash} dicts."""
    return _load_index_manifest()


def get_chunk_counts() -> dict[str, int]:
    """Returns {file_path: chunk_count} by reading ChromaDB metadata in one call."""
    if GLOBAL_VECTORSTORE is None:
        return {}
    try:
        results = GLOBAL_VECTORSTORE.get(include=["metadatas"])
        counts: dict[str, int] = {}
        for meta in results.get("metadatas", []):
            source = meta.get("source", "")
            if source:
                counts[source] = counts.get(source, 0) + 1
        return counts
    except Exception as e:
        logger.warning(f"Could not fetch chunk counts: {e}")
        return {}


def trigger_reindex() -> int:
    """
    Queues every file in the input folder for re-ingestion via the background worker.
    Returns the number of files queued.
    """
    count = 0
    for root, _, files in os.walk(DOC_FOLDER):
        for file in files:
            if file.startswith("~$"):
                continue
            if file.lower().endswith(SUPPORTED_EXTENSIONS):
                INGESTION_QUEUE.put(("update", os.path.join(root, file)))
                count += 1
    return count


def get_retriever(k=4, source_filter: Optional[str] = None, tag_filter: Optional[str] = None):
    """Returns a retriever interface from the current global vectorstore.

    Args:
        k: Number of documents to retrieve.
        source_filter: If provided, restricts results to documents whose
            ``source`` metadata field matches this full file path.
        tag_filter: If provided, restricts results to documents whose
            ``tags`` metadata field contains this tag string.
    """
    global GLOBAL_VECTORSTORE
    if GLOBAL_VECTORSTORE is None:
        initialize_vectorstore()

    search_kwargs: dict = {"k": k, "fetch_k": 20, "lambda_mult": 0.5}
    if source_filter:
        search_kwargs["filter"] = {"source": source_filter}
    if tag_filter:
        tag_f = {"tags": {"$contains": tag_filter}}
        if "filter" in search_kwargs:
            search_kwargs["filter"] = {"$and": [search_kwargs["filter"], tag_f]}
        else:
            search_kwargs["filter"] = tag_f

    return GLOBAL_VECTORSTORE.as_retriever(
        search_type="mmr",
        search_kwargs=search_kwargs,
    )


# --- PART 2: MODELS & PROMPTS ---

# Initialize LLMs
thinking_llm = ChatOllama(model=THINKING_OLLAMA_MODEL, temperature=0)
fast_llm = ChatOllama(model=FAST_OLLAMA_MODEL, temperature=0)


# Grader
class GradeDocuments(BaseModel):
    """Binary score for relevance check on retrieved documents."""
    binary_score: str = Field(description="Documents are relevant to the question, 'yes' or 'no'")


structured_llm_grader = fast_llm.with_structured_output(GradeDocuments)
grader_system = """You are a grader assessing whether a retrieved document contains information needed to answer a user question.
Grade as 'yes' ONLY if the document directly contains information that addresses the specific question — not merely shares the same general topic or setting.
Grade as 'no' if the document's content is tangential or does not actually help answer the question.
Give a binary score 'yes' or 'no'."""
grader_prompt = ChatPromptTemplate.from_messages(
    [("system", grader_system), ("human", "Retrieved document: \n\n {document} \n\n User question: {question}")]
)
grader_chain = grader_prompt | structured_llm_grader

rag_system_prompt = """You are a helpful assistant. Answer the user's question based ONLY on the context provided below.
Each context chunk is labelled with a reference number: [1], [2], etc.

Rules:
1. Base your answer SOLELY on what is explicitly written in the provided chunks.
2. Cite a chunk [N] inline ONLY if the specific fact you are stating appears literally in that chunk's text.
3. Do NOT attribute information to a source if that information does not appear in that source's text.
4. If the context does not contain the answer, say "I don't know."
5. Do NOT add a Sources list at the end — one will be appended automatically.
"""

# RAG Generator
rag_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", rag_system_prompt),
        ("human", "Question: {question} \n\n Context: {context} \n\n Answer:"),
    ]
)
rag_chain = rag_prompt | thinking_llm | StrOutputParser()

# Rewriter
rewrite_system = "You are a question re-writer that converts an input question to a better version for vector retrieval."
rewrite_prompt = ChatPromptTemplate.from_messages(
    [("system", rewrite_system), ("human", "Initial question: \n\n {question} \n Formulate an improved question.")]
)
rewriter_chain = rewrite_prompt | fast_llm | StrOutputParser()

# Multi-query — generate alternative phrasings to improve retrieval recall
multi_query_system = (
    "You are an AI assistant. Generate {n} alternative phrasings of the following question "
    "to improve document retrieval. Each phrasing should approach the topic differently. "
    "Output only the rephrased questions, one per line, numbered."
)
multi_query_prompt = ChatPromptTemplate.from_messages(
    [("system", multi_query_system), ("human", "{question}")]
)
multi_query_chain = multi_query_prompt | fast_llm | StrOutputParser()

# Faithfulness checker — verifies generated answer is grounded in retrieved context
class FaithfulnessScore(BaseModel):
    """Binary faithfulness check for the generated answer."""
    is_grounded: bool = Field(
        description="True if every factual claim in the answer is directly supported by the provided context"
    )


faithfulness_system = (
    "You are a faithfulness checker. Given a context and an answer, decide whether every "
    "factual claim in the answer is directly supported by the context. "
    "Return true only if the answer contains no facts that are absent from or contradict the context."
)
faithfulness_prompt = ChatPromptTemplate.from_messages([
    ("system", faithfulness_system),
    ("human", "Context:\n{context}\n\nAnswer:\n{answer}"),
])
faithfulness_check_chain = faithfulness_prompt | fast_llm.with_structured_output(FaithfulnessScore)


# --- PART 3: GRAPH NODES (OPTIMIZED) ---

class GraphState(TypedDict):
    question: str
    generation: str
    documents: List[Document]
    loop_step: int
    scoped_source: NotRequired[Optional[str]]         # None = search all documents
    stream_queue: NotRequired[Optional[queue.Queue]]  # queue for live token streaming
    tag_filter: NotRequired[Optional[str]]            # optional tag to restrict retrieval
    history: NotRequired[Optional[list]]              # list of {"role", "content"} dicts


def retrieve(state):
    question = state["question"]
    scoped_source = state.get("scoped_source")
    tag_filter = state.get("tag_filter")
    logger.debug(f"[RETRIEVE] question: {question!r}" +
                 (f" | scoped to: {os.path.basename(scoped_source)}" if scoped_source else ""))

    # Generate alternative phrasings to improve recall; fall back gracefully
    queries = [question]
    try:
        variants_text = multi_query_chain.invoke({"question": question, "n": 2})
        for line in variants_text.splitlines():
            variant = re.sub(r"^\d+[\.\)\-]\s*", "", line.strip())
            if variant and variant not in queries:
                queries.append(variant)
    except Exception:
        pass  # use original query only if LLM call fails
    queries = queries[:3]
    logger.debug(f"[RETRIEVE] {len(queries)} query variant(s): {queries}")

    # Vector search: one ranked list per query variant
    retriever = get_retriever(k=4, source_filter=scoped_source, tag_filter=tag_filter)
    vector_lists: list[list[Document]] = [list(retriever.invoke(q)) for q in queries]

    # BM25 keyword search on the original question only
    bm25_results = _bm25_search(question, k=8, source_filter=scoped_source, tag_filter=tag_filter)

    # Merge all result lists with Reciprocal Rank Fusion
    documents: List[Document] = _rrf_merge(vector_lists + [bm25_results])[:6]

    # Optional cross-encoder reranking (only if RERANK_MODEL is configured)
    documents = _rerank_documents(question, documents)

    sources = [os.path.basename(d.metadata.get("source", "?")) for d in documents]
    logger.debug(f"[RETRIEVE] {len(documents)} doc(s) after hybrid+RRF (rerank={'on' if RERANK_MODEL else 'off'}): {sources}")
    return {"documents": documents, "question": question}


def grade_documents(state):
    question = state["question"]
    documents = state["documents"]
    logger.debug(f"[GRADE] grading {len(documents)} doc(s) against question: {question!r}")

    if not documents:
        logger.debug("[GRADE] no documents to grade")
        return {"documents": [], "question": question}

    batch_inputs = [{"question": question, "document": d.page_content} for d in documents]
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(batch_inputs)) as pool:
        scores = list(pool.map(grader_chain.invoke, batch_inputs))

    filtered_docs = []
    for i, score in enumerate(scores):
        source = os.path.basename(documents[i].metadata.get("source", "?"))
        verdict = score.binary_score
        logger.debug(f"[GRADE] doc {i + 1} ({source}): {verdict.upper()}")
        if verdict == "yes":
            filtered_docs.append(documents[i])

    logger.debug(f"[GRADE] {len(filtered_docs)}/{len(documents)} docs passed")
    return {"documents": filtered_docs, "question": question}


def generate_rag(state):
    question = state["question"]
    documents = state["documents"]
    history: list = state.get("history") or []

    # Build context — expand each chunk with its neighbors for richer LLM input (item 2)
    formatted_context_list = []
    for i, doc in enumerate(documents):
        source_name = os.path.basename(doc.metadata.get("source", "Unknown"))
        location = _citation_location(doc.metadata)
        loc_str = f" ({location})" if location else ""
        expanded_text = _expand_doc_context(doc)
        formatted_context_list.append(f"[{i + 1}] {source_name}{loc_str}\n{expanded_text}")

    full_context_string = "\n\n---\n\n".join(formatted_context_list)

    # Prepend conversation history so the model can refer to prior turns (item 6)
    history_prefix = ""
    if history:
        lines = [f"{t.get('role', 'user').capitalize()}: {t.get('content', '')}" for t in history]
        history_prefix = "Previous conversation:\n" + "\n".join(lines) + "\n\n"
    context_with_history = history_prefix + full_context_string

    sq: Optional[queue.Queue] = state.get("stream_queue")

    logger.debug(
        f"[GENERATE] → {THINKING_OLLAMA_MODEL}"
        f"\n  question : {question!r}"
        f"\n  history  : {len(history)} turn(s)"
        f"\n  context  : {len(full_context_string)} chars across {len(documents)} chunk(s)"
        f"\n{full_context_string}"
    )

    generation = ""
    for chunk in rag_chain.stream({"context": context_with_history, "question": question}):
        generation += chunk
        if sq is not None:
            sq.put(chunk)

    # Faithfulness check — warn if any claims are unsupported by the retrieved context (item 3)
    faithfulness_warning = ""
    if documents:
        try:
            result = faithfulness_check_chain.invoke({"context": full_context_string, "answer": generation})
            if not result.is_grounded:
                faithfulness_warning = (
                    "\n\n⚠️ *Note: Some claims in this answer may not be fully supported by the source material.*"
                )
        except Exception:
            pass  # faithfulness check is best-effort; never block the response

    if faithfulness_warning:
        generation += faithfulness_warning
        if sq is not None:
            sq.put(faithfulness_warning)

    # Append a Sources section using page → line priority
    if documents:
        source_lines = []
        for i, doc in enumerate(documents):
            meta = doc.metadata
            fname = os.path.basename(meta.get("source", "Unknown"))
            loc = _citation_location(meta)
            source_lines.append(f"[{i + 1}] {fname}" + (f" — {loc}" if loc else ""))
        footer = "\n\n---\n**Sources:**\n" + "\n".join(source_lines)
        generation += footer
        if sq is not None:
            sq.put(footer)

    if sq is not None:
        sq.put(None)  # sentinel — signals streaming is complete

    logger.debug(f"[GENERATE] ← response:\n{generation}")
    return {"generation": generation, "documents": documents}


def transform_query(state):
    question = state["question"]
    documents = state["documents"]
    loop_step = state.get("loop_step", 0)
    logger.debug(f"[REWRITE] → {FAST_OLLAMA_MODEL} | attempt {loop_step + 1} | input: {question!r}")
    better_question = rewriter_chain.invoke({"question": question})
    logger.debug(f"[REWRITE] ← output: {better_question!r}")
    return {"documents": documents, "question": better_question, "loop_step": loop_step + 1}


def decide_to_generate(state):
    filtered_documents = state["documents"]
    loop_step = state.get("loop_step", 0)

    if not filtered_documents:
        if loop_step >= 3:
            return "generate"
        return "transform_query"
    return "generate"


# --- PART 4: WORKFLOW BUILD ---

workflow = StateGraph(GraphState)
workflow.add_node("retrieve", retrieve)
workflow.add_node("grade_documents", grade_documents)
workflow.add_node("generate_rag", generate_rag)
workflow.add_node("transform_query", transform_query)

workflow.add_edge(START, "retrieve")
workflow.add_edge("retrieve", "grade_documents")
workflow.add_conditional_edges(
    "grade_documents",
    decide_to_generate,
    {"transform_query": "transform_query", "generate": "generate_rag"},
)
workflow.add_edge("transform_query", "retrieve")
workflow.add_edge("generate_rag", END)

app = workflow.compile()


# --- ENTRY POINTS ---

_CACHE_TTL_SECONDS = 3600  # cached answers expire after 1 hour
_CACHE_DB_PATH = os.path.join(CHROMA_DB_PATH, "query_cache.db")
_CACHE_CONN: Optional["sqlite3.Connection"] = None
_CACHE_CONN_LOCK = threading.Lock()


def _get_cache_conn():
    """Return the shared SQLite cache connection, creating it on first call."""
    global _CACHE_CONN
    if _CACHE_CONN is None:
        os.makedirs(CHROMA_DB_PATH, exist_ok=True)
        _CACHE_CONN = sqlite3.connect(_CACHE_DB_PATH, check_same_thread=False)
        _CACHE_CONN.execute(
            "CREATE TABLE IF NOT EXISTS query_cache "
            "(key TEXT PRIMARY KEY, result TEXT NOT NULL, created_at REAL NOT NULL)"
        )
        _CACHE_CONN.commit()
    return _CACHE_CONN


def _cache_clear() -> None:
    """Remove all entries from the persistent query cache. Used in tests and admin ops."""
    try:
        with _CACHE_CONN_LOCK:
            conn = _get_cache_conn()
            conn.execute("DELETE FROM query_cache")
            conn.commit()
    except Exception as e:
        logger.warning(f"[CACHE] Clear error: {e}")


def _manifest_hash() -> str:
    """Return a short hash of the current index manifest.

    Used as part of the cache key so that cached answers are automatically
    invalidated whenever documents are added, removed, or updated.
    """
    raw = json.dumps(_load_index_manifest(), sort_keys=True)
    return hashlib.md5(raw.encode()).hexdigest()


def _make_cache_key(
    query: str,
    scoped_source: Optional[str],
    tag_filter: Optional[str],
    history: Optional[list],
    include_sources: bool,
) -> str:
    parts = json.dumps(
        {"q": query, "s": scoped_source, "t": tag_filter,
         "h": history, "m": _manifest_hash(), "i": include_sources},
        sort_keys=True,
    )
    return hashlib.md5(parts.encode()).hexdigest()


def _cache_get(key: str):
    try:
        with _CACHE_CONN_LOCK:
            row = _get_cache_conn().execute(
                "SELECT result, created_at FROM query_cache WHERE key = ?", (key,)
            ).fetchone()
        if row is None:
            return None
        result_json, ts = row
        if time.time() - ts > _CACHE_TTL_SECONDS:
            with _CACHE_CONN_LOCK:
                conn = _get_cache_conn()
                conn.execute("DELETE FROM query_cache WHERE key = ?", (key,))
                conn.commit()
            return None
        return json.loads(result_json)
    except Exception as e:
        logger.warning(f"[CACHE] Read error: {e}")
        return None


def _cache_put(key: str, result) -> None:
    try:
        result_json = json.dumps(result)
        with _CACHE_CONN_LOCK:
            conn = _get_cache_conn()
            conn.execute(
                "INSERT OR REPLACE INTO query_cache (key, result, created_at) VALUES (?, ?, ?)",
                (key, result_json, time.time()),
            )
            conn.commit()
    except Exception as e:
        logger.warning(f"[CACHE] Write error: {e}")


def _run_rag_graph(inputs: dict, include_sources: bool):
    """Execute the compiled RAG graph and return a formatted answer."""
    config = {"recursion_limit": 25}
    final_state = None
    generate_rag_state = None

    try:
        for output in app.stream(inputs, config=config):
            for key, value in output.items():
                final_state = value
                if key == "generate_rag":
                    generate_rag_state = value
    except Exception as e:
        logger.error(f"Graph execution error: {e}")
        return f"Error: {e}" if not include_sources else {"error": str(e)}

    result_state = generate_rag_state if generate_rag_state is not None else final_state
    if result_state is None:
        no_answer = "No answer generated."
        return no_answer if not include_sources else {"answer": no_answer, "citations": []}

    final_answer = result_state.get("generation", "No answer generated.")
    used_docs = result_state.get("documents", [])

    sources_data = []
    for doc in used_docs:
        meta = doc.metadata
        fname = os.path.basename(meta.get("source", "Unknown"))
        location = _format_location(meta)
        sources_data.append({
            "file": fname,
            "location": location,
            "snippet": doc.page_content[:50] + "...",
        })

    if include_sources:
        return {"answer": final_answer, "citations": sources_data}
    return final_answer


def query_documents(user_input: str, include_sources: bool = False,
                    stream_queue: Optional[queue.Queue] = None,
                    tag_filter: Optional[str] = None,
                    history: Optional[list] = None):
    """Search all indexed documents and return an answer."""
    if GLOBAL_VECTORSTORE is None:
        initialize_vectorstore()

    key = _make_cache_key(user_input, None, tag_filter, history, include_sources)
    cached = _cache_get(key)
    if cached is not None:
        logger.debug(f"[CACHE] HIT: {user_input!r}")
        if stream_queue is not None:
            stream_queue.put(cached["answer"] if include_sources else cached)
            stream_queue.put(None)
        return cached

    inputs: dict = {"question": user_input, "loop_step": 0}
    if stream_queue is not None:
        inputs["stream_queue"] = stream_queue
    if tag_filter:
        inputs["tag_filter"] = tag_filter
    if history:
        inputs["history"] = history
    result = _run_rag_graph(inputs, include_sources)
    _cache_put(key, result)
    return result


def query_documents_scoped(user_input: str, filename: str, include_sources: bool = False,
                           stream_queue: Optional[queue.Queue] = None,
                           tag_filter: Optional[str] = None):
    """Search within a single named document and return an answer.

    Args:
        user_input: The user's question.
        filename: The basename of the indexed file to restrict the search to.
        include_sources: If True, return a dict with 'answer' and 'citations'.
        stream_queue: Optional queue that receives token chunks as they are generated.
        tag_filter: Optional tag string to further restrict the search.
    """
    if GLOBAL_VECTORSTORE is None:
        initialize_vectorstore()

    manifest = _load_index_manifest()
    matching = [p for p in manifest if os.path.basename(p) == filename]
    if not matching:
        msg = f"No indexed document named '{filename}' found. Use `/status` to see available documents."
        return {"answer": msg, "citations": []} if include_sources else msg

    key = _make_cache_key(user_input, matching[0], tag_filter, None, include_sources)
    cached = _cache_get(key)
    if cached is not None:
        logger.debug(f"[CACHE] HIT: {user_input!r} (scoped: {filename})")
        if stream_queue is not None:
            stream_queue.put(cached["answer"] if include_sources else cached)
            stream_queue.put(None)
        return cached

    inputs: dict = {"question": user_input, "loop_step": 0, "scoped_source": matching[0]}
    if stream_queue is not None:
        inputs["stream_queue"] = stream_queue
    if tag_filter:
        inputs["tag_filter"] = tag_filter
    result = _run_rag_graph(inputs, include_sources)
    _cache_put(key, result)
    return result


def similarity_search(
    query: str,
    k: int = 10,
    source_filter: Optional[str] = None,
    tag_filter: Optional[str] = None,
) -> list[tuple]:
    """Return top-k (Document, relevance_score) pairs without LLM generation.

    relevance_score is in [0, 1] — higher means more similar.
    source_filter, if given, must be the full file path stored in the manifest.
    Returns an empty list if the vectorstore is not ready or an error occurs.
    """
    if GLOBAL_VECTORSTORE is None:
        return []
    filter_dict = {"source": source_filter} if source_filter else None
    if tag_filter:
        tag_f = {"tags": {"$contains": tag_filter}}
        if filter_dict is not None:
            filter_dict = {"$and": [filter_dict, tag_f]}
        else:
            filter_dict = tag_f
    try:
        return GLOBAL_VECTORSTORE.similarity_search_with_relevance_scores(
            query, k=k, filter=filter_dict
        )
    except Exception as e:
        logger.error(f"Similarity search error: {e}")
        return []


def get_chunks_for_file(file_path: str) -> list[dict]:
    """Return all stored chunks for a file as a list of {content, metadata} dicts.

    Results are sorted by start_index when available, otherwise by page.
    Returns an empty list if the vectorstore is not ready or an error occurs.
    """
    if GLOBAL_VECTORSTORE is None:
        return []
    try:
        results = GLOBAL_VECTORSTORE.get(
            where={"source": file_path},
            include=["metadatas", "documents"],
        )
        chunks = [
            {"content": doc, "metadata": meta}
            for doc, meta in zip(
                results.get("documents", []),
                results.get("metadatas", []),
            )
        ]
        chunks.sort(key=lambda c: (
            c["metadata"].get("page", 0),
            c["metadata"].get("start_index", 0),
        ))
        return chunks
    except Exception as e:
        logger.error(f"Error fetching chunks for {file_path}: {e}")
        return []


if __name__ == "__main__":
    query_documents("Who is best friends with Dagbert?", True)