video_decode/tasks/decode.py

from utils.params import DecodeContentParam, SceneEnum, ContentTypeEnum, CapabilityEnum
from models.task import WorkflowTask
from models.decode_task_result import WorkflowDecodeTaskResult
from utils.sync_mysql_help import mysql
from pymysql.cursors import DictCursor
from loguru import logger
import sys
from datetime import datetime
from typing import Dict, Any, Optional

logger.add(sink=sys.stderr, level="ERROR", backtrace=True, diagnose=True)

# 错误码常量
ERROR_CODE_SUCCESS = 0
ERROR_CODE_FAILED = -1
ERROR_CODE_TASK_CREATE_FAILED = 2001


def _build_error_response(code: int, reason: str) -> Dict[str, Any]:
    """构建错误响应"""
    return {
        "code": code,
        "task_id": None,
        "reason": reason
    }


def _build_success_response(task_id: str) -> Dict[str, Any]:
    """构建成功响应"""
    return {
        "code": ERROR_CODE_SUCCESS,
        "task_id": task_id,
        "reason": ""
    }


def _create_workflow_task(scene: SceneEnum) -> Optional[WorkflowTask]:
    """创建工作流任务"""
    try:
        task = WorkflowTask.create_task(
            scene=scene,
            capability=CapabilityEnum.DECODE,
            root_task_id=''
        )
        logger.info(f"创建解构任务成功，task_id: {task.task_id}")
        return task
    except Exception as e:
        logger.error(f"创建解构任务失败: {str(e)}")
        return None


def _initialize_task_result(task_id: str, content) -> Optional[WorkflowDecodeTaskResult]:
    """初始化任务结果"""
    try:
        result = WorkflowDecodeTaskResult.create_result(
            task_id=task_id,
            content=content
        )
        logger.info(f"初始化任务结果成功，task_id: {task_id}")
        return result
    except Exception as e:
        logger.error(f"初始化任务结果失败，task_id: {task_id}, error: {str(e)}")
        return None


def _trigger_decode_workflow(task_id: str) -> bool:
    """触发解构工作流"""
    try:
        invoke_decode_workflow(task_id)
        logger.info(f"发起解构任务成功，task_id: {task_id}")
        return True
    except Exception as e:
        logger.error(f"发起解构任务失败，task_id: {task_id}, error: {str(e)}")
        return False


def _check_quota(scene: SceneEnum, capability: CapabilityEnum = CapabilityEnum.DECODE) -> bool:
    """检查配额是否充足（并发安全版本）"""
    try:
        # 获取今天的日期，格式为 YYYYMMDD
        quota_date = datetime.now().strftime('%Y%m%d')
        scene_value = scene.value
        capability_value = capability.value
        
        # 使用事务和 SELECT FOR UPDATE 确保并发安全
        pool = mysql.get_pool()
        with pool.connection() as conn:
            try:
                # 开始事务
                conn.begin()
                
                with conn.cursor(DictCursor) as cursor:
                    # 1. 使用 SELECT FOR UPDATE 锁定配额记录（如果存在）
                    select_sql = """
                        SELECT quota_limit, used_count, locked 
                        FROM workflow_daily_quota 
                        WHERE quota_date = %s AND scene = %s AND capability = %s
                        FOR UPDATE
                    """
                    cursor.execute(select_sql, (quota_date, scene_value, capability_value))
                    quota_record = cursor.fetchone()
                    
                    # 2. 如果没找到，创建一条新记录，quota_limit 默认为 10
                    if not quota_record:
                        insert_sql = """
                            INSERT INTO workflow_daily_quota (scene, capability, quota_date, quota_limit, used_count, locked)
                            VALUES (%s, %s, %s, %s, %s, %s)
                        """
                        cursor.execute(insert_sql, (scene_value, capability_value, quota_date, 10, 0, 0))
                        quota_limit = 10
                        current_used_count = 0
                        is_locked = 0
                    else:
                        quota_limit = quota_record.get('quota_limit', 10)
                        current_used_count = quota_record.get('used_count', 0)
                        is_locked = quota_record.get('locked', 0)
                    
                    # 3. 检查配额是否被锁定或已用完
                    if is_locked == 1 or current_used_count >= quota_limit:
                        conn.rollback()
                        logger.warning(f"配额已锁定或已用完，scene={scene_value}, capability={capability_value}, date={quota_date}, used={current_used_count}, limit={quota_limit}")
                        return False
                    
                    # 4. 查询 workflow_task 表中今天已创建的任务数（实际使用量）
                    count_sql = """
                        SELECT COUNT(*) as task_count 
                        FROM workflow_task 
                        WHERE scene = %s 
                        AND capability = %s 
                        AND DATE(created_time) = CURDATE()
                    """
                    cursor.execute(count_sql, (scene_value, capability_value))
                    count_result = cursor.fetchone()
                    actual_used_count = count_result.get('task_count', 0) if count_result else 0
                    
                    # 5. 如果实际使用量 >= 配额限制，返回 False
                    if actual_used_count >= quota_limit:
                        conn.rollback()
                        logger.warning(f"配额不足，scene={scene_value}, capability={capability_value}, used={actual_used_count}, limit={quota_limit}")
                        return False
                    
                    # 6. 配额充足，更新 used_count（原子操作）
                    update_sql = """
                        UPDATE workflow_daily_quota 
                        SET used_count = used_count + 1 
                        WHERE quota_date = %s AND scene = %s AND capability = %s
                        AND used_count < quota_limit AND locked = 0
                    """
                    cursor.execute(update_sql, (quota_date, scene_value, capability_value))
                    
                    # 检查是否更新成功（受影响的行数）
                    if cursor.rowcount == 0:
                        conn.rollback()
                        logger.warning(f"配额更新失败，可能已被其他请求占用，scene={scene_value}, capability={capability_value}")
                        return False
                    
                    # 提交事务
                    conn.commit()
                    logger.info(f"配额检查通过，scene={scene_value}, capability={capability_value}, used={actual_used_count + 1}, limit={quota_limit}")
                    return True
                    
            except Exception as e:
                conn.rollback()
                raise e
                
    except Exception as e:
        logger.error(f"配额检查失败: {str(e)}")
        # 发生异常时，为了不影响业务，返回 True 允许继续执行
        return True


