"""PostgreSQL + pgvector vector backend for SaaS deployments."""

import logging

from sqlalchemy import text
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from sulci_core.config import SulciConfig

logger = logging.getLogger(__name__)


class PgVectorBackend:
    """Vector storage using PostgreSQL + pgvector extension.

    Uses its own async engine / connection pool (QueuePool by default).
    """

    def __init__(self, config: SulciConfig) -> None:
        self._config = config
        self._dimensions = config.embedding_dimensions
        self._engine = None
        self._session_factory = None

    async def initialize(self) -> None:
        """Create the pgvector extension and atom_vectors table."""
        db_url = self._config.database_url
        if not db_url:
            raise ValueError("SULCI_DATABASE_URL is required for postgres backend")

        self._engine = create_async_engine(db_url, echo=False)
        self._session_factory = async_sessionmaker(self._engine, expire_on_commit=False)

        async with self._engine.begin() as conn:
            await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
            await conn.execute(
                text(
                    f"CREATE TABLE IF NOT EXISTS atom_vectors ("
                    f"  atom_id TEXT PRIMARY KEY,"
                    f"  embedding vector({self._dimensions}) NOT NULL"
                    f")"
                )
            )
            # IVFFlat index for fast cosine search
            await conn.execute(
                text(
                    "CREATE INDEX IF NOT EXISTS idx_atom_vectors_embedding "
                    "ON atom_vectors USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100)"
                )
            )

        logger.info("pgvector backend initialized (dimensions=%d)", self._dimensions)

    async def upsert(self, atom_id: str, embedding: list[float], metadata: dict, document: str) -> None:
        emb_str = "[" + ",".join(str(v) for v in embedding) + "]"
        async with self._session_factory() as session, session.begin():
            await session.execute(
                text(
                    "INSERT INTO atom_vectors (atom_id, embedding) "
                    "VALUES (:id, :emb::vector) "
                    "ON CONFLICT (atom_id) DO UPDATE SET embedding = EXCLUDED.embedding"
                ),
                {"id": atom_id, "emb": emb_str},
            )

    async def upsert_batch(self, items: list[tuple[str, list[float], dict, str]]) -> None:
        if not items:
            return
        async with self._session_factory() as session, session.begin():
            for atom_id, embedding, _meta, _doc in items:
                emb_str = "[" + ",".join(str(v) for v in embedding) + "]"
                await session.execute(
                    text(
                        "INSERT INTO atom_vectors (atom_id, embedding) "
                        "VALUES (:id, :emb::vector) "
                        "ON CONFLICT (atom_id) DO UPDATE SET embedding = EXCLUDED.embedding"
                    ),
                    {"id": atom_id, "emb": emb_str},
                )

    async def delete(self, atom_id: str) -> None:
        async with self._session_factory() as session, session.begin():
            await session.execute(text("DELETE FROM atom_vectors WHERE atom_id = :id"), {"id": atom_id})

    async def delete_batch(self, atom_ids: list[str]) -> None:
        if not atom_ids:
            return
        async with self._session_factory() as session, session.begin():
            for atom_id in atom_ids:
                await session.execute(text("DELETE FROM atom_vectors WHERE atom_id = :id"), {"id": atom_id})

    async def update_metadata(self, atom_id: str, metadata: dict) -> None:
        """No-op — metadata lives in the relational table."""
        pass

    async def search(
        self, query_embedding: list[float], limit: int = 10, where: dict | None = None
    ) -> list[tuple[str, float]]:
        emb_str = "[" + ",".join(str(v) for v in query_embedding) + "]"
        fetch_limit = limit * 5 if where else limit

        async with self._session_factory() as session:
            if where:
                conditions = []
                params: dict = {"emb": emb_str, "lim": fetch_limit}
                for key, value in where.items():
                    if key == "is_active":
                        conditions.append(f"ka.is_active = :f_{key}")
                        params[f"f_{key}"] = value
                    elif key == "atom_type":
                        conditions.append(f"ka.atom_type = :f_{key}")
                        params[f"f_{key}"] = value
                where_clause = " AND ".join(conditions)
                sql = (
                    f"SELECT v.atom_id, v.embedding <=> :emb::vector AS distance "
                    f"FROM atom_vectors v "
                    f"INNER JOIN knowledge_atoms ka ON ka.id = v.atom_id "
                    f"WHERE {where_clause} "
                    f"ORDER BY distance "
                    f"LIMIT :lim"
                )
            else:
                params = {"emb": emb_str, "lim": fetch_limit}
                sql = (
                    "SELECT atom_id, embedding <=> :emb::vector AS distance "
                    "FROM atom_vectors "
                    "ORDER BY distance "
                    "LIMIT :lim"
                )

            result = await session.execute(text(sql), params)
            rows = result.fetchall()

        return [(row[0], float(row[1])) for row in rows][:limit]

    async def get_all_with_embeddings(self, where: dict | None = None, limit: int = 500) -> dict:
        async with self._session_factory() as session:
            if where:
                conditions = []
                params: dict = {"lim": limit}
                for key, value in where.items():
                    if key == "is_active":
                        conditions.append(f"ka.is_active = :f_{key}")
                        params[f"f_{key}"] = value
                    elif key == "atom_type":
                        conditions.append(f"ka.atom_type = :f_{key}")
                        params[f"f_{key}"] = value
                where_clause = " AND ".join(conditions)
                sql = (
                    f"SELECT v.atom_id, v.embedding::text "
                    f"FROM atom_vectors v "
                    f"INNER JOIN knowledge_atoms ka ON ka.id = v.atom_id "
                    f"WHERE {where_clause} "
                    f"LIMIT :lim"
                )
            else:
                params = {"lim": limit}
                sql = "SELECT atom_id, embedding::text FROM atom_vectors LIMIT :lim"

            result = await session.execute(text(sql), params)
            rows = result.fetchall()

        ids = []
        embeddings = []
        for row in rows:
            ids.append(row[0])
            # pgvector text format: [0.1,0.2,...]
            emb_text = row[1].strip("[]")
            emb = [float(v) for v in emb_text.split(",")]
            embeddings.append(emb)

        return {"ids": ids, "embeddings": embeddings}

    def distance_to_similarity(self, distance: float) -> float:
        """pgvector <=> cosine distance: 0 = identical, 2 = opposite."""
        return 1.0 - (distance / 2.0)

    async def purge(self) -> None:
        async with self._session_factory() as session, session.begin():
            await session.execute(text("DELETE FROM atom_vectors"))
