"""
pgvector_store.py
-----------------
Optimised PostgreSQL + pgvector persistence backend for HIEROMEM.
"""

from __future__ import annotations

import json
from datetime import datetime
from typing import Any, Iterable, List, Optional, Tuple, TypeVar

import networkx as nx
import numpy as np
import psycopg2
import psycopg2.extras
from psycopg2.extensions import connection as PGConnection

from hieromem_core.episode import Episode
from .memory_store import MemoryStore

_TSettings = TypeVar("_TSettings", bound=Any)


class PGVectorStore(MemoryStore):
    """Persistence and vector search backend using Postgres + pgvector."""

    def __init__(
        self,
        dsn: str,
        dim: int = 768,
        *,
        debug: bool = False,
        index_strategy: str = "ivfflat",
        ivfflat_lists: int = 200,
        ivfflat_probes: int = 10,
        hnsw_m: int = 16,
        hnsw_ef_construction: int = 128,
        hnsw_ef_search: int = 64,
    ) -> None:
        super().__init__(debug=debug)
        self.dsn = dsn
        self.dim = int(dim)
        self.index_strategy = index_strategy if index_strategy in {
            "ivfflat", "hnsw", "both"} else "ivfflat"
        self.ivfflat_lists = max(4, int(ivfflat_lists))
        self.ivfflat_probes = max(1, int(ivfflat_probes))
        self.hnsw_m = max(4, int(hnsw_m))
        self.hnsw_ef_construction = max(16, int(hnsw_ef_construction))
        self.hnsw_ef_search = max(8, int(hnsw_ef_search))

        self.conn: PGConnection = psycopg2.connect(self.dsn)
        self.conn.autocommit = True
        self._ensure_schema()
        self._log("PGVectorStore connected and schema verified")

    # ------------------------------------------------------------------ #
    # Schema setup
    # ------------------------------------------------------------------ #
    def _ensure_schema(self) -> None:
        with self.conn.cursor() as cur:
            cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
            cur.execute(
                f"""
                CREATE TABLE IF NOT EXISTS episodes (
                    id TEXT PRIMARY KEY,
                    embedding vector({self.dim}),
                    graph JSONB,
                    tocc TIMESTAMP,
                    meta JSONB,
                    parent TEXT,
                    trust REAL
                )
                """
            )
            cur.execute(
                f"ALTER TABLE episodes ADD COLUMN IF NOT EXISTS embedding vector({self.dim})")
            cur.execute(
                "CREATE INDEX IF NOT EXISTS idx_trust ON episodes (trust)")
            cur.execute(
                "CREATE INDEX IF NOT EXISTS idx_tocc ON episodes (tocc)")
            if self.index_strategy in {"ivfflat", "both"}:
                cur.execute(
                    """
                    CREATE INDEX IF NOT EXISTS idx_episodes_embedding_ivfflat
                    ON episodes USING ivfflat (embedding vector_cosine_ops)
                    WITH (lists = %s)
                    """,
                    (self.ivfflat_lists,),
                )
            if self.index_strategy in {"hnsw", "both"}:
                cur.execute(
                    """
                    CREATE INDEX IF NOT EXISTS idx_episodes_embedding_hnsw
                    ON episodes USING hnsw (embedding vector_cosine_ops)
                    WITH (m = %s, ef_construction = %s)
                    """,
                    (self.hnsw_m, self.hnsw_ef_construction),
                )

    # ------------------------------------------------------------------ #
    # Encoding helpers
    # ------------------------------------------------------------------ #
    def _encode_episode(self,
                        ep: Episode) -> Tuple[str,
                                              List[float],
                                              str,
                                              datetime,
                                              str,
                                              Optional[str],
                                              float]:
        self.validate_episode(ep)
        G = ep.G
        graph_data = {
            "nodes": [{"id": n, **(G.nodes[n] or {})} for n in G.nodes],
            "edges": [{"u": u, "v": v, **(G.edges[u, v] or {})} for u, v in G.edges],
        }
        emb = [float(x) for x in ep.e.tolist()]
        return (
            ep.id,
            emb,
            json.dumps(graph_data),
            ep.tocc,
            json.dumps(ep.meta),
            ep.parent,
            float(ep.trust),
        )

    def _decode_episode(self, row: Tuple[Any, ...]) -> Episode:
        eid, e_vec, graph_json, tocc, meta_json, parent, trust = row
        G_data = json.loads(graph_json)
        G = nx.DiGraph()
        for n in G_data.get("nodes", []):
            nid = n.pop("id")
            G.add_node(nid, **n)
        for e in G_data.get("edges", []):
            u, v = e.pop("u"), e.pop("v")
            G.add_edge(u, v, **e)
        meta_raw = json.loads(meta_json or "{}")
        if not isinstance(meta_raw, dict):
            meta_raw = {}
        meta = dict(meta_raw)
        e = np.array(e_vec, dtype=np.float32)
        return Episode(
            id=eid,
            e=e,
            G=G,
            tocc=tocc,
            meta=meta,
            parent=parent,
            trust=trust)

    # ------------------------------------------------------------------ #
    # CRUD
    # ------------------------------------------------------------------ #
    def save_episode(self, ep: Episode) -> None:
        self.save_many([ep])

    def save_many(self, episodes: Iterable[Episode]) -> None:
        batch = [self._encode_episode(ep) for ep in episodes]
        if not batch:
            return
        with self.conn.cursor() as cur:
            psycopg2.extras.execute_values(
                cur,
                """
                INSERT INTO episodes (id, embedding, graph, tocc, meta, parent, trust)
                VALUES %s
                ON CONFLICT (id) DO UPDATE SET
                    embedding = EXCLUDED.embedding,
                    graph = EXCLUDED.graph,
                    tocc = EXCLUDED.tocc,
                    meta = EXCLUDED.meta,
                    parent = EXCLUDED.parent,
                    trust = EXCLUDED.trust
                """,
                batch,
            )
        self._log(f"Saved batch of {len(batch)} episodes")

    def load_all(self) -> List[Episode]:
        with self.conn.cursor() as cur:
            cur.execute(
                "SELECT id, embedding, graph, tocc, meta, parent, trust FROM episodes")
            rows = cur.fetchall()
        eps = [self._decode_episode(r) for r in rows]
        self._log(f"Loaded {len(eps)} episodes")
        return eps

    def load_by_id(self, episode_id: str) -> Optional[Episode]:
        with self.conn.cursor() as cur:
            cur.execute(
                "SELECT id, embedding, graph, tocc, meta, parent, trust FROM episodes WHERE id=%s",
                (episode_id,),
            )
            row = cur.fetchone()
        return self._decode_episode(row) if row else None

    def update_trust(self, episode_id: str, trust: float) -> None:
        with self.conn.cursor() as cur:
            cur.execute("UPDATE episodes SET trust=%s WHERE id=%s",
                        (float(trust), episode_id))
        self._log(f"Updated trust for {episode_id} → {trust:.3f}")

    def remove(self, episode_id: str) -> None:
        with self.conn.cursor() as cur:
            cur.execute("DELETE FROM episodes WHERE id=%s", (episode_id,))
        self._log(f"Removed episode {episode_id}")

    def remove_many(self, episode_ids: Iterable[str]) -> None:
        ids = list(episode_ids)
        if not ids:
            return
        with self.conn.cursor() as cur:
            psycopg2.extras.execute_batch(
                cur, "DELETE FROM episodes WHERE id=%s", [
                    (i,) for i in ids])
        self._log(f"Removed {len(ids)} episodes")

    def clear(self) -> None:
        with self.conn.cursor() as cur:
            cur.execute("DELETE FROM episodes")
        self._log("Cleared all episodes from Postgres store")

    # ------------------------------------------------------------------ #
    # Vector Search
    # ------------------------------------------------------------------ #
    def search_similar(self, query_vec: np.ndarray,
                       k: int = 5) -> List[Tuple[str, float]]:
        k = max(1, min(int(k), 200))
        qlist = [float(x) for x in query_vec.tolist()]
        with self.conn.cursor() as cur:
            if self.index_strategy in {"ivfflat", "both"}:
                cur.execute("SET LOCAL ivfflat.probes = %s",
                            (self.ivfflat_probes,))
            if self.index_strategy in {"hnsw", "both"}:
                cur.execute("SET LOCAL hnsw.ef_search = %s",
                            (self.hnsw_ef_search,))
            cur.execute(
                "SELECT id, 1 - (embedding <=> %s::vector) AS score "
                "FROM episodes ORDER BY embedding <=> %s::vector ASC LIMIT %s",
                (qlist, qlist, k),
            )
            results = cur.fetchall()
        return [(r[0], float(r[1])) for r in results]

    # ------------------------------------------------------------------ #
    # Context management
    # ------------------------------------------------------------------ #
    def close(self) -> None:
        self.conn.close()
        self._log("Postgres connection closed")

    # ------------------------------------------------------------------ #
    # Convenience constructors
    # ------------------------------------------------------------------ #
    @classmethod
    def from_settings(
        cls,
        settings: _TSettings,
        *,
        dim: int,
        debug: bool = False,
        **kwargs: Any,
    ) -> "PGVectorStore":
        dsn = getattr(
            settings,
            "pgvector_dsn",
            None) or kwargs.pop(
            "dsn",
            None)
        if not dsn:
            raise ValueError("pgvector DSN missing from settings")
        return cls(dsn, dim=dim, debug=debug, **kwargs)