def decode_topic(param: DecodeContentParam) -> Dict[str, Any]:
    """选题解构方法"""
    try:
        # 前置配额检查，用于超出每天解构次数时，直接返回错误
        if not _check_quota(param.scene, CapabilityEnum.DECODE):
            return _build_error_response(
                ERROR_CODE_FAILED,
                "配额不足"
            )

        # 步骤1: 创建工作流task任务
        task = _create_workflow_task(param.scene)
        if not task or not task.task_id:
            return _build_error_response(
                ERROR_CODE_TASK_CREATE_FAILED,
                "创建解构任务失败"
            )
        
        # 步骤2: 初始化任务结果
        result = _initialize_task_result(task.task_id, param.content)
        if not result or not result.task_id:
            return _build_error_response(
                ERROR_CODE_FAILED,
                "初始化任务结果失败"
            )
        
        # 步骤3: 触发解构工作流
        if not _trigger_decode_workflow(task.task_id):
            return _build_error_response(
                ERROR_CODE_FAILED,
                "发起解构任务失败"
            )
        
        # 所有步骤成功
        return _build_success_response(task.task_id)
        
    except Exception as e:
        logger.error(f"选题解构失败: {str(e)}")
        return _build_error_response(
            ERROR_CODE_TASK_CREATE_FAILED,
            f"解构任务创建失败: {str(e)}"
        )


def begin_decode_task(param: DecodeContentParam) -> Dict[str, Any]:
    """根据场景分发任务"""
    if param.scene == SceneEnum.TOPIC:
        return decode_topic(param)
    elif param.scene == SceneEnum.CREATION:
        # TODO: 实现创作场景
        return _build_error_response(ERROR_CODE_FAILED, "创作场景暂未实现")
    elif param.scene == SceneEnum.PRODUCTION:
        # TODO: 实现制作场景
        return _build_error_response(ERROR_CODE_FAILED, "制作场景暂未实现")
    else:
        return _build_error_response(ERROR_CODE_FAILED, f"未知场景: {param.scene}")


def invoke_decode_workflow(task_id: str) -> Dict[str, Any]:
    """发起解构任务"""
    try:
        # TODO: 调用实际的工作流接口
        return {
            "code": ERROR_CODE_SUCCESS,
            "task_id": task_id,
            "reason": ""
        }
    except Exception as e:
        logger.error(f"发起解构任务失败: {str(e)}")
        return {
            "code": ERROR_CODE_FAILED,
            "task_id": None,
            "reason": f"解构任务执行失败: {str(e)}"
        }