from __future__ import annotations

from pathlib import Path
import shutil
from typing import Any, Callable

from Mama.models import ExecutionContext
from Mama.ports import VectorStoreFactoryPort, VectorStorePort

try:
    from langchain_community.vectorstores import FAISS  # type: ignore
except Exception:  # pragma: no cover
    from langchain.vectorstores import FAISS  # type: ignore


class FaissStore(VectorStorePort):
    def __init__(self, embeddings: Any, index_path: Path | None = None):
        self._embeddings = embeddings
        self._index_path = index_path
        self._store = None

        if index_path and index_path.exists():
            self._store = FAISS.load_local(
                str(index_path),
                embeddings=embeddings,
                allow_dangerous_deserialization=True,
            )

    def _save(self) -> None:
        if self._index_path and self._store:
            self._index_path.mkdir(parents=True, exist_ok=True)
            self._store.save_local(str(self._index_path))

    def _clear_index(self) -> None:
        if self._index_path and self._index_path.exists():
            shutil.rmtree(self._index_path, ignore_errors=True)

    def add_documents(self, documents: list[Any]) -> int:
        if not documents:
            return 0
        if self._store is None:
            self._store = FAISS.from_documents(documents=documents, embedding=self._embeddings)
        else:
            self._store.add_documents(documents=documents)
        self._save()
        return len(documents)

    def similarity_search(
        self,
        query: str,
        *,
        k: int,
        search_type: str,
        filters: dict[str, Any] | None = None,
    ) -> list[Any]:
        if self._store is None:
            return []
        # FAISS does not natively support metadata filtering in this adapter.
        # We fetch a wider candidate set, then apply metadata post-filter.
        search_k = max(20, int(k) * 8) if filters else max(1, int(k))
        retriever = self._store.as_retriever(
            search_type=search_type or "similarity",
            search_kwargs={"k": search_k},
        )
        try:
            docs = retriever.invoke(query)
        except Exception:
            legacy_getter = getattr(retriever, "get_relevant_documents", None)
            if callable(legacy_getter):
                docs = legacy_getter(query)  # type: ignore[misc]
            else:
                docs = []

        if filters:
            docs = [doc for doc in docs if _matches_filters(doc.metadata or {}, filters)]
        return docs[: max(1, int(k))]

    def list_documents(self, limit: int | None = None) -> list[dict[str, Any]]:
        if self._store is None:
            return []
        out: list[dict[str, Any]] = []
        for doc_id, doc in self._store.docstore._dict.items():  # type: ignore[attr-defined]
            out.append(
                {
                    "id": doc_id,
                    "source": doc.metadata.get("source", ""),
                    "page_content": doc.page_content,
                    "metadata": dict(doc.metadata),
                }
            )
            if limit and len(out) >= limit:
                break
        return out

    def delete_documents(self, filters: dict[str, Any]) -> int:
        if self._store is None:
            return 0
        if not filters:
            return 0

        existing_items = list(self._store.docstore._dict.items())  # type: ignore[attr-defined]
        to_remove_ids = [
            doc_id
            for doc_id, doc in existing_items
            if _matches_filters(doc.metadata or {}, filters)
        ]
        if not to_remove_ids:
            return 0

        to_remove_set = set(to_remove_ids)
        remaining_docs = [
            doc
            for doc_id, doc in existing_items
            if doc_id not in to_remove_set
        ]

        if remaining_docs:
            self._store = FAISS.from_documents(documents=remaining_docs, embedding=self._embeddings)
            self._save()
        else:
            self._store = None
            self._clear_index()
        return len(to_remove_ids)

    def stats(self) -> dict[str, Any]:
        if self._store is None:
            return {"documents": 0, "vectors": 0, "dimension": 0, "disk_size": 0}

        size = 0
        if self._index_path and self._index_path.exists():
            for child in self._index_path.glob("**/*"):
                if child.is_file():
                    size += child.stat().st_size

        faiss_index = getattr(self._store, "_faiss_index", None) or getattr(self._store, "index", None)
        dim = getattr(faiss_index, "d", 0) if faiss_index is not None else 0

        return {
            "documents": len(self._store.docstore._dict),  # type: ignore[attr-defined]
            "vectors": len(self._store.index_to_docstore_id),
            "dimension": dim,
            "disk_size": size,
        }

    def identifier(self) -> str:
        if self._index_path:
            return str(self._index_path)
        return "faiss://memory"


def _build_faiss_local(ctx: ExecutionContext, embeddings: Any) -> VectorStorePort:
    params = ctx.storage.params
    index_path = params.get("index_path")
    if not index_path:
        base_path = params.get("base_path")
        if not base_path:
            raise ValueError("faiss_local requires storage.params.index_path or storage.params.base_path")
        tenant = ctx.tenant_id or "default"
        index_path = str(Path(base_path) / tenant / ctx.kb_id)

    return FaissStore(embeddings=embeddings, index_path=Path(index_path))


def _build_faiss_memory(_: ExecutionContext, embeddings: Any) -> VectorStorePort:
    return FaissStore(embeddings=embeddings, index_path=None)


class DynamicVectorStoreFactory(VectorStoreFactoryPort):
    def __init__(self):
        self._builders: dict[str, Callable[[ExecutionContext, Any], VectorStorePort]] = {
            "faiss_local": _build_faiss_local,
            "faiss_memory": _build_faiss_memory,
        }

    def register_backend(
        self,
        backend: str,
        builder: Callable[[ExecutionContext, Any], VectorStorePort],
    ) -> None:
        self._builders[backend] = builder

    def get_store(self, ctx: ExecutionContext, embeddings: Any) -> VectorStorePort:
        backend = (ctx.storage.backend or "").strip().lower()
        if not backend:
            raise ValueError("storage backend is required")
        builder = self._builders.get(backend)
        if builder is None:
            supported = ", ".join(sorted(self._builders.keys()))
            raise ValueError(f"Unsupported vector store backend '{backend}'. Supported: {supported}")
        return builder(ctx, embeddings)


def _matches_filters(metadata: dict[str, Any], filters: dict[str, Any]) -> bool:
    for key, expected in filters.items():
        if expected in (None, "", []):
            continue
        actual = metadata.get(key)
        if key == "tags":
            if not _tags_match(actual, expected):
                return False
            continue
        if not _value_matches(actual, expected):
            return False
    return True


def _value_matches(actual: Any, expected: Any) -> bool:
    # list filter: match if any expected value matches the actual field
    if isinstance(expected, list):
        return any(_value_matches(actual, item) for item in expected)

    # string filter: case-insensitive "contains" semantics
    if isinstance(expected, str):
        expected_norm = expected.strip().lower()
        if not expected_norm:
            return True
        if actual is None:
            return False
        return expected_norm in str(actual).strip().lower()

    # scalar fallback: strict equality
    return actual == expected


def _tags_match(actual: Any, expected: Any) -> bool:
    if isinstance(expected, str):
        expected_tags = {expected.strip().lower()} if expected.strip() else set()
    else:
        expected_tags = {str(item).strip().lower() for item in (expected or []) if str(item).strip()}

    if not expected_tags:
        return True

    if isinstance(actual, str):
        actual_tags = {item.strip().lower() for item in actual.split(",") if item.strip()}
    else:
        actual_tags = {str(item).strip().lower() for item in (actual or []) if str(item).strip()}

    return expected_tags.issubset(actual_tags)
