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

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 update_task_status(self, task_id: str, status: str, result_url: str = None, error_message: str = None):
        """
        更新任务状态
        """
        try:
            query = """
            UPDATE t_gen_tasks
            SET status = %s, result_url = %s, error_message = %s, complete_time = NOW()
            WHERE task_id = %s
            """
            self.cursor.execute(query, (status, result_url, error_message, task_id))
            self.connection.commit()
        except Exception as e:
            print(f"Failed to update task status: {str(e)}")
            raise

    def get_task_status(self, task_id: str):
        """
        查询任务状态
        """
        try:
            query = """
            SELECT task_id, keyword, status, submit_time, complete_time, result_url, error_message
            FROM t_gen_tasks
            WHERE task_id = %s
            """
            self.cursor.execute(query, (task_id,))
            result = self.cursor.fetchone()
            if result:
                return {
                    "task_id": result[0],
                    "keyword": result[1],
                    "status": result[2],
                    "submit_time": result[3],
                    "complete_time": result[4],
                    "result_url": result[5],
                    "error_message": result[6]
                }
            return None
        except Exception as e:
            print(f"Failed to get task status: {str(e)}")
            raise

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