webapp/app/services/vas_task_service.py

# app/services/task_service.py

from datetime import datetime, date, timedelta
from typing import List, Optional

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, or_, and_


from app.utils.search import apply_keyword_search_stmt
from app.utils.pagination import paginate
from app.core.biz_exception import NotFoundError,BizLogicError
from app.models.vas_task import VasTask
from app.models.order import VasOrder
from app.schemas.vas_task import VasTaskCreate, VasTaskUpdate
from app.services.task_handlers import task_processor


class VasTaskService:

    @staticmethod
    async def create(db: AsyncSession, data: VasTaskCreate) -> VasTask:
        rec = VasTask(
            **data.dict(),
            status="pending",
            created_at=datetime.utcnow(),
        )
        db.add(rec)
        await db.commit()
        await db.refresh(rec)
        return rec
    
    @staticmethod
    async def pop_vas_task(session: AsyncSession, routing_key: str, cooldown_seconds: int = 60):
        """
        异步获取任务，支持冷却期机制。
        
        :param session: 数据库异步会话
        :param routing_key: 队列键值
        :param cooldown_seconds: 失败后的冷却时间（秒），默认60秒
        :return: VasTask 对象 or None
        """
        # 计算冷却截止时间：当前时间 - 冷却秒数
        # 只有 updated_at 早于这个时间的重试任务，才会被提取
        cutoff_time = datetime.utcnow() - timedelta(seconds=cooldown_seconds)

        try:
            # --- 构造查询语句 ---
            stmt = (
                select(VasTask)
                .where(VasTask.routing_key == routing_key)
                .where(VasTask.status == 'pending')
                # === 核心逻辑：冷却期筛选 ===
                .where(
                    or_(
                        # 情况1：这是一个全新任务 (从未尝试过)
                        VasTask.attempt_count == 0,
                        
                        # 情况2：这是一个重试任务，且距离上次更新（失败）已经过了冷却期
                        and_(
                            VasTask.attempt_count > 0,
                            VasTask.updated_at < cutoff_time
                        )
                    )
                )
                # 排序：优先级优先(0假设是最高优先级?)，其次是先创建的
                # 注意：根据你的业务，priority 可能 desc 才是高优先级，这里按 asc 写
                .order_by(VasTask.priority.desc(), VasTask.id.asc())
                .limit(1)
                # MySQL 8.0+ 必加，跳过被锁定的行
                .with_for_update(skip_locked=True)
            )

            result = await session.execute(stmt)
            task = result.scalar_one_or_none()

            # --- 更新状态 ---
            if task:
                task.status = 'running'     # 标记为已被抓取
                task.attempt_count += 1     # 增加尝试次数
                task.updated_at = datetime.utcnow() # 更新时间（重置冷却计时起点）
                
                await session.commit()
                # session.begin() 结束时自动 commit
                return task
            raise NotFoundError(message="Task not found")
        except Exception as e:
            # 记录日志
            raise e
        
    @staticmethod
    async def get_expiring_tasks(db: AsyncSession, threshold_days: int = 3):
        """
        获取即将过期或已过期的活跃任务
        :param threshold_days: 预警阈值，默认 7 天内到期
        """
        # 1. 查出所有活跃任务
        stmt = select(VasTask).where(
            VasTask.status.in_(['pending', 'running', 'grabbed'])
        )
        tasks = (await db.execute(stmt)).scalars().all()
        
        results = []
        today = date.today()

        for task in tasks:
            user_inputs = task.user_inputs or {}
            end_date_str = user_inputs.get("expected_end_date")
            
            # 如果没有截止日期，跳过
            if not end_date_str:
                continue

            try:
                # 解析日期
                end_date = datetime.strptime(end_date_str, "%Y-%m-%d").date()
                
                # 计算剩余天数
                delta = (end_date - today).days
                
                # 筛选条件: 已过期 (delta < 0) 或 即将过期 (delta <= threshold)
                if delta <= threshold_days:
                    # 获取客户姓名
                    first = user_inputs.get("first_name", "")
                    last = user_inputs.get("last_name", "")
                    email = user_inputs.get("email", "")
                    social_media_account = user_inputs.get("social_media_account", "")
                    full_name = f"{first} {last}".strip() or "未知客户"

                    results.append({
                        "id": task.id,
                        "order_id": task.order_id,
                        "routing_key": task.routing_key,
                        "status": task.status,
                        "social_media_account": social_media_account,
                        "customer_name": full_name,
                        "expected_end_date": end_date_str,
                        "email": email,
                        "days_left": delta
                    })
            except ValueError:
                continue # 日期格式错误则忽略

        # 按剩余天数升序排列 (最急的在前面)
        results.sort(key=lambda x: x["days_left"])
        return results

    @staticmethod
    async def list_task(
        db: AsyncSession,
        status: Optional[str] = None,
        routing_key: Optional[str] = None,
        script_version: Optional[str] = None,
        keyword: Optional[str] = None,
        page: int = 0,
        size: int = 10,
    ):
        stmt = select(VasTask)

        if status:
            stmt = stmt.where(VasTask.status == status)

        if routing_key:
            stmt = stmt.where(VasTask.routing_key == routing_key)

        if script_version:
            stmt = stmt.where(VasTask.script_version == script_version)

        stmt = apply_keyword_search_stmt(
            stmt=stmt,
            model=VasTask,
            keyword=keyword,
            fields=["order_id", "routing_key", "user_inputs"],
        )

        stmt = stmt.order_by(
            VasTask.priority.desc(),
            VasTask.id.asc(),
        )

        return await paginate(db, stmt, page, size)

    @staticmethod
    async def update(
        db: AsyncSession,
        id: int,
        payload: VasTaskUpdate,
    ) -> VasTask:
        stmt = select(VasTask).where(VasTask.id == id)
        result = await db.execute(stmt)
        obj = result.scalar_one_or_none()

        if not obj:
            raise NotFoundError("Task not exist")

        data = payload.dict(exclude_unset=True)

        for key, value in data.items():
            setattr(obj, key, value)

        await db.commit()
        await db.refresh(obj)
        return obj

    @staticmethod
    async def get_active_task_by_order_id(
        db: AsyncSession,
        order_id: str,
    ) -> List[VasTask]:
        stmt = select(VasTask).where(
            VasTask.status == "pending",
            VasTask.order_id == order_id,
        )
        result = await db.execute(stmt)
        return result.scalars().all()

    @staticmethod
    async def return_to_queue(db: AsyncSession, id: int) -> VasTask:
        stmt = select(VasTask).where(VasTask.id == id)
        result = await db.execute(stmt)
        rec = result.scalar_one_or_none()

        if not rec:
            raise NotFoundError("Task not exist")
        
        if rec.status == "pending":
            raise BizLogicError("Task is in queue already")
        
        rec.status = "pending"
        await db.commit()
        await db.refresh(rec)
        return rec

    @staticmethod
    async def manual_confirm(db: AsyncSession, id: int) -> VasTask:
        stmt = select(VasTask).where(VasTask.id == id)
        result = await db.execute(stmt)
        task = result.scalar_one_or_none()

        if not task:
            raise NotFoundError("Task not exist")

        task.status = "completed"

        order_stmt = select(VasOrder).where(VasOrder.id == task.order_id)
        order_result = await db.execute(order_stmt)
        order = order_result.scalar_one_or_none()

        if not order:
            raise NotFoundError("Order not exist")

        order.status = "completed"
        
        await task_processor.execute(task.routing_key, db, task, order)
        await db.commit()
        await db.refresh(task)
        return task