from mysql.connector import pooling, Error
import logging
from contextlib import contextmanager
import time

# ========== 日志配置 ==========
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s"
)
logger = logging.getLogger("MySQLHelper")

class MySQLHelper:
    def __init__(self, host, user, password, database, pool_size=5):
        """
        初始化 MySQL 连接池
        """
        try:
            self.pool = pooling.MySQLConnectionPool(
                pool_name="mypool",
                pool_size=pool_size,
                pool_reset_session=True,
                host=host,
                user=user,
                password=password,
                database=database,
                charset="utf8mb4"
            )
            logger.info("✅ MySQL 连接池已创建，大小=%s", pool_size)
        except Error as e:
            logger.error("❌ 创建连接池失败: %s", e)
            raise

    @contextmanager
    def get_conn(self):
        """
        获取连接并在使用完毕后释放
        """
        conn = None
        try:
            conn = self.pool.get_connection()
            yield conn
        except Error as e:
            logger.error("数据库连接错误: %s", e)
            raise
        finally:
            if conn:
                conn.close()

    def execute(self, sql, params=None, commit=False, retry=3):
        """
        执行 INSERT/UPDATE/DELETE 等操作
        """
        for attempt in range(1, retry + 1):
            try:
                with self.get_conn() as conn:
                    cursor = conn.cursor()
                    cursor.execute(sql, params or ())
                    if commit:
                        conn.commit()
                    affected = cursor.rowcount
                    cursor.close()
                    return affected
            except Error as e:
                logger.warning("执行 SQL 失败 (尝试 %s/%s): %s", attempt, retry, e)
                time.sleep(1)
                if attempt == retry:
                    raise

    def fetch_all(self, sql, params=None, retry=3):
        """
        查询多条记录
        """
        for attempt in range(1, retry + 1):
            try:
                with self.get_conn() as conn:
                    cursor = conn.cursor(dictionary=True)
                    cursor.execute(sql, params or ())
                    result = cursor.fetchall()
                    cursor.close()
                    return result
            except Error as e:
                logger.warning("查询失败 (尝试 %s/%s): %s", attempt, retry, e)
                time.sleep(1)
                if attempt == retry:
                    raise

    def fetch_one(self, sql, params=None, retry=3):
        """
        查询单条记录
        """
        result = self.fetch_all(sql, params, retry)
        return result[0] if result else None
    
    def bulk_insert(self, sql: str, rows: list[tuple]):
        with self.get_conn() as conn:
            cursor = conn.cursor()
            cursor.executemany(sql, rows)
            conn.commit()
            cursor.close()

    @contextmanager
    def transaction(self):
        """
        事务上下文管理器
        用法:
        with db.transaction() as cursor:
            cursor.execute(...)
            cursor.execute(...)
        """
        with self.get_conn() as conn:
            try:
                cursor = conn.cursor()
                yield cursor
                conn.commit()
            except:
                conn.rollback()
                raise
            finally:
                cursor.close()

# ================= 使用示例 =================
if __name__ == "__main__":
    db = MySQLHelper(
        host="localhost",
        user="root",
        password="123456",
        database="test_db",
        pool_size=5
    )

    # 插入数据
    db.execute("INSERT INTO users(name, age) VALUES (%s, %s)", ("Alice", 25), commit=True)

    # 查询数据
    users = db.fetch_all("SELECT * FROM users WHERE age > %s", (18,))
    logger.info("查询结果: %s", users)

    # 事务示例
    try:
        with db.transaction() as cur:
            cur.execute("UPDATE users SET age = age + 1 WHERE name = %s", ("Alice",))
            cur.execute("INSERT INTO logs(message) VALUES (%s)", ("Alice age updated",))
    except Error as e:
        logger.error("事务失败: %s", e)