import pymysql
from dbutils.pooled_db import PooledDB
from loguru import logger


class DatabasePool:
    def __init__(self, DB_HOST, DB_PORT, DB_DATABASE, DB_USER, DB_PASSWORD):
        self.pool = PooledDB(creator=pymysql, maxconnections=8, mincached=2, maxcached=5,
                             # maxshared=3,
                             blocking=True,
                             host=DB_HOST, port=DB_PORT, database=DB_DATABASE, user=DB_USER, password=DB_PASSWORD,
                             charset='utf8mb4', ping=1)

    def _get_connection(self):
        try:
            return self.pool.connection()
        except Exception as e:
            logger.error("Failed to get connection from pool: %s" % e)
            return None

    @staticmethod
    def to_dict(data):
        if isinstance(data, dict):
            return data
        elif hasattr(data, '__table__'):
            return {column.name: getattr(data, column.name) for column in data.__table__.columns if
                    getattr(data, column.name) is not None}
        else:
            raise TypeError("Unsupported data type for conversion to dictionary.")

    def _prepare_data_list(self, data_dict, data_list):
        if data_dict is not None:
            return [self.to_dict(data_dict)]
        elif data_list is not None:
            return [self.to_dict(item) for item in data_list]
        else:
            raise ValueError("Either data_dict or data_list must be provided.")

    def query_sql(self, sql):
        """执行查询语句（SELECT），返回 list[dict]；失败返回 None。"""
        conn = self._get_connection()
        if not conn:
            return None
        try:
            with conn.cursor() as cursor:
                result = cursor.execute(sql)
                if not result:
                    return []
                rows = cursor.fetchall()
                columns = [column[0] for column in cursor.description]
                return [dict(zip(columns, row)) for row in rows]
        except Exception as e:
            conn.rollback()
            logger.error(f"SQL执行失败：{e}")
            return None
        finally:
            conn.close()

    def execute_many_or_loop(self, table, sql, values, return_ids=False, operation='insert'):
        """
        执行批量操作或循环操作以返回ID。
        :param table: 表名。
        :param sql: SQL语句模板。
        :param values: 值列表，每个元素都是一个记录的值。
        :param return_ids: 是否返回每条记录的ID。
        :param operation: 操作类型，'insert' or 'update'。
        :return: 如果return_ids为True，则返回ID列表；否则返回None。
        """
        ids = []
        query_result = False
        conn = self._get_connection()
        if conn is None:
            return ids
        try:
            with conn.cursor() as cursor:
                if return_ids and operation == 'insert':
                    for value in values:
                        try:
                            cursor.execute(sql, value)
                            conn.commit()
                            provided_id = value.get("id") if "id" in value else cursor.lastrowid
                            ids.append(provided_id)
                            logger.info(f"{operation} success【{table}】-> {ids}")
                        except cursor.IntegrityError as err:
                            pass
                        except Exception as e:
                            conn.rollback()
                            logger.error(f"{operation} failed: {e} for {value}")
                elif operation == "query":
                    cursor.execute(sql, values)
                    query_result = cursor.fetchall()
                    columns = [column[0] for column in cursor.description]
                    rows_as_dict = [dict(zip(columns, row)) for row in query_result]
                    return rows_as_dict
                else:
                    if values:
                        cursor.executemany(sql, values)
                        logger.info(f"{operation} success【{table}】-> {values[0]} ")
                        conn.commit()
                    else:
                        cursor.execute(sql)
                        conn.commit()
        except cursor.IntegrityError as err:
            pass
        except Exception as e:
            if str(e) == "数据重复":
                raise Exception("数据重复")
            else:
                conn.rollback()
                logger.error(f"Database operation failed:{e}")
        finally:
            conn.close()
            if query_result:
                return query_result
            return ids if return_ids and ids else None

    def _get_ids_for_update(self, table, condition):
        condition_parts = ' AND '.join([f"{k} = %s" for k in condition.keys()])
        sql = f"SELECT id FROM {table} WHERE {condition_parts}"
        values = list(condition.values())
        ids = None
        conn = self._get_connection()
        if conn is None:
            return ids
        try:
            with conn.cursor() as cursor:
                cursor.execute(sql, values)
                ids = [row[0] for row in cursor.fetchall()]
        except Exception as e:
            logger.error("Failed to get IDs for update: %s", e)
        finally:
            conn.close()
            return ids

    def insert(self, table, data_dict: dict = None, data_list=None, return_ids: bool = False):
        """
            插入mysql数据库。
            :param table: 表名。
            :param data_dict: 单个插入数据 str。
            :param data_list: 批量插入数据 list。
            :param return_ids: 是否返回每条记录的ID bool。
            :return: 如果return_ids为True，则返回ID列表；否则返回None。
        """
        ids = None
        try:
            data_dicts = self._prepare_data_list(data_dict, data_list)
            keys = data_dicts[0].keys()
            values = [[item[key] for key in keys] for item in data_dicts]  # values只是元素列表
            placeholders = ', '.join(['%s'] * len(keys))
            columns = ', '.join(keys)
            sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
            ids = self.execute_many_or_loop(table, sql, values, return_ids, 'insert')
        except Exception as E:
            logger.error("insert failed: %s" % E)
        finally:
            return ids

    def insert_exception(self, table, data_dict: dict = None, data_list=None):
        data_dicts = self._prepare_data_list(data_dict, data_list)
        keys = data_dicts[0].keys()
        values = [[item[key] for key in keys] for item in data_dicts]
        placeholders = ', '.join(['%s'] * len(keys))
        columns = ', '.join(keys)
        sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
        conn = self._get_connection()
        try:
            with conn.cursor() as cursor:
                for value in values:
                    cursor.execute(sql, value)
                    conn.commit()
        finally:
            conn.close()

    def update(self, table, condition, data_dict=None, data_list=None, return_ids=False, update_all=False):
        """
        优化后：支持模糊查询 + 新增update_all参数控制更新范围（保留初始简洁风格）
        :param table: 表名
        :param condition: 更新条件（dict）：- 精确匹配：{'url': 'twzdsj'} → url = 'twzdsj'
        - 模糊匹配：{'url__like': 'twzdsj'} → url LIKE '%twzdsj%'
        :param data_dict: 单个更新数据（dict）
        :param data_list: 批量更新数据（list[dict]）
        :param return_ids: 是否返回更新的ID列表
        :param update_all: 是否更新所有符合条件的记录（True=全部，False=仅1条，默认False）
        :return: 更新的ID列表（return_ids=True时）/None
        """
        ids = None
        if not condition:
            logger.error("更新条件不能为空，禁止全表更新")
            return None
        if return_ids:
            ids = self._get_ids_for_update(table, condition)
            if not ids:
                logger.warning(f"update【{table}】未找到数据，条件：{condition}")
                return None
        try:
            data_dicts = self._prepare_data_list(data_dict, data_list)
            if not data_dicts:
                logger.error("更新数据不能为空")
                return None
            cond_parts, cond_values = [], []
            for k, v in condition.items():
                if '__like' in k:
                    cond_parts.append(f"{k.replace('__like', '')} LIKE %s")
                    cond_values.append(f'%{v}%')
                else:
                    cond_parts.append(f"{k} = %s")
                    cond_values.append(v)
            condition_sql = ' AND '.join(cond_parts)
            set_parts = ', '.join([f"{k} = %s" for k in data_dicts[0].keys()])
            limit_clause = " LIMIT 1" if not update_all else ""
            sql = f"UPDATE {table} SET {set_parts} WHERE {condition_sql}{limit_clause}"
            values = [(list(item.values()) + cond_values) for item in data_dicts]
            self.execute_many_or_loop(table, sql, values, operation='update')
            logger.info(f"update【{table}】成功：条件={condition}，更新{'所有' if update_all else '1条'}记录")
        except Exception as E:
            logger.error(f"update failed: {E}")
            ids = None
        finally:
            return ids

    def query(self, table, condition, query_criteria=None, get_results=False, sql=None):
        """
            更新mysql数据库。
            :param table: 表名。
            :param condition: 查询字段 dict。
            :param query_criteria: 指定查询特定的值 list ,strx`x`
            :return: 如果return_ids为True，则返回ID列表；否则返回None。
        """
        query_result = False
        try:
            if not sql:
                if condition:
                    condition_str = ' AND '.join([f"{key} = %s" for key in condition.keys()])
                if query_criteria:
                    if isinstance(query_criteria, list):
                        query_criteria = query_criteria
                        key_str = ' ,'.join([key for key in query_criteria])
                    elif isinstance(query_criteria, str):
                        key_str = query_criteria
                    else:
                        key_str = ' ,'.join([key for key in condition.keys()])
                else:
                    key_str = "*"
                if condition:
                    sql = f"SELECT {key_str} FROM {table} WHERE {condition_str}"
                    values = list(condition.values())
                else:
                    sql = f"SELECT {key_str} FROM {table}"
                    values = list()
            else:
                values = list()
            query_result = self.execute_many_or_loop(table, sql, values, operation='query')
            if not get_results and query_result:
                query_result = True
        except Exception as E:
            logger.error("update failed: %s" % E)
        finally:
            return query_result

    def upsert(self, table, data_dict=None, data_list=None):
        """
        插入或更新数据库（兼容缺失字段）
        :param table: 表名
        :param data_dict: 单条数据字典
        :param data_list: 批量数据列表
        注意： data_dict和data_list 必须传入完整的数据, 否则数据库中数据会被更新为None
        """
        try:
            data_dicts = self._prepare_data_list(data_dict, data_list)
            all_keys = set()
            for item in data_dicts:
                all_keys.update(item.keys())
            all_keys = list(all_keys)
            for item in data_dicts:
                for key in all_keys:
                    if key not in item:
                        item[key] = None
            keys = all_keys
            values = [[item[key] for key in keys] for item in data_dicts]
            placeholders = ', '.join(['%s'] * len(keys))
            columns = ', '.join(keys)
            updates = ', '.join([f"{key} = VALUES({key})" for key in keys])
            sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {updates}"
            self.execute_many_or_loop(table, sql, values, operation='upsert')
            # logger.info(f"upsert【{table}】成功，处理{len(data_dicts)}条数据")
        except Exception as E:
            logger.error(f"Upsert failed: {E}")

    def query_count(self, table):
        sql = f"SELECT COUNT(*) FROM {table}"
        query_result = self.execute_many_or_loop(table, sql, list(), operation='query')
        if query_result:
            data = list(query_result[0].values())[0]
        else:
            data = 0
        return data


