import os
import mysql.connector
from mysql.connector import Error


class MySQLDB:
    def __init__(self, host=None, port=None, username=None, password=None, database=None):
        self.host = host or os.getenv('MYSQL_HOST', 'localhost')
        self.port = port or int(os.getenv('MYSQL_PORT', '3306'))
        self.username = username or os.getenv('MYSQL_USER', 'myquant')
        self.password = password or os.getenv('MYSQL_PASSWORD', '')
        self.database = database or os.getenv('MYSQL_DATABASE', 'myquant')
        self.connection = None
        self.connect()

    def connect(self):
        """连接数据库"""
        try:
            self.connection = mysql.connector.connect(
                host=self.host,
                port=self.port,
                user=self.username,
                password=self.password,
                database=self.database
            )
            if self.connection.is_connected():
                print("成功连接到MySQL数据库")
        except Error as e:
            print(f"连接错误: {e}")

    def insert_data(self, table, data):
        """插入数据"""
        try:
            cursor = self.connection.cursor()
            columns = ', '.join(data.keys())
            placeholders = ', '.join(['%s'] * len(data))
            query = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"

            cursor.execute(query, tuple(data.values()))
            self.connection.commit()
            print(f"成功插入数据，影响行数: {cursor.rowcount}")
            return cursor.lastrowid
        except Error as e:
            print(f"插入错误: {e}")
            return None
        finally:
            if cursor:
                cursor.close()

    def query_data(self, table, columns="*", where=None, params=None):
        """查询数据"""
        try:
            cursor = self.connection.cursor(dictionary=True)

            query = f"SELECT {columns} FROM {table}"
            if where:
                query += f" WHERE {where}"
            print(query)
            cursor.execute(query, params or ())
            result = cursor.fetchall()
            return result
        except Error as e:
            print(f"查询错误: {e}")
            return []
        finally:
            if cursor:
                cursor.close()

    def update_data(self, table, data, where, params=None):
        """更新数据"""
        try:
            cursor = self.connection.cursor()

            set_clause = ', '.join([f"{key} = %s" for key in data.keys()])
            query = f"UPDATE {table} SET {set_clause} WHERE {where}"

            all_params = tuple(data.values()) + (params if params else ())

            cursor.execute(query, all_params)
            self.connection.commit()
            print(f"成功更新数据，影响行数: {cursor.rowcount}")
            return cursor.rowcount
        except Error as e:
            print(f"更新错误: {e}")
            return 0
        finally:
            if cursor:
                cursor.close()

    def close(self):
        """关闭数据库连接"""
        if self.connection and self.connection.is_connected():
            self.connection.close()
            print("数据库连接已关闭")