"""PostgreSQL cache backend."""
from __future__ import annotations

import asyncio
import json
import time
from typing import Any

from logos_adk.cache import CacheBackend, CacheEntry, _cosine_similarity
from logos_adk.persistence import get_postgresql_schema_version, set_postgresql_schema_version
from logos_adk.postgres_utils import pooled_connect, resolve_database_url, safe_identifier


CACHE_SCHEMA_KIND = "cache_backend"
CACHE_SCHEMA_VERSION = 1


class PostgreSQLCacheBackend(CacheBackend):
    """PostgreSQL-backed cache with TTL support."""

    def __init__(self, dsn: str | None = None, table_name: str = "cache") -> None:
        self._dsn = resolve_database_url(dsn)
        if not self._dsn:
            raise ValueError(
                "PostgreSQL cache backend requires a DSN or DATABASE_URL/LOGOS_ADK_DATABASE_URL"
            )
        self._table_name = safe_identifier(table_name)

    def _connect(self) -> Any:
        return pooled_connect(self._dsn)

    def _ensure_schema(self, connection: Any) -> None:
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self._table_name} (
                    key TEXT PRIMARY KEY,
                    response_data JSONB NOT NULL,
                    created_at TEXT NOT NULL,
                    hit_count INTEGER NOT NULL DEFAULT 0,
                    vector JSONB,
                    scope TEXT,
                    expires_at DOUBLE PRECISION
                )
                """
            )
            cursor.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self._table_name}_scope "
                f"ON {self._table_name} (scope)"
            )
        recorded_version = get_postgresql_schema_version(connection, CACHE_SCHEMA_KIND)
        if recorded_version != CACHE_SCHEMA_VERSION:
            set_postgresql_schema_version(connection, CACHE_SCHEMA_KIND, CACHE_SCHEMA_VERSION)
        connection.commit()

    def _get_sync(self, key: str) -> CacheEntry | None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    SELECT response_data, created_at, hit_count, vector, scope, expires_at
                    FROM {self._table_name}
                    WHERE key = %s
                    """,
                    (key,),
                )
                row = cursor.fetchone()
                if row is None:
                    return None

                response_data, created_at, hit_count, vector, scope, expires_at = row
                if expires_at is not None and time.time() > expires_at:
                    cursor.execute(
                        f"DELETE FROM {self._table_name} WHERE key = %s",
                        (key,),
                    )
                    connection.commit()
                    return None

                cursor.execute(
                    f"UPDATE {self._table_name} SET hit_count = hit_count + 1 WHERE key = %s",
                    (key,),
                )
            connection.commit()

        new_count = hit_count + 1
        if isinstance(response_data, str):
            response_data = json.loads(response_data)
        if isinstance(vector, str):
            vector = json.loads(vector)
        return CacheEntry(
            response_data=response_data,
            created_at=created_at,
            hit_count=new_count,
            vector=vector,
            scope=scope,
        )

    async def get(self, key: str) -> CacheEntry | None:
        return await asyncio.to_thread(self._get_sync, key)

    def _put_sync(self, key: str, entry: CacheEntry, ttl: int | None) -> None:
        expires_at = (time.time() + ttl) if ttl is not None else None
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    INSERT INTO {self._table_name}
                    (key, response_data, created_at, hit_count, vector, scope, expires_at)
                    VALUES (%s, %s::jsonb, %s, %s, %s::jsonb, %s, %s)
                    ON CONFLICT (key) DO UPDATE SET
                        response_data = EXCLUDED.response_data,
                        created_at = EXCLUDED.created_at,
                        hit_count = EXCLUDED.hit_count,
                        vector = EXCLUDED.vector,
                        scope = EXCLUDED.scope,
                        expires_at = EXCLUDED.expires_at
                    """,
                    (
                        key,
                        json.dumps(entry.response_data),
                        entry.created_at,
                        entry.hit_count,
                        json.dumps(entry.vector) if entry.vector is not None else None,
                        entry.scope,
                        expires_at,
                    ),
                )
            connection.commit()

    async def put(self, key: str, entry: CacheEntry, ttl: int | None) -> None:
        await asyncio.to_thread(self._put_sync, key, entry, ttl)

    def _invalidate_sync(self, key: str) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(f"DELETE FROM {self._table_name} WHERE key = %s", (key,))
            connection.commit()

    async def invalidate(self, key: str) -> None:
        await asyncio.to_thread(self._invalidate_sync, key)

    def _clear_sync(self) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(f"DELETE FROM {self._table_name}")
            connection.commit()

    async def clear(self) -> None:
        await asyncio.to_thread(self._clear_sync)

    def _get_stats_sync(self) -> dict[str, Any]:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"SELECT COUNT(*), COALESCE(SUM(hit_count), 0) FROM {self._table_name}"
                )
                row = cursor.fetchone() or (0, 0)
        total_entries, total_hits = row
        return {
            "backend": "PostgreSQL",
            "total_entries": int(total_entries or 0),
            "total_hits": int(total_hits or 0),
        }

    async def get_stats(self) -> dict[str, Any]:
        return await asyncio.to_thread(self._get_stats_sync)

    def _search_similar_sync(
        self,
        vector: list[float],
        threshold: float,
        limit: int = 1,
        scope: str | None = None,
    ) -> list[CacheEntry]:
        query = (
            f"SELECT response_data, created_at, hit_count, vector, scope, expires_at "
            f"FROM {self._table_name} WHERE vector IS NOT NULL"
        )
        params: list[Any] = []
        if scope is not None:
            query += " AND scope = %s"
            params.append(scope)

        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(query, params)
                rows = cursor.fetchall()

        results: list[tuple[float, CacheEntry]] = []
        now = time.time()
        for response_data, created_at, hit_count, stored_vec, stored_scope, expires_at in rows:
            if expires_at is not None and now > expires_at:
                continue
            if isinstance(response_data, str):
                response_data = json.loads(response_data)
            if isinstance(stored_vec, str):
                stored_vec = json.loads(stored_vec)
            sim = _cosine_similarity(vector, stored_vec)
            if sim >= threshold:
                results.append(
                    (
                        sim,
                        CacheEntry(
                            response_data=response_data,
                            created_at=created_at,
                            hit_count=hit_count,
                            vector=stored_vec,
                            scope=stored_scope,
                        ),
                    )
                )

        results.sort(key=lambda item: item[0], reverse=True)
        return [entry for _, entry in results[:limit]]

    async def search_similar(
        self,
        vector: list[float],
        threshold: float,
        limit: int = 1,
        scope: str | None = None,
    ) -> list[CacheEntry]:
        return await asyncio.to_thread(self._search_similar_sync, vector, threshold, limit, scope)
