"""VectorBackend protocol and implementations for filtered vector search.

FAISS has no metadata filtering. These backends support pre-filtering by
type/date/status before vector search.

Implementations:
- QdrantBackend: qdrant-client based (pip install core-memory[qdrant])
- ChromaDBBackend: chromadb based (pip install core-memory[chromadb])
- PgvectorBackend: psycopg + pgvector (pip install core-memory[pgvector])
"""
from __future__ import annotations

import json
import logging
from typing import Any, Protocol, runtime_checkable

logger = logging.getLogger(__name__)


@runtime_checkable
class VectorBackend(Protocol):
    """Protocol for vector search backends with metadata filtering."""

    def upsert(self, bead_id: str, embedding: list[float], metadata: dict[str, Any]) -> None:
        """Insert or update a bead's embedding with metadata."""
        ...

    def search(
        self,
        query_embedding: list[float],
        k: int = 8,
        filters: dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """Search for similar beads with optional metadata filters.

        Filters support: type, status, session_id, created_after, created_before.
        Returns list of {"bead_id": str, "score": float, "metadata": dict}.
        """
        ...

    def delete(self, bead_id: str) -> None:
        """Delete a bead's embedding."""
        ...

    def update_metadata(self, bead_id: str, metadata: dict[str, Any]) -> None:
        """Update metadata for an existing embedding without changing its vector."""
        ...

    def count(self) -> int:
        """Return the number of indexed embeddings."""
        ...


def _bead_id_to_qdrant_id(bead_id: str) -> str:
    """Map an arbitrary bead ID to a UUID5. Deterministic and collision-resistant.

    Qdrant embedded mode (local QdrantClient) requires UUID or integer point IDs.
    Server mode accepts arbitrary strings, but we use UUID5 everywhere for consistency.
    """
    import uuid as _uuid
    return str(_uuid.uuid5(_uuid.UUID("6ba7b810-9dad-11d1-80b4-00c04fd430c8"), str(bead_id)))


class QdrantBackend:
    """Qdrant vector search backend.

    Embedded mode (zero-ops): pass path=".beads/qdrant" — no server required.
    Server mode: pass url="http://localhost:6333" or set CORE_MEMORY_QDRANT_URL.

    Requires: pip install core-memory[qdrant]
    """

    def __init__(
        self,
        collection_name: str = "core_memory_beads",
        url: str | None = None,
        path: str | None = None,
        dimensions: int = 1536,
    ):
        try:
            from qdrant_client import QdrantClient
            from qdrant_client.models import Distance, VectorParams
        except ImportError:
            raise ImportError("Qdrant backend requires: pip install core-memory[qdrant]")

        if url:
            self._client = QdrantClient(url=url)
        elif path:
            import os
            os.makedirs(path, exist_ok=True)
            self._client = QdrantClient(path=path)
        else:
            self._client = QdrantClient(":memory:")

        self._collection = collection_name

        # Ensure collection exists for regular vector mode.
        # Skip when dimensions=0: caller will use upsert_texts()/client.add() which
        # creates a fastembed-configured collection — pre-creating with VectorParams
        # would produce an incompatible collection type.
        if int(dimensions) > 0:
            expected_dim = int(dimensions)
            collections = [c.name for c in self._client.get_collections().collections]
            collection_exists = collection_name in collections
            if collection_exists:
                existing_dim = self._collection_vector_size(collection_name)
                # Qdrant collections have fixed vector params. If this collection
                # was created by FastEmbed/client.add(), or by a different external
                # model dimension, recreate it before upserting provider vectors so
                # mode/model switches do not fail until a manual collection drop.
                if existing_dim != expected_dim:
                    logger.warning(
                        "core-memory: recreating incompatible Qdrant collection %s (existing_dim=%s, expected_dim=%s)",
                        collection_name,
                        existing_dim,
                        expected_dim,
                    )
                    self._client.delete_collection(collection_name=collection_name)
                    collection_exists = False
            if not collection_exists:
                self._client.create_collection(
                    collection_name=collection_name,
                    vectors_config=VectorParams(size=expected_dim, distance=Distance.COSINE),
                )

    def _collection_vector_size(self, collection_name: str) -> int | None:
        try:
            info = self._client.get_collection(collection_name)
            vectors = getattr(getattr(getattr(info, "config", None), "params", None), "vectors", None)
            if isinstance(vectors, dict):
                sizes = [int(getattr(v, "size", 0) or 0) for v in vectors.values()]
                return sizes[0] if len(sizes) == 1 and sizes[0] > 0 else None
            size = int(getattr(vectors, "size", 0) or 0)
            return size if size > 0 else None
        except Exception:
            return None

    def upsert(self, bead_id: str, embedding: list[float], metadata: dict[str, Any]) -> None:
        from qdrant_client.models import PointStruct
        # Store original bead_id in payload for retrieval; use UUID5 as point ID.
        point_id = _bead_id_to_qdrant_id(bead_id)
        payload = {**metadata, "bead_id": bead_id}
        self._client.upsert(
            collection_name=self._collection,
            points=[PointStruct(id=point_id, vector=embedding, payload=payload)],
        )

    def upsert_texts(self, bead_ids: list[str], texts: list[str], metadatas: list[dict[str, Any]]) -> None:
        """Batch upsert using FastEmbed native indexing — no external embedding API required."""
        ids = [_bead_id_to_qdrant_id(bid) for bid in bead_ids]
        payloads = [{**meta, "bead_id": bid} for bid, meta in zip(bead_ids, metadatas)]
        self._client.add(
            collection_name=self._collection,
            documents=texts,
            metadata=payloads,
            ids=ids,
        )

    def search(
        self,
        query_embedding: list[float],
        k: int = 8,
        filters: dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        from qdrant_client.models import FieldCondition, Filter, MatchValue, Range

        qdrant_filter = None
        if filters:
            conditions = []
            if "type" in filters:
                conditions.append(FieldCondition(key="type", match=MatchValue(value=filters["type"])))
            if "status" in filters:
                conditions.append(FieldCondition(key="status", match=MatchValue(value=filters["status"])))
            if "session_id" in filters:
                conditions.append(FieldCondition(key="session_id", match=MatchValue(value=filters["session_id"])))
            if "retrieval_eligible" in filters:
                conditions.append(FieldCondition(key="retrieval_eligible", match=MatchValue(value=bool(filters["retrieval_eligible"]))))
            if "created_after" in filters:
                conditions.append(FieldCondition(key="created_at", range=Range(gte=filters["created_after"])))
            if "created_before" in filters:
                conditions.append(FieldCondition(key="created_at", range=Range(lte=filters["created_before"])))
            if conditions:
                qdrant_filter = Filter(must=conditions)

        response = self._client.query_points(
            collection_name=self._collection,
            query=query_embedding,
            limit=k,
            query_filter=qdrant_filter,
            with_payload=True,
        )
        return [
            # Recover original bead_id from payload (UUID5 point ID is internal)
            {"bead_id": str((r.payload or {}).get("bead_id") or r.id), "score": float(r.score), "metadata": dict(r.payload or {})}
            for r in response.points
        ]

    def delete(self, bead_id: str) -> None:
        from qdrant_client.models import PointIdsList
        point_id = _bead_id_to_qdrant_id(bead_id)
        self._client.delete(collection_name=self._collection, points_selector=PointIdsList(points=[point_id]))

    def update_metadata(self, bead_id: str, metadata: dict[str, Any]) -> None:
        point_id = _bead_id_to_qdrant_id(bead_id)
        self._client.set_payload(
            collection_name=self._collection,
            payload=dict(metadata or {}),
            points=[point_id],
        )

    def count(self) -> int:
        info = self._client.get_collection(self._collection)
        return info.points_count or 0

    def hybrid_search(
        self,
        query: str,
        k: int = 8,
        filters: dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """Sparse+dense hybrid search using FastEmbed. One Qdrant query replaces FAISS + lexical.py."""
        try:
            from qdrant_client.models import FieldCondition, Filter, MatchValue, Prefetch, Query
        except ImportError:
            raise ImportError("Qdrant hybrid search requires: pip install 'qdrant-client[fastembed]'")

        must_conditions = []
        if filters:
            for key in ("type", "status", "session_id"):
                if key in filters:
                    must_conditions.append(FieldCondition(key=key, match=MatchValue(value=filters[key])))
            if "retrieval_eligible" in filters:
                must_conditions.append(
                    FieldCondition(key="retrieval_eligible", match=MatchValue(value=bool(filters["retrieval_eligible"])))
                )

        qdrant_filter = Filter(must=must_conditions) if must_conditions else None

        results = self._client.query(
            collection_name=self._collection,
            query_text=query,
            query_filter=qdrant_filter,
            limit=k,
        )
        return [
            {"bead_id": str((r.payload or {}).get("bead_id") or r.id), "score": float(r.score), "metadata": dict(r.payload or {})}
            for r in results
        ]

    def retrieve_by_ids(self, bead_ids: list[str]) -> list[dict[str, Any]]:
        """Point lookup by bead ID — O(1) per batch."""
        if not bead_ids:
            return []
        point_ids = [_bead_id_to_qdrant_id(bid) for bid in bead_ids]
        points = self._client.retrieve(collection_name=self._collection, ids=point_ids, with_payload=True)
        return [
            {"bead_id": str((p.payload or {}).get("bead_id") or p.id), "metadata": dict(p.payload or {})}
            for p in points
        ]


class ChromaDBBackend:
    """ChromaDB vector search backend.

    Requires: pip install core-memory[chromadb]
    """

    def __init__(self, collection_name: str = "core_memory_beads", persist_directory: str | None = None):
        try:
            import chromadb
        except ImportError:
            raise ImportError("ChromaDB backend requires: pip install core-memory[chromadb]")

        if persist_directory:
            self._client = chromadb.PersistentClient(path=persist_directory)
        else:
            self._client = chromadb.Client()
        self._collection = self._client.get_or_create_collection(name=collection_name)

    def upsert(self, bead_id: str, embedding: list[float], metadata: dict[str, Any]) -> None:
        # ChromaDB requires metadata values to be str/int/float/bool
        clean_meta = {}
        for k, v in metadata.items():
            if isinstance(v, (str, int, float, bool)):
                clean_meta[k] = v
            elif v is not None:
                clean_meta[k] = str(v)
        self._collection.upsert(ids=[bead_id], embeddings=[embedding], metadatas=[clean_meta])

    def search(
        self,
        query_embedding: list[float],
        k: int = 8,
        filters: dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        where_filter = None
        if filters:
            conditions = {}
            for key in ("type", "status", "session_id"):
                if key in filters:
                    conditions[key] = filters[key]
            if conditions:
                if len(conditions) == 1:
                    key, val = next(iter(conditions.items()))
                    where_filter = {key: val}
                else:
                    where_filter = {"$and": [{k: v} for k, v in conditions.items()]}

        results = self._collection.query(
            query_embeddings=[query_embedding],
            n_results=k,
            where=where_filter,
        )
        ids = results.get("ids", [[]])[0]
        distances = results.get("distances", [[]])[0]
        metadatas = results.get("metadatas", [[]])[0]

        return [
            {"bead_id": bid, "score": 1.0 - dist, "metadata": dict(meta or {})}
            for bid, dist, meta in zip(ids, distances, metadatas)
        ]

    def delete(self, bead_id: str) -> None:
        self._collection.delete(ids=[bead_id])

    def update_metadata(self, bead_id: str, metadata: dict[str, Any]) -> None:
        clean_meta = {}
        for k, v in (metadata or {}).items():
            if isinstance(v, (str, int, float, bool)):
                clean_meta[k] = v
            elif v is not None:
                clean_meta[k] = str(v)
        self._collection.update(ids=[bead_id], metadatas=[clean_meta])

    def count(self) -> int:
        return self._collection.count()


class PgvectorBackend:
    """PostgreSQL + pgvector backend for filtered vector search.

    Requires: pip install core-memory[pgvector]
    (psycopg[binary] + pgvector extension on the database)

    Connection via CORE_MEMORY_PG_DSN env var or dsn parameter.
    """

    def __init__(self, dsn: str | None = None, table_name: str = "core_memory_beads", dimensions: int = 1536):
        import os
        try:
            import psycopg
        except ImportError:
            raise ImportError("pgvector backend requires: pip install core-memory[pgvector]")

        self._dsn = dsn or os.environ.get("CORE_MEMORY_PG_DSN", "")
        if not self._dsn:
            raise ValueError("PostgreSQL DSN required: set CORE_MEMORY_PG_DSN or pass dsn= parameter")
        self._table = table_name
        self._dim = dimensions
        self._conn = psycopg.connect(self._dsn, autocommit=True)

        # Ensure pgvector extension and table exist
        self._conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
        self._conn.execute(f"""
            CREATE TABLE IF NOT EXISTS {self._table} (
                bead_id TEXT PRIMARY KEY,
                embedding vector({self._dim}),
                type TEXT,
                status TEXT,
                session_id TEXT,
                created_at TEXT,
                metadata JSONB DEFAULT '{{}}'::jsonb
            )
        """)
        self._conn.execute(f"CREATE INDEX IF NOT EXISTS idx_{self._table}_type ON {self._table}(type)")
        self._conn.execute(f"CREATE INDEX IF NOT EXISTS idx_{self._table}_status ON {self._table}(status)")

    def upsert(self, bead_id: str, embedding: list[float], metadata: dict[str, Any]) -> None:
        self._conn.execute(
            f"""
            INSERT INTO {self._table} (bead_id, embedding, type, status, session_id, created_at, metadata)
            VALUES (%s, %s::vector, %s, %s, %s, %s, %s::jsonb)
            ON CONFLICT (bead_id) DO UPDATE SET
                embedding = EXCLUDED.embedding,
                type = EXCLUDED.type,
                status = EXCLUDED.status,
                session_id = EXCLUDED.session_id,
                created_at = EXCLUDED.created_at,
                metadata = EXCLUDED.metadata
            """,
            (
                bead_id,
                str(embedding),
                metadata.get("type"),
                metadata.get("status"),
                metadata.get("session_id"),
                metadata.get("created_at"),
                json.dumps(metadata),
            ),
        )

    def search(
        self,
        query_embedding: list[float],
        k: int = 8,
        filters: dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        conditions = []
        where_params: list[Any] = []
        query_vec = str(query_embedding)

        if filters:
            if "type" in filters:
                conditions.append("type = %s")
                where_params.append(filters["type"])
            if "status" in filters:
                conditions.append("status = %s")
                where_params.append(filters["status"])
            if "session_id" in filters:
                conditions.append("session_id = %s")
                where_params.append(filters["session_id"])
            if "created_after" in filters:
                conditions.append("created_at >= %s")
                where_params.append(filters["created_after"])
            if "created_before" in filters:
                conditions.append("created_at <= %s")
                where_params.append(filters["created_before"])

        where = (" WHERE " + " AND ".join(conditions)) if conditions else ""
        params: list[Any] = [query_vec] + where_params + [query_vec, k]

        cur = self._conn.execute(
            f"""
            SELECT bead_id, 1 - (embedding <=> %s::vector) AS score, metadata
            FROM {self._table}{where}
            ORDER BY embedding <=> %s::vector
            LIMIT %s
            """,
            params,
        )

        results = []
        for row in cur.fetchall():
            results.append({
                "bead_id": row[0],
                "score": float(row[1]),
                "metadata": json.loads(row[2]) if isinstance(row[2], str) else dict(row[2] or {}),
            })
        return results

    def delete(self, bead_id: str) -> None:
        self._conn.execute(f"DELETE FROM {self._table} WHERE bead_id = %s", (bead_id,))

    def update_metadata(self, bead_id: str, metadata: dict[str, Any]) -> None:
        md = dict(metadata or {})
        self._conn.execute(
            f"""
            UPDATE {self._table}
            SET type = %s,
                status = %s,
                session_id = %s,
                created_at = %s,
                metadata = %s::jsonb
            WHERE bead_id = %s
            """,
            (
                md.get("type"),
                md.get("status"),
                md.get("session_id"),
                md.get("created_at"),
                json.dumps(md),
                bead_id,
            ),
        )

    def count(self) -> int:
        cur = self._conn.execute(f"SELECT COUNT(*) FROM {self._table}")
        return cur.fetchone()[0]


def create_vector_backend(backend_type: str = "qdrant", **kwargs: Any) -> VectorBackend:
    """Factory for creating vector search backends.

    Args:
        backend_type: "qdrant", "chromadb", or "pgvector"
        **kwargs: Passed to the backend constructor. For qdrant, pass url= for server
                  mode or path= for embedded mode (zero-ops, recommended default).
    """
    if backend_type == "qdrant":
        return QdrantBackend(**kwargs)
    if backend_type in ("chromadb", "chroma"):
        return ChromaDBBackend(**kwargs)
    if backend_type in ("pgvector", "postgres", "postgresql"):
        return PgvectorBackend(**kwargs)
    raise ValueError(f"Unknown vector backend: {backend_type}. Supported: qdrant, chromadb, pgvector")
