import mysql.connector
from AiService.Config.Config import *
import mysql.connector.pooling

class TaskModel:
    # 创建连接池
    connection_pool = mysql.connector.pooling.MySQLConnectionPool(
        pool_name="task_pool",
        pool_size=5,  # 连接池大小
        **MYSQL_CONFIG
    )

    def __init__(self):
        # 从连接池获取连接
        self.connection = self.connection_pool.get_connection()
        self.cursor = self.connection.cursor()

    def insert_task(self, task_id: str, keyword: str):
        """
        插入任务记录到 t_gen_tasks 表
        """
        try:
            query = """
            INSERT INTO t_gen_tasks (task_id, keyword, status)
            VALUES (%s, %s, %s)
            """
            self.cursor.execute(query, (task_id, keyword, 'pending'))
            self.connection.commit()
        except Exception as e:
            print(f"Failed to insert task: {str(e)}")
            raise

    def close(self):
        """
        将连接归还到连接池
        """
        self.cursor.close()
        self.connection.close()