"""PostgreSQL-backed memory backend."""
from __future__ import annotations

import asyncio
from datetime import datetime
import json
import re
from typing import Any

from logos_adk.errors import MemoryBackendError
from logos_adk.memory import MemoryBackend, MemoryEntry
from logos_adk.persistence import get_postgresql_schema_version, set_postgresql_schema_version
from logos_adk.postgres_utils import pooled_connect, safe_identifier


MEMORY_SCHEMA_KIND = "memory_backend"
MEMORY_SCHEMA_VERSION = 1


class PostgreSQLBackend(MemoryBackend):
    """Persistent memory backend using PostgreSQL."""

    def __init__(self, dsn: str, table_name: str = "memories") -> None:
        self.dsn = dsn
        self.table_name = safe_identifier(table_name)

    async def store(self, entry: MemoryEntry) -> None:
        await asyncio.to_thread(self._store_sync, entry)

    async def search(
        self,
        query: str,
        limit: int = 10,
        session_id: str | None = None,
    ) -> list[MemoryEntry]:
        return await asyncio.to_thread(self._search_sync, query, limit, session_id)

    async def delete(self, entry_id: str) -> None:
        await asyncio.to_thread(self._delete_sync, entry_id)

    def _connect(self) -> Any:
        try:
            return pooled_connect(self.dsn)
        except Exception as exc:
            raise MemoryBackendError(
                f"Failed to open PostgreSQL memory backend: {exc}"
            ) from exc

    def _ensure_schema(self, connection: Any) -> None:
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    id TEXT PRIMARY KEY,
                    content TEXT NOT NULL,
                    embedding JSONB,
                    metadata JSONB NOT NULL DEFAULT '{{}}'::jsonb,
                    created_at TIMESTAMPTZ NOT NULL,
                    session_id TEXT
                )
                """
            )
            cursor.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self.table_name}_session_id "
                f"ON {self.table_name} (session_id)"
            )
            cursor.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self.table_name}_created_at "
                f"ON {self.table_name} (created_at DESC)"
            )
        recorded_version = get_postgresql_schema_version(connection, MEMORY_SCHEMA_KIND)
        if recorded_version != MEMORY_SCHEMA_VERSION:
            set_postgresql_schema_version(connection, MEMORY_SCHEMA_KIND, MEMORY_SCHEMA_VERSION)
        connection.commit()

    def _store_sync(self, entry: MemoryEntry) -> None:
        try:
            with self._connect() as connection:
                self._ensure_schema(connection)
                with connection.cursor() as cursor:
                    cursor.execute(
                        f"""
                        INSERT INTO {self.table_name}
                        (id, content, embedding, metadata, created_at, session_id)
                        VALUES (%s, %s, %s::jsonb, %s::jsonb, %s, %s)
                        ON CONFLICT (id) DO UPDATE SET
                            content = EXCLUDED.content,
                            embedding = EXCLUDED.embedding,
                            metadata = EXCLUDED.metadata,
                            created_at = EXCLUDED.created_at,
                            session_id = EXCLUDED.session_id
                        """,
                        (
                            entry.id,
                            entry.content,
                            json.dumps(entry.embedding) if entry.embedding is not None else None,
                            json.dumps(entry.metadata),
                            entry.created_at,
                            entry.session_id,
                        ),
                    )
                connection.commit()
        except Exception as exc:
            raise MemoryBackendError(f"Failed to store memory entry: {exc}") from exc

    def _search_sync(
        self,
        query: str,
        limit: int,
        session_id: str | None,
    ) -> list[MemoryEntry]:
        normalized_query = query.strip().lower()
        keywords = [
            token
            for token in re.findall(r"[a-zA-Z0-9_]+", normalized_query)
            if len(token) >= 4
        ]

        try:
            with self._connect() as connection:
                self._ensure_schema(connection)
                clauses: list[str] = []
                params: list[object] = []

                if normalized_query:
                    content_matchers = ["LOWER(content) LIKE LOWER(%s)"]
                    content_params: list[object] = [f"%{normalized_query}%"]
                    if keywords:
                        content_matchers.extend("LOWER(content) LIKE LOWER(%s)" for _ in keywords)
                        content_params.extend(f"%{keyword}%" for keyword in keywords)

                    clauses.append("(" + " OR ".join(content_matchers) + ")")
                    params.extend(content_params)

                if session_id is None:
                    clauses.append("session_id IS NULL")
                else:
                    clauses.append("session_id = %s")
                    params.append(session_id)

                where_sql = " AND ".join(clauses) if clauses else "TRUE"
                params.append(limit)

                with connection.cursor() as cursor:
                    cursor.execute(
                        f"""
                        SELECT id, content, embedding, metadata, created_at, session_id
                        FROM {self.table_name}
                        WHERE {where_sql}
                        ORDER BY created_at DESC
                        LIMIT %s
                        """,
                        params,
                    )
                    rows = cursor.fetchall()
        except Exception as exc:
            raise MemoryBackendError(f"Failed to search memory entries: {exc}") from exc

        return [self._row_to_entry(row) for row in rows]

    def _delete_sync(self, entry_id: str) -> None:
        try:
            with self._connect() as connection:
                self._ensure_schema(connection)
                with connection.cursor() as cursor:
                    cursor.execute(
                        f"DELETE FROM {self.table_name} WHERE id = %s",
                        (entry_id,),
                    )
                connection.commit()
        except Exception as exc:
            raise MemoryBackendError(f"Failed to delete memory entry: {exc}") from exc

    def _row_to_entry(self, row: Any) -> MemoryEntry:
        entry_id, content, embedding, metadata, created_at, session_id = row
        if isinstance(embedding, str):
            embedding = json.loads(embedding)
        if isinstance(metadata, str):
            metadata = json.loads(metadata)
        return MemoryEntry(
            id=entry_id,
            content=content,
            embedding=embedding,
            metadata=metadata or {},
            created_at=created_at if isinstance(created_at, datetime) else datetime.fromisoformat(created_at),
            session_id=session_id,
        )
