"""
pgvector.py
-----------
PGVector implementation of VectorBackend.
"""

from typing import Any, Dict, List, Optional, Sequence, Tuple
import numpy as np

try:
    import psycopg
    from pgvector.psycopg import register_vector
except ImportError:
    psycopg = None

from hieromem_core.backends.interface import VectorBackend


class PGVectorBackend(VectorBackend):
    def __init__(self, dsn: str):
        if not psycopg:
            raise ImportError(
                "psycopg[binary] and pgvector are required for PGVectorBackend"
            )
        self.dsn = dsn
        self._init_db()

    def _init_db(self):
        with psycopg.connect(self.dsn) as conn:
            conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS embeddings (
                    id TEXT PRIMARY KEY,
                    embedding vector(768),
                    metadata JSONB
                )
            """
            )
            conn.commit()

    def add(
        self,
        ids: Sequence[str],
        embeddings: np.ndarray,
        metadata: Optional[List[Dict[str, Any]]] = None,
    ) -> None:
        if not metadata:
            metadata = [{} for _ in ids]

        with psycopg.connect(self.dsn) as conn:
            register_vector(conn)
            with conn.cursor() as cur:
                data = [
                    (eid, emb, meta)
                    for eid, emb, meta in zip(ids, embeddings, metadata)
                ]
                cur.executemany(
                    "INSERT INTO embeddings (id, embedding, metadata) VALUES (%s, %s, %s) "
                    "ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding, metadata = EXCLUDED.metadata",
                    data,
                )
            conn.commit()

    def search(
        self,
        query_embedding: np.ndarray,
        k: int = 10,
        filter: Optional[Dict[str, Any]] = None,
    ) -> List[Tuple[str, float, Dict[str, Any]]]:
        query_sql = "SELECT id, embedding <=> %s as dist, metadata FROM embeddings"
        params = [query_embedding]

        if filter:
            # Simple JSONB containment
            query_sql += " WHERE metadata @> %s"
            import json

            params.append(json.dumps(filter))

        query_sql += " ORDER BY dist LIMIT %s"
        params.append(k)

        with psycopg.connect(self.dsn) as conn:
            register_vector(conn)
            results = conn.execute(query_sql, params).fetchall()

        return [(r[0], float(r[1]), r[2]) for r in results]

    def get(
        self, ids: Sequence[str]
    ) -> List[Optional[Tuple[str, np.ndarray, Dict[str, Any]]]]:
        with psycopg.connect(self.dsn) as conn:
            register_vector(conn)
            # Use unnest to maintain order or dict lookup
            res = conn.execute(
                "SELECT id, embedding, metadata FROM embeddings WHERE id = ANY(%s)",
                (list(ids),),
            ).fetchall()

        lookup = {r[0]: (r[0], np.array(r[1]), r[2]) for r in res}
        return [lookup.get(eid) for eid in ids]

    def delete(self, ids: Sequence[str]) -> None:
        with psycopg.connect(self.dsn) as conn:
            conn.execute("DELETE FROM embeddings WHERE id = ANY(%s)", (list(ids),))
            conn.commit()

    def count(self) -> int:
        with psycopg.connect(self.dsn) as conn:
            return conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0]

    def clear(self) -> None:
        with psycopg.connect(self.dsn) as conn:
            conn.execute("TRUNCATE TABLE embeddings")
            conn.commit()
