"""
author: 馒头
email: neihanshenshou@163.com
"""
from pymysql import OperationalError, ProgrammingError
from pymysql.connections import Connection
from pymysql.cursors import DictCursor


class MySQL:
    def __init__(self, user="root", password="123456", database="mantou", host="127.0.0.1", port=3306, typo=DictCursor):

        self.user = user
        self.password = password
        self.database = database
        self.host = host
        self.port = port
        self.cursor_type = typo
        self.conn = None
        self._init_conn()

    def _init_conn(self):
        """初始化数据库连接"""
        try:
            self.conn = Connection(
                user=self.user,
                password=self.password,
                host=self.host,
                port=self.port,
                database=self.database,
                cursorclass=self.cursor_type,
                autocommit=False
            )
        except (OperationalError, ProgrammingError) as e:
            raise RuntimeError(f"数据库连接失败: {e}") from e

    def __enter__(self):
        return self

    def execute(self, sql=None):
        """执行 SQL 语句，返回结果列表"""
        cursor = None
        try:
            if not self.conn or not self.conn.open:
                self._init_conn()

            cursor = self.conn.cursor()
            sql = sql.strip() if (sql and isinstance(sql, str)) else "select * from lick_dog;"
            cursor.execute(sql)

            if not sql.lower().startswith("select"):
                self.conn.commit()

            data = cursor.fetchall()
            return list(data)

        except (OperationalError, ProgrammingError) as e:
            # SQL 执行失败时回滚事务
            if self.conn:
                self.conn.rollback()
            raise RuntimeError(f"SQL 执行失败: {sql}\n错误信息: {e}") from e

        finally:
            if cursor:
                cursor.close()

    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器退出"""
        if exc_type is not None and self.conn:
            self.conn.rollback()

        if self.conn and self.conn.open:
            self.conn.close()
