AgentEvalTool/backend/agenteval/intelligent_eval/task_queue.py

"""Task queue for intelligent evaluations (任务队列).

Platform scans executing evals every minute and enqueues tasks for
OpenClaw workers to pick up. Tasks are prioritized by:
1. Time slot due (时段到期)
2. Session deficit (欠账多)
3. Wait time (等待时间长)
"""

from datetime import timedelta
from typing import Optional

from sqlmodel import Session, select

from agenteval.intelligent_eval.models import IntelligentEvalStatus
from agenteval.storage.db import (
    IntelligentEvalDB,
    IntelligentEvalSessionDB,
    IntelligentEvalTaskQueueDB,
    as_utc,
    utc_now,
)


def _is_slot_due(slot: dict, current_offset: timedelta) -> bool:
    """Check if a time slot is due (时段到期).

    Args:
        slot: Time slot from plan.time_distribution (e.g., {"time_slot": "8-10h", "sessions": 2})
        current_offset: Time elapsed since eval started

    Returns:
        True if the slot's start time has passed
    """
    time_slot = slot.get("time_slot", "")
    if not time_slot:
        return False

    # Parse time slot (e.g., "8-10h" -> 8 hours)
    try:
        start_hour = int(time_slot.split("-")[0].replace("h", ""))
        slot_start = timedelta(hours=start_hour)
        return current_offset >= slot_start
    except (ValueError, IndexError):
        return False


def _calculate_session_deficit(eval_db: IntelligentEvalDB, session: Session) -> int:
    """Calculate session deficit (欠账).

    Returns:
        Number of sessions that should exist but don't
    """
    if not eval_db.plan:
        return 0

    plan = eval_db.get_plan()
    time_distribution = plan.get("time_distribution", [])

    # Count current sessions
    current_sessions = session.exec(
        select(IntelligentEvalSessionDB).where(IntelligentEvalSessionDB.eval_id == eval_db.id)
    ).all()
    current_count = len(current_sessions)

    # Calculate how many sessions should exist by now
    current_offset = utc_now() - as_utc(eval_db.started_at) if eval_db.started_at else timedelta(0)
    should_have = 0

    for slot in time_distribution:
        if _is_slot_due(slot, current_offset):
            should_have += slot.get("sessions", 0)

    # Deficit = should have - current
    deficit = max(0, should_have - current_count)
    return deficit


def _calculate_priority(eval_db: IntelligentEvalDB, session: Session) -> int:
    """Calculate task priority (越小越优先).

    Priority rules:
    - Base priority: 100
    - Time slot due: -50
    - Session deficit: -10 per session
    - Wait time: -1 per 10 minutes (max -20)
    """
    priority = 100

    # Check if any time slot is due
    if eval_db.plan and eval_db.started_at:
        plan = eval_db.get_plan()
        time_distribution = plan.get("time_distribution", [])
        current_offset = utc_now() - as_utc(eval_db.started_at)

        for slot in time_distribution:
            if _is_slot_due(slot, current_offset):
                priority -= 50
                break

    # Session deficit
    deficit = _calculate_session_deficit(eval_db, session)
    priority -= deficit * 10

    # Wait time
    if eval_db.started_at:
        wait_minutes = (utc_now() - as_utc(eval_db.started_at)).total_seconds() / 60
        priority -= min(int(wait_minutes / 10), 20)

    return max(priority, 1)


def _get_attention_reason(eval_db: IntelligentEvalDB, session: Session) -> Optional[str]:
    """Determine why this eval needs attention.

    Returns:
        Reason string, or None if no attention needed
    """
    if not eval_db.plan or not eval_db.started_at:
        return None

    plan = eval_db.get_plan()
    time_distribution = plan.get("time_distribution", [])
    current_offset = utc_now() - as_utc(eval_db.started_at)

    # Check if any time slot is due
    for slot in time_distribution:
        if _is_slot_due(slot, current_offset):
            deficit = _calculate_session_deficit(eval_db, session)
            if deficit > 0:
                return "slot_due"

    # Check if all sessions completed (need analysis)
    sessions = session.exec(
        select(IntelligentEvalSessionDB).where(IntelligentEvalSessionDB.eval_id == eval_db.id)
    ).all()

    if sessions and all(s.status == "completed" for s in sessions):
        estimated_sessions = plan.get("estimated_sessions", 0)
        if len(sessions) >= estimated_sessions:
            return "all_sessions_completed"

    return None


def _has_pending_task(eval_id: str, session: Session) -> bool:
    """Check if eval already has a pending task (去重)."""
    existing = session.exec(
        select(IntelligentEvalTaskQueueDB).where(
            IntelligentEvalTaskQueueDB.eval_id == eval_id,
            IntelligentEvalTaskQueueDB.status == "pending",
        )
    ).first()
    return existing is not None


def scan_and_enqueue_tasks(session: Session) -> int:
    """Scan all executing evals and enqueue tasks.

    Returns:
        Number of tasks enqueued
    """
    # Get all executing evals
    evals = session.exec(
        select(IntelligentEvalDB).where(IntelligentEvalDB.status == IntelligentEvalStatus.EXECUTING.value)
    ).all()

    enqueued = 0

    for eval_db in evals:
        # Check if eval needs attention
        reason = _get_attention_reason(eval_db, session)
        if reason is None:
            continue

        # Check if already has pending task (去重)
        if _has_pending_task(eval_db.id, session):
            continue

        # Calculate priority
        priority = _calculate_priority(eval_db, session)

        # Create task
        task = IntelligentEvalTaskQueueDB(
            eval_id=eval_db.id,
            status="pending",
            priority=priority,
            reason=reason,
            created_at=utc_now(),
            updated_at=utc_now(),
        )
        session.add(task)
        enqueued += 1

    session.commit()
    return enqueued


def get_next_task(session: Session) -> Optional[IntelligentEvalTaskQueueDB]:
    """Get next pending task (highest priority).

    Returns:
        Task with lowest priority value (highest priority), or None
    """
    task = session.exec(
        select(IntelligentEvalTaskQueueDB)
        .where(IntelligentEvalTaskQueueDB.status == "pending")
        .order_by(IntelligentEvalTaskQueueDB.priority, IntelligentEvalTaskQueueDB.created_at)
        .limit(1)
    ).first()
    return task


def assign_task(task_id: str, cron_id: str, session: Session) -> bool:
    """Assign a task to a cron.

    Returns:
        True if assigned successfully, False if task not found or already assigned
    """
    task = session.get(IntelligentEvalTaskQueueDB, task_id)
    if task is None or task.status != "pending":
        return False

    task.status = "assigned"
    task.assigned_cron_id = cron_id
    task.assigned_at = utc_now()
    task.updated_at = utc_now()
    session.commit()
    return True


def complete_task(task_id: str, success: bool, error: Optional[str], session: Session) -> bool:
    """Mark a task as completed or failed.

    Returns:
        True if completed successfully, False if task not found
    """
    task = session.get(IntelligentEvalTaskQueueDB, task_id)
    if task is None:
        return False

    task.status = "completed" if success else "failed"
    task.completed_at = utc_now()
    task.error = error
    task.updated_at = utc_now()
    session.commit()
    return True