import pymysql
from config import DB_CONFIG

class Db:
    def __enter__(self):
        # 使用配置文件中的信息
        self.db = pymysql.connect(**DB_CONFIG)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 关闭数据库连接
        self.db.close()

    # 插入数据
    def insert_data(self, table_name, data):
        cursor = self.db.cursor()
        sql = f"INSERT INTO {table_name} ({', '.join(data.keys())}) VALUES ({', '.join(['%s'] * len(data))})"
        cursor.execute(sql, list(data.values()))
        self.db.commit()
        inserted_id = cursor.lastrowid  # 获取插入行的 ID
        cursor.close()
        return inserted_id

    # 查询数据
    def select_data(self, table_name, columns="*", condition=None):
        cursor = self.db.cursor(pymysql.cursors.DictCursor)
        sql = f"SELECT {columns} FROM {table_name}"
        if condition:
            sql += f" WHERE {condition}"
        cursor.execute(sql)
        result = cursor.fetchall()
        cursor.close()
        return result

    # 删除数据
    def delete_data(self, table_name, condition):
        cursor = self.db.cursor()
        sql = f"DELETE FROM {table_name} WHERE {condition}"
        cursor.execute(sql)
        self.db.commit()
        cursor.close()

    # 更新数据
    def update_data(self, table_name, data, condition):
        cursor = self.db.cursor()
        set_clause = ", ".join([f"{key} = %s" for key in data.keys()])
        sql = f"UPDATE {table_name} SET {set_clause} WHERE {condition}"
        cursor.execute(sql, list(data.values()))
        self.db.commit()
        cursor.close()
    # 判断数据是否存在
    def is_data_exist(self, table_name, condition):
        cursor = self.db.cursor()
        sql = f"SELECT COUNT(*) FROM {table_name} WHERE {condition}"
        cursor.execute(sql)
        result = cursor.fetchone()
        cursor.close()
        return result[0] > 0