"""PostgreSQL向量数据库客户端实现."""

from typing import Any, Optional

from embedding_sever.db.base import VectorDBClient


class PostgresClient(VectorDBClient):
    """PostgreSQL向量数据库客户端.

    使用pgvector扩展存储向量，支持向量相似度搜索。
    需要安装pgvector扩展: CREATE EXTENSION IF NOT EXISTS vector;
    """

    def __init__(self, pg_url: Optional[str] = None):
        """初始化PG客户端.

        Args:
            pg_url: PostgreSQL连接URL，默认从配置读取
        """
        self.pg_url = pg_url
        self._pool: Optional[Any] = None

    async def connect(self) -> None:
        """建立PG连接池."""
        from psycopg_pool import AsyncConnectionPool

        self._pool = AsyncConnectionPool(
            conninfo=self.pg_url,
            min_size=1,
            max_size=10,
        )

        # 验证连接并启用pgvector
        async with self._pool.connection() as conn:
            await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
            await conn.commit()

    async def disconnect(self) -> None:
        """断开PG连接."""
        if self._pool:
            await self._pool.close()
            self._pool = None

    async def health_check(self) -> bool:
        """检查PG连接健康状态."""
        if not self._pool:
            return False
        try:
            async with self._pool.connection() as conn:
                cursor = await conn.execute("SELECT 1")
                result = await cursor.fetchone()
                return result is not None and result[0] == 1
        except Exception:
            return False


# 全局PG客户端实例
_pg_client: Optional[PostgresClient] = None


def get_pg_client() -> PostgresClient:
    """获取全局PG客户端实例（单例模式）.

    Returns:
        PostgresClient: PG客户端实例
    """
    global _pg_client
    if _pg_client is None:
        _pg_client = PostgresClient()
    return _pg_client