"""
数据库连接池管理
"""

from typing import Optional

import psycopg2
from psycopg2 import pool
from psycopg2.extras import RealDictCursor


class DatabaseConnectionPool:
    """PostgreSQL 连接池"""

    _instance: Optional["DatabaseConnectionPool"] = None
    _pool: Optional[pool.ThreadedConnectionPool] = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def initialize(
        self,
        connection_string: str,
        minconn: int = 1,
        maxconn: int = 10,
    ):
        """初始化连接池"""
        if self._pool is None:
            self._pool = pool.ThreadedConnectionPool(
                minconn=minconn,
                maxconn=maxconn,
                dsn=connection_string,
            )

    def get_connection(self):
        """从连接池获取连接"""
        if self._pool is None:
            raise RuntimeError("Connection pool not initialized")
        return self._pool.getconn()

    def put_connection(self, conn):
        """将连接归还到连接池"""
        if self._pool is not None:
            self._pool.putconn(conn)

    def close_all(self):
        """关闭所有连接"""
        if self._pool is not None:
            self._pool.closeall()
            self._pool = None


# 全局连接池实例
db_pool = DatabaseConnectionPool()


class PooledConnection:
    """连接池上下文管理器"""

    def __init__(self, cursor_factory=RealDictCursor):
        self.cursor_factory = cursor_factory
        self.conn = None

    def __enter__(self):
        self.conn = db_pool.get_connection()
        if self.cursor_factory:
            self.conn.cursor_factory = self.cursor_factory
        return self.conn

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.conn:
            if exc_type is not None:
                self.conn.rollback()
            db_pool.put_connection(self.conn)
        return False


def get_db_connection():
    """获取数据库连接（用于依赖注入）"""
    return PooledConnection()