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

class TaskModel:
    def __init__(self):
        # 初始化 MySQL 连接
        self.connection = mysql.connector.connect(**MYSQL_CONFIG)
        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):
        """
        关闭 MySQL 连接
        """
        self.cursor.close()
        self.connection.close()