"""PostgresBackend — direct Postgres/Neon driver using psycopg."""

from __future__ import annotations

import logging
from typing import Any

from psycopg.rows import dict_row
from psycopg.types.json import Jsonb
from psycopg_pool import ConnectionPool

from ogham.config import settings
from ogham.retry import with_retry

logger = logging.getLogger(__name__)

# Allowlist of valid memory table columns to prevent SQL injection via dynamic column names
_ALLOWED_MEMORY_COLUMNS = frozenset(
    {
        "content",
        "embedding",
        "metadata",
        "source",
        "profile",
        "tags",
        "expires_at",
        "access_count",
        "last_accessed_at",
        "confidence",
    }
)

# Columns to SELECT when we don't want the embedding or fts vectors.
_COLUMNS = (
    "id, content, metadata, source, profile, tags,"
    " created_at, updated_at, expires_at, access_count, last_accessed_at, confidence"
)


def _embedding_literal(embedding: list[float]) -> str:
    """Format as Postgres vector literal '[0.1,0.2,...]'."""
    return "[" + ",".join(str(v) for v in embedding) + "]"


class PostgresBackend:
    """DatabaseBackend implementation using psycopg + connection pool."""

    def __init__(self) -> None:
        self._pool: ConnectionPool | None = None

    def _get_pool(self) -> ConnectionPool:
        if self._pool is None:
            if not settings.database_url:
                raise RuntimeError("DATABASE_URL is required for PostgresBackend")
            self._pool = ConnectionPool(
                conninfo=settings.database_url,
                min_size=1,
                max_size=5,
                kwargs={"row_factory": dict_row},
            )
            self._ensure_columns()
        return self._pool

    def _ensure_columns(self) -> None:
        """Auto-add columns introduced in newer versions.

        Runs on first connection so upgraders don't need manual migrations.
        All statements use IF NOT EXISTS -- safe to run repeatedly.
        """
        migrations = [
            # v0.7.0: importance, surprise, compression
            "ALTER TABLE memories ADD COLUMN IF NOT EXISTS importance real DEFAULT 0.5",
            "ALTER TABLE memories ADD COLUMN IF NOT EXISTS surprise real DEFAULT 0.5",
            "ALTER TABLE memories ADD COLUMN IF NOT EXISTS compression_level integer DEFAULT 0",
            "ALTER TABLE memories ADD COLUMN IF NOT EXISTS original_content text",
            # v0.8.0: temporal columns
            "ALTER TABLE memories ADD COLUMN IF NOT EXISTS occurrence_period tstzrange",
            "ALTER TABLE memories ADD COLUMN IF NOT EXISTS recurrence_days integer[]",
        ]
        try:
            with self._pool.connection() as conn:  # type: ignore[union-attr]
                for sql in migrations:
                    conn.execute(sql)
                conn.commit()
        except Exception as e:
            import logging

            logging.getLogger(__name__).warning("Auto-migration skipped: %s", e)

    # ── Helper ────────────────────────────────────────────────────────

    def _execute(
        self,
        query: str,
        params: dict[str, Any] | None = None,
        *,
        fetch: str = "all",
    ) -> Any:
        """Central query execution.

        fetch: "all" -> list[dict], "one" -> dict|None,
               "scalar" -> single value, "none" -> None
        """
        with self._get_pool().connection() as conn:
            with conn.cursor() as cur:
                cur.execute(query, params)
                if fetch == "none":
                    return None
                if fetch == "scalar":
                    row = cur.fetchone()
                    if row is None:
                        return None
                    # dict_row returns a dict — get the first value
                    return next(iter(row.values()))
                if fetch == "one":
                    return cur.fetchone()
                # fetch == "all"
                return cur.fetchall()

    # ── Core CRUD ─────────────────────────────────────────────────────

    def store_memory(
        self,
        content: str,
        embedding: list[float],
        profile: str,
        metadata: dict[str, Any] | None = None,
        source: str | None = None,
        tags: list[str] | None = None,
        expires_at: str | None = None,
        importance: float = 0.5,
        surprise: float = 0.5,
        recurrence_days: list[int] | None = None,
    ) -> dict[str, Any]:
        cols = [
            "content",
            "embedding",
            "profile",
            "metadata",
            "source",
            "tags",
            "importance",
            "surprise",
        ]
        vals = [
            "%(content)s",
            "%(embedding)s::vector",
            "%(profile)s",
            "%(metadata)s",
            "%(source)s",
            "%(tags)s",
            "%(importance)s",
            "%(surprise)s",
        ]
        params: dict[str, Any] = {
            "content": content,
            "embedding": _embedding_literal(embedding),
            "profile": profile,
            "metadata": Jsonb(metadata or {}),
            "source": source,
            "tags": tags or [],
            "importance": importance,
            "surprise": surprise,
        }
        if expires_at is not None:
            cols.append("expires_at")
            vals.append("%(expires_at)s")
            params["expires_at"] = expires_at
        if recurrence_days is not None:
            cols.append("recurrence_days")
            vals.append("%(recurrence_days)s")
            params["recurrence_days"] = recurrence_days

        sql = f"INSERT INTO memories ({', '.join(cols)}) VALUES ({', '.join(vals)}) RETURNING *"
        row = self._execute(sql, params, fetch="one")
        if row is None:
            raise RuntimeError("INSERT returned no data")
        # Drop large vector fields from result
        row.pop("embedding", None)
        row.pop("fts", None)
        return row

    def store_memories_batch(self, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Insert multiple memories. Each row comes pre-formatted from export_import.py."""
        if not rows:
            return []
        results: list[dict[str, Any]] = []
        with self._get_pool().connection() as conn:
            with conn.cursor() as cur:
                for row in rows:
                    cols = list(row.keys())
                    for col in cols:
                        if col not in _ALLOWED_MEMORY_COLUMNS:
                            raise ValueError(f"Unknown column: {col!r}")
                    placeholders = []
                    params: dict[str, Any] = {}
                    for col in cols:
                        param_name = col
                        val = row[col]
                        if col == "embedding":
                            placeholders.append(f"%({param_name})s::vector")
                            # Already a string from export_import.py
                            params[param_name] = (
                                val if isinstance(val, str) else _embedding_literal(val)
                            )
                        elif col == "metadata":
                            placeholders.append(f"%({param_name})s")
                            params[param_name] = Jsonb(val) if not isinstance(val, Jsonb) else val
                        else:
                            placeholders.append(f"%({param_name})s")
                            params[param_name] = val

                    sql = (
                        f"INSERT INTO memories ({', '.join(cols)})"
                        f" VALUES ({', '.join(placeholders)}) RETURNING *"
                    )
                    cur.execute(sql, params)
                    result = cur.fetchone()
                    if result:
                        result.pop("embedding", None)
                        result.pop("fts", None)
                        results.append(result)
        return results

    def update_memory(
        self, memory_id: str, updates: dict[str, Any], profile: str
    ) -> dict[str, Any]:
        if not updates:
            raise ValueError("No updates provided")
        for key in updates:
            if key not in _ALLOWED_MEMORY_COLUMNS:
                raise ValueError(f"Unknown column: {key!r}")
        set_clauses = []
        params: dict[str, Any] = {"id": memory_id, "profile": profile}
        for key, val in updates.items():
            param = f"u_{key}"
            if key == "metadata":
                set_clauses.append(f"{key} = %({param})s")
                params[param] = Jsonb(val)
            elif key == "embedding":
                set_clauses.append(f"{key} = %({param})s::vector")
                params[param] = val if isinstance(val, str) else _embedding_literal(val)
            else:
                set_clauses.append(f"{key} = %({param})s")
                params[param] = val

        sql = (
            f"UPDATE memories SET {', '.join(set_clauses)}"
            f" WHERE id = %(id)s AND profile = %(profile)s"
            f" RETURNING {_COLUMNS}"
        )
        row = self._execute(sql, params, fetch="one")
        if row is None:
            raise KeyError(f"Memory {memory_id!r} not found in profile {profile!r}")
        return row

    def get_memory_by_id(self, memory_id: str, profile: str) -> dict[str, Any] | None:
        sql = f"SELECT {_COLUMNS} FROM memories WHERE id = %(id)s AND profile = %(profile)s"
        return self._execute(sql, {"id": memory_id, "profile": profile}, fetch="one")

    def delete_memory(self, memory_id: str, profile: str) -> bool:
        row = self._execute(
            "DELETE FROM memories WHERE id = %(id)s AND profile = %(profile)s RETURNING id",
            {"id": memory_id, "profile": profile},
            fetch="one",
        )
        return row is not None

    # ── Search & Retrieval ────────────────────────────────────────────

    @with_retry(max_attempts=2, base_delay=0.3)
    def search_memories(
        self,
        query_embedding: list[float],
        profile: str,
        threshold: float | None = None,
        limit: int | None = None,
        tags: list[str] | None = None,
        source: str | None = None,
    ) -> list[dict[str, Any]]:
        params: dict[str, Any] = {
            "embedding": _embedding_literal(query_embedding),
            "threshold": threshold or settings.default_match_threshold,
            "limit": limit or settings.default_match_count,
            "tags": tags,
            "source": source,
            "profile": profile,
        }
        return self._execute(
            "SELECT * FROM match_memories("
            "  %(embedding)s::vector, %(threshold)s, %(limit)s,"
            "  %(tags)s, %(source)s, %(profile)s"
            ")",
            params,
        )

    @with_retry(max_attempts=2, base_delay=0.3)
    def hybrid_search_memories(
        self,
        query_text: str,
        query_embedding: list[float],
        profile: str,
        limit: int | None = None,
        tags: list[str] | None = None,
        source: str | None = None,
        profiles: list[str] | None = None,
    ) -> list[dict[str, Any]]:
        params: dict[str, Any] = {
            "query_text": query_text,
            "embedding": _embedding_literal(query_embedding),
            "limit": limit or settings.default_match_count,
            "profile": profile,
            "tags": tags,
            "source": source,
            "profiles": profiles,
        }
        return self._execute(
            "SELECT * FROM hybrid_search_memories("
            "  %(query_text)s, %(embedding)s::vector, %(limit)s,"
            "  %(profile)s, %(tags)s, %(source)s,"
            "  0.3, 0.7, 10, %(profiles)s"
            ")",
            params,
        )

    @with_retry(max_attempts=2, base_delay=0.3)
    def list_recent_memories(
        self,
        profile: str,
        limit: int = 10,
        source: str | None = None,
        tags: list[str] | None = None,
    ) -> list[dict[str, Any]]:
        where = ["profile = %(profile)s", "(expires_at IS NULL OR expires_at > now())"]
        params: dict[str, Any] = {"profile": profile, "limit": limit}
        if source:
            where.append("source = %(source)s")
            params["source"] = source
        if tags:
            where.append("tags && %(tags)s")
            params["tags"] = tags

        sql = (
            f"SELECT {_COLUMNS} FROM memories"
            f" WHERE {' AND '.join(where)}"
            " ORDER BY created_at DESC LIMIT %(limit)s"
        )
        return self._execute(sql, params)

    @with_retry(max_attempts=2, base_delay=0.3)
    def get_all_memories_full(self, profile: str) -> list[dict[str, Any]]:
        rows: list[dict[str, Any]] = []
        batch = 1000
        last_created_at: str | None = None
        last_id: str | None = None
        while True:
            params: dict[str, Any] = {"profile": profile, "batch": batch}
            conditions = [
                "profile = %(profile)s",
                "(expires_at IS NULL OR expires_at > now())",
            ]
            if last_created_at is not None:
                conditions.append(
                    "(created_at > %(last_created_at)s OR "
                    "(created_at = %(last_created_at)s AND id > %(last_id)s))"
                )
                params["last_created_at"] = last_created_at
                params["last_id"] = last_id
            where = " AND ".join(conditions)
            sql = (
                f"SELECT {_COLUMNS} FROM memories WHERE {where}"
                " ORDER BY created_at, id LIMIT %(batch)s"
            )
            page = self._execute(sql, params)
            rows.extend(page)
            if len(page) < batch:
                break
            last_row = page[-1]
            last_created_at = str(last_row["created_at"])
            last_id = str(last_row["id"])
        return rows

    @with_retry(max_attempts=2, base_delay=0.3)
    def get_all_memories_content(self, profile: str | None = None) -> list[dict[str, Any]]:
        rows: list[dict[str, Any]] = []
        batch = 1000
        last_id: str | None = None
        while True:
            params: dict[str, Any] = {"batch": batch}
            conditions: list[str] = []
            if profile:
                conditions.append("profile = %(profile)s")
                params["profile"] = profile
            if last_id is not None:
                conditions.append("id > %(last_id)s")
                params["last_id"] = last_id
            where = (" WHERE " + " AND ".join(conditions)) if conditions else ""
            sql = f"SELECT id, content FROM memories{where} ORDER BY id LIMIT %(batch)s"
            page = self._execute(sql, params)
            rows.extend(page)
            if len(page) < batch:
                break
            last_id = str(page[-1]["id"])
        return rows

    # ── Batch & Embedding Ops ─────────────────────────────────────────

    @with_retry(max_attempts=2, base_delay=0.3)
    def batch_check_duplicates(
        self,
        query_embeddings: list[list[float]],
        profile: str,
        threshold: float = 0.8,
    ) -> list[bool]:
        if not query_embeddings:
            return []
        emb_literals = [_embedding_literal(e) for e in query_embeddings]
        result = self._execute(
            "SELECT batch_check_duplicates(  %(embeddings)s::vector[], %(threshold)s, %(profile)s)",
            {"embeddings": emb_literals, "threshold": threshold, "profile": profile},
            fetch="scalar",
        )
        return result if result is not None else []

    def batch_update_embeddings(self, ids: list[str], embeddings: list[list[float]]) -> int:
        emb_literals = [_embedding_literal(e) for e in embeddings]
        result = self._execute(
            "SELECT batch_update_embeddings(%(ids)s::uuid[], %(embeddings)s::vector[])",
            {"ids": ids, "embeddings": emb_literals},
            fetch="scalar",
        )
        return result if isinstance(result, int) else 0

    # ── Access & Confidence ───────────────────────────────────────────

    def record_access(self, memory_ids: list[str]) -> None:
        if not memory_ids:
            return
        self._execute(
            "SELECT record_access(%(ids)s::uuid[])",
            {"ids": memory_ids},
            fetch="none",
        )

    def update_confidence(self, memory_id: str, signal: float, profile: str) -> float:
        result = self._execute(
            "SELECT update_confidence(%(id)s, %(signal)s, %(profile)s)",
            {"id": memory_id, "signal": signal, "profile": profile},
            fetch="scalar",
        )
        return result if isinstance(result, float) else 0.5

    # ── Profile & Stats ───────────────────────────────────────────────

    @with_retry(max_attempts=2, base_delay=0.3)
    def get_memory_stats(self, profile: str) -> dict[str, Any]:
        result = self._execute(
            "SELECT get_memory_stats_sql(%(profile)s)",
            {"profile": profile},
            fetch="scalar",
        )
        if not result:
            return {"profile": profile, "total": 0, "sources": {}, "top_tags": []}
        if isinstance(result, dict):
            return result
        # Should already be a dict from jsonb, but handle list edge case
        return result[0] if isinstance(result, list) else result

    @with_retry(max_attempts=2, base_delay=0.3)
    def list_profiles(self) -> list[dict[str, Any]]:
        return self._execute("SELECT * FROM get_profile_counts()")

    def get_profile_ttl(self, profile: str) -> int | None:
        row = self._execute(
            "SELECT ttl_days FROM profile_settings WHERE profile = %(profile)s",
            {"profile": profile},
            fetch="one",
        )
        if row is None:
            return None
        return row.get("ttl_days")

    def set_profile_ttl(self, profile: str, ttl_days: int | None) -> dict[str, Any]:
        row = self._execute(
            "INSERT INTO profile_settings (profile, ttl_days)"
            " VALUES (%(profile)s, %(ttl_days)s)"
            " ON CONFLICT (profile) DO UPDATE SET ttl_days = EXCLUDED.ttl_days"
            " RETURNING *",
            {"profile": profile, "ttl_days": ttl_days},
            fetch="one",
        )
        return row

    def cleanup_expired(self, profile: str) -> int:
        result = self._execute(
            "SELECT cleanup_expired_memories(%(profile)s)",
            {"profile": profile},
            fetch="scalar",
        )
        return result if isinstance(result, int) else 0

    def count_expired(self, profile: str) -> int:
        result = self._execute(
            "SELECT count_expired_memories(%(profile)s)",
            {"profile": profile},
            fetch="scalar",
        )
        return result if isinstance(result, int) else 0

    # ── Relationships / Knowledge Graph ───────────────────────────────

    @with_retry(max_attempts=2, base_delay=0.3)
    def auto_link_memory(
        self,
        memory_id: str,
        embedding: list[float],
        profile: str,
        threshold: float = 0.85,
        max_links: int = 5,
    ) -> int:
        result = self._execute(
            "SELECT auto_link_memory("
            "  %(memory_id)s, %(embedding)s::vector,"
            "  %(threshold)s, %(max_links)s, %(profile)s"
            ")",
            {
                "memory_id": memory_id,
                "embedding": _embedding_literal(embedding),
                "threshold": threshold,
                "max_links": max_links,
                "profile": profile,
            },
            fetch="scalar",
        )
        return result if isinstance(result, int) else 0

    @with_retry(max_attempts=2, base_delay=0.3)
    def link_unlinked_memories(
        self,
        profile: str,
        threshold: float = 0.85,
        max_links: int = 5,
        batch_size: int = 100,
    ) -> int:
        result = self._execute(
            "SELECT link_unlinked_memories("
            "  %(profile)s, %(threshold)s, %(max_links)s, %(batch_size)s"
            ")",
            {
                "profile": profile,
                "threshold": threshold,
                "max_links": max_links,
                "batch_size": batch_size,
            },
            fetch="scalar",
        )
        return result if isinstance(result, int) else 0

    @with_retry(max_attempts=2, base_delay=0.3)
    def explore_memory_graph(
        self,
        query_text: str,
        query_embedding: list[float],
        profile: str,
        limit: int = 5,
        depth: int = 1,
        min_strength: float = 0.5,
        tags: list[str] | None = None,
        source: str | None = None,
    ) -> list[dict[str, Any]]:
        return self._execute(
            "SELECT * FROM explore_memory_graph("
            "  %(query_text)s, %(embedding)s::vector, %(profile)s,"
            "  %(limit)s, %(depth)s, %(min_strength)s,"
            "  %(tags)s, %(source)s"
            ")",
            {
                "query_text": query_text,
                "embedding": _embedding_literal(query_embedding),
                "profile": profile,
                "limit": limit,
                "depth": depth,
                "min_strength": min_strength,
                "tags": tags,
                "source": source,
            },
        )

    def create_relationship(
        self,
        source_id: str,
        target_id: str,
        relationship: str,
        strength: float = 1.0,
        created_by: str = "user",
        metadata: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        row = self._execute(
            "INSERT INTO memory_relationships"
            " (source_id, target_id, relationship, strength, created_by, metadata)"
            " VALUES (%(source_id)s, %(target_id)s, %(relationship)s::relationship_type,"
            "         %(strength)s, %(created_by)s, %(metadata)s)"
            " RETURNING *",
            {
                "source_id": source_id,
                "target_id": target_id,
                "relationship": relationship,
                "strength": strength,
                "created_by": created_by,
                "metadata": Jsonb(metadata or {}),
            },
            fetch="one",
        )
        if row is None:
            raise RuntimeError("Insert returned no data for relationship")
        return row

    @with_retry(max_attempts=2, base_delay=0.3)
    def get_related_memories(
        self,
        memory_id: str,
        depth: int = 1,
        min_strength: float = 0.5,
        relationship_types: list[str] | None = None,
        limit: int = 20,
    ) -> list[dict[str, Any]]:
        return self._execute(
            "SELECT * FROM get_related_memories("
            "  %(memory_id)s, %(depth)s, %(min_strength)s,"
            "  %(types)s::relationship_type[], %(limit)s"
            ")",
            {
                "memory_id": memory_id,
                "depth": depth,
                "min_strength": min_strength,
                "types": relationship_types,
                "limit": limit,
            },
        )
