"""Safe loading of on-disk FAISS indexes.

LangChain's ``FAISS.load_local(..., allow_dangerous_deserialization=True)``
deserializes the ``<index>.pkl`` docstore companion with a raw
``pickle.load`` (see ``langchain_community/vectorstores/faiss.py``). Pickle
executes arbitrary code during deserialization, so a tampered ``.pkl``
dropped into the index directory yields arbitrary code execution when the
index is next loaded.

The ``.faiss`` and ``.pkl`` files are loaded independently with no
cross-check, so a content checksum over only the ``.faiss`` file (which is
what the file-integrity system records) provides **no** protection against a
``.pkl``-only swap. Rather than gate the dangerous load behind a checksum,
this module removes the dangerous deserialization entirely: it replicates
``load_local`` but unpickles the docstore through a restricted unpickler that
only permits the two classes a legitimate FAISS docstore contains
(:class:`InMemoryDocstore` and :class:`Document`). Any other global
(``os.system``, ``subprocess.Popen``, a crafted ``__reduce__`` payload, …)
raises :class:`pickle.UnpicklingError` before it can execute.

The payload is the 2-tuple ``(docstore, index_to_docstore_id)`` written by
``FAISS.save_local``; the metadata it carries is plain JSON-ish scalars
(str/int/float/bool/None/list/dict), none of which require a class global to
unpickle. If a future LangChain/Pydantic version changes the on-disk format
to need additional classes, loading fails closed (a clear error) rather than
silently — and the round-trip test in
``tests/research_library/services/test_faiss_safe_load.py`` will catch it so
the allow-list can be widened deliberately.
"""

import pickle
from pathlib import Path
from typing import Any

from langchain_community.vectorstores import FAISS
from langchain_community.vectorstores.faiss import dependable_faiss_import
from langchain_core.embeddings import Embeddings

# (module, qualname) pairs that a legitimate FAISS docstore pickle resolves
# via ``Unpickler.find_class``. Determined empirically by round-tripping a
# real ``save_local`` payload with varied metadata; only these two appear.
_ALLOWED_GLOBALS = frozenset(
    {
        ("langchain_community.docstore.in_memory", "InMemoryDocstore"),
        ("langchain_core.documents.base", "Document"),
    }
)


# Subclass the PURE-PYTHON unpickler (``pickle._Unpickler``), NOT the default
# C ``pickle.Unpickler``. We need to override ``get_extension`` to refuse the
# copyreg extension opcodes (EXT1/EXT2/EXT4), and only the pure-Python unpickler
# routes those opcodes through an overridable ``get_extension`` method — the C
# unpickler resolves them internally. (Why ``get_extension`` and not just
# ``find_class``: both unpicklers normally route EXT through ``find_class`` on a
# cold cache, but ``get_extension`` short-circuits on a warm process-global
# ``copyreg._extension_cache`` and skips ``find_class`` entirely. Refusing
# ``get_extension`` outright closes that path regardless of cache/registry
# state, instead of relying on the extension registry happening to be empty.)
# ``_Unpickler`` has been the stable pure-Python name since Python 3.0; if a
# future runtime drops it, this fails loudly at import (caught by tests) rather
# than silently falling back to the C unpickler.
class _RestrictedFaissUnpickler(pickle._Unpickler):
    """Unpickler that only resolves the classes a FAISS docstore contains.

    Defense is by construction: a malicious pickle can never obtain an
    attacker-chosen callable to invoke, because every opcode that resolves a
    global or callable is refused or allow-listed:

    - ``find_class`` (GLOBAL/STACK_GLOBAL/INST/OBJ, then REDUCE/NEWOBJ/BUILD)
      allow-lists only ``InMemoryDocstore`` and ``Document`` — both inert.
    - ``get_extension`` (copyreg EXT1/EXT2/EXT4) is refused outright; a
      legitimate FAISS docstore pickle never uses copyreg extensions.
    - ``persistent_load`` (PERSID/BINPERSID) is refused outright; ``save_local``
      never emits persistent ids.
    """

    def find_class(self, module: str, name: str) -> Any:
        if (module, name) in _ALLOWED_GLOBALS:
            return super().find_class(module, name)
        raise pickle.UnpicklingError(
            f"Refusing to unpickle disallowed object '{module}.{name}' from "
            f"FAISS docstore (possible tampering)."
        )

    def get_extension(self, code: int) -> Any:
        raise pickle.UnpicklingError(
            f"Refusing to resolve copyreg extension code {code} from FAISS "
            f"docstore (possible tampering)."
        )

    def persistent_load(self, pid: Any) -> Any:
        raise pickle.UnpicklingError(
            "Refusing to resolve a persistent id from FAISS docstore "
            "(possible tampering)."
        )


def safe_load_faiss(
    folder_path: str,
    embeddings: Embeddings,
    *,
    index_name: str,
    normalize_L2: bool = False,
) -> FAISS:
    """Load a FAISS index from disk without dangerous deserialization.

    Drop-in replacement for ``FAISS.load_local(...,
    allow_dangerous_deserialization=True)`` for the index/docstore layout
    this project writes via ``FAISS.save_local``.

    Args:
        folder_path: Directory containing ``<index_name>.faiss`` and
            ``<index_name>.pkl``.
        embeddings: Embeddings used for subsequent queries.
        index_name: Base name of the index files (the ``.faiss`` stem).
        normalize_L2: Whether vectors are L2-normalized (must match how the
            index was built).

    Returns:
        A :class:`FAISS` vector store.

    Raises:
        pickle.UnpicklingError: If the ``.pkl`` references any class outside
            the allow-list (tampering or an unexpected format change).
    """
    path = Path(folder_path)

    # The .faiss file is a binary FAISS index (not pickle); read it as usual.
    faiss = dependable_faiss_import()
    index = faiss.read_index(str(path / f"{index_name}.faiss"))

    # The .pkl is the only attack surface — unpickle it under restriction.
    with open(path / f"{index_name}.pkl", "rb") as f:
        docstore, index_to_docstore_id = _RestrictedFaissUnpickler(f).load()

    return FAISS(
        embeddings,
        index,
        docstore,
        index_to_docstore_id,
        normalize_L2=normalize_L2,
    )
