rag/api/task_routes.py

"""
异步任务查询 API

提供任务状态的 JSON 轮询和 SSE 流式两种查询方式。

路由列表:
    GET /tasks              : 任务列表
    GET /tasks/<id>         : 任务状态（JSON）
    GET /tasks/<id>/progress: 任务进度（SSE 流式）
    GET /tasks/stats        : 任务统计

后端组调用方式:
    1. POST /sync → 返回 {"task_id": "xxx", ...}
    2. GET /tasks/xxx → 轮询状态，直到 status 为 completed 或 failed

dev-ui 调用方式:
    1. POST /sync → 返回 task_id
    2. GET /tasks/xxx/progress → SSE 流式接收进度事件
"""

import json
import time
import logging
from typing import Tuple, Any
from flask import Blueprint, request, Response, stream_with_context
from auth.gateway import require_gateway_auth
from api.response_utils import success_response, error_response
from core.status_codes import SUCCESS, NOT_FOUND

logger = logging.getLogger(__name__)

task_bp = Blueprint('tasks', __name__)


def _get_registry():
    """获取任务注册表"""
    from core.task_registry import get_registry
    return get_registry()


# ==================== 任务查询 API ====================

@task_bp.route('/tasks', methods=['GET'])
@require_gateway_auth
def list_tasks() -> Tuple[Any, int]:
    """
    获取任务列表

    查询参数:
        status: 过滤状态（running / completed / failed / pending）
        type: 过滤类型（sync / reindex / upload / batch_upload / exam_generate / exam_grade）
        limit: 返回数量限制（默认 50）

    Returns:
        {"success": true, "data": {"tasks": [...], "total": N}}
    """
    registry = _get_registry()
    registry.maybe_cleanup()

    status = request.args.get('status')
    task_type = request.args.get('type')
    limit = request.args.get('limit', 50, type=int)

    tasks = registry.list_tasks(status=status, task_type=task_type, limit=limit)
    return success_response(
        data={
            'tasks': [t.to_dict() for t in tasks],
            'total': len(tasks)
        },
        status_code=SUCCESS,
        message="查询成功"
    )


@task_bp.route('/tasks/stats', methods=['GET'])
@require_gateway_auth
def task_stats() -> Tuple[Any, int]:
    """
    获取任务统计信息

    Returns:
        {"success": true, "data": {"total": N, "by_status": {...}, "by_type": {...}}}
    """
    registry = _get_registry()
    return success_response(
        data=registry.get_stats(),
        status_code=SUCCESS,
        message="查询成功"
    )


@task_bp.route('/tasks/<task_id>', methods=['GET'])
@require_gateway_auth
def get_task(task_id: str) -> Tuple[Any, int]:
    """
    获取单个任务状态（JSON 轮询接口）

    后端组推荐使用此接口轮询任务进度。

    轮询建议:
        - 间隔 1-2 秒
        - 当 status 为 completed 或 failed 时停止轮询

    Returns:
        成功: {"success": true, "data": {task详情}}
        未找到: {"error": "任务不存在"}
    """
    registry = _get_registry()
    task = registry.get_task(task_id)

    if not task:
        return error_response(
            "TASK_NOT_FOUND", NOT_FOUND,
            f"任务不存在: {task_id}",
            http_status=404
        )

    return success_response(
        data=task.to_dict(),
        status_code=SUCCESS,
        message="查询成功"
    )


@task_bp.route('/tasks/<task_id>/progress', methods=['GET'])
@require_gateway_auth
def task_progress_stream(task_id: str):
    """
    SSE 流式任务进度推送

    dev-ui 前端推荐使用此接口，实时接收任务进度事件。

    SSE 事件类型:
        - start: 任务开始
        - progress: 进度更新（含 progress/current/total/stage/message）
        - complete: 任务完成（含完整结果）
        - error: 任务失败（含错误信息）
        - heartbeat: 每 15 秒发送一次保活

    Returns:
        text/event-stream
    """
    registry = _get_registry()
    task = registry.get_task(task_id)

    if not task:
        return error_response(
            "TASK_NOT_FOUND", NOT_FOUND,
            f"任务不存在: {task_id}",
            http_status=404
        )

    def generate_sse():
        """SSE 生成器"""
        sent_complete = False

        while True:
            # 取出缓冲区事件
            events = task.drain_events()

            for event in events:
                yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"
                if event.get('type') in ('complete', 'error'):
                    sent_complete = True

            # 如果任务已结束且事件已发送完毕，退出
            if sent_complete:
                break

            # 如果任务已完成但没有事件，发送最终状态后退出
            if task.status in ('completed', 'failed') and not events:
                final_event = {
                    'type': 'complete' if task.status == 'completed' else 'error',
                    'data': task.to_dict()
                }
                yield f"data: {json.dumps(final_event, ensure_ascii=False)}\n\n"
                break

            # 心跳保活
            yield f": heartbeat\n\n"
            time.sleep(1)

    return Response(
        stream_with_context(generate_sse()),
        mimetype='text/event-stream',
        headers={
            'Cache-Control': 'no-cache',
            'X-Accel-Buffering': 'no',
            'Connection': 'keep-alive'
        }
    )