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 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
        
        # 1. 查询 workflow_daily_quota 表，查找今天的配额记录
        sql = """
            SELECT quota_limit, used_count, locked 
            FROM workflow_daily_quota 
            WHERE quota_date = %s AND scene = %s AND capability = %s
        """
        quota_record = mysql.fetchone(sql, (quota_date, scene_value, capability_value))
        
        # 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)
            """
            mysql.execute(insert_sql, (scene_value, capability_value, quota_date, 10, 0, 0))
            quota_limit = 10
        else:
            quota_limit = quota_record.get('quota_limit', 10)
            # 如果配额被锁定，直接返回 False
            if quota_record.get('locked', 0) == 1 or quota_record.get('used_count', 0) >= quota_record.get('quota_limit', 10):
                logger.warning(f"配额已锁定，scene={scene_value}, capability={capability_value}, date={quota_date}")
                return False
        
        # 3. 查询 workflow_task 表中今天已创建的任务数
        count_sql = """
            SELECT COUNT(*) as task_count 
            FROM workflow_task 
            WHERE scene = %s 
            AND capability = %s 
            AND DATE(created_time) = CURDATE()
        """
        count_result = mysql.fetchone(count_sql, (scene_value, capability_value))
        used_count = count_result.get('task_count', 0) if count_result else 0
        
        # 4. 如果已创建数 < quota_limit，返回 True，否则返回 False
        if used_count < quota_limit:
            logger.info(f"配额检查通过，scene={scene_value}, capability={capability_value}, used={used_count}, limit={quota_limit}")
            return True
        else:
            logger.warning(f"配额不足，scene={scene_value}, capability={capability_value}, used={used_count}, limit={quota_limit}")
            return False
            
    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)}"
        }