from flask import Flask, Response, stream_with_context, request, jsonify
from flask_cors import CORS
import json
import time
import random
import threading
import queue
import uuid

app = Flask(__name__)
CORS(app)

# 存储每个任务的队列: {task_id: queue.Queue}
task_queues = {}


def send_to_client(task_id, message):
    """向指定任务的消息队列发送消息"""
    if task_id in task_queues:
        task_queues[task_id].put(message)


def run_background_task(task_id, params):
    """模拟后台任务执行过程"""
    try:
        # 【步骤1】初始化任务环境
        # 说明：接收传入的参数，准备执行所需的资源和上下文
        send_to_client(task_id, {
            "step": 1,
            "status": "running",
            "message": "任务开始，正在初始化...",
            "comment": "初始化：加载任务参数，准备执行环境",
            "params": params,
        })
        time.sleep(1)

        # 【步骤2】数据预处理
        # 说明：对输入数据进行校验、清洗和格式转换，为后续计算做准备
        send_to_client(task_id, {
            "step": 2,
            "status": "running",
            "message": "正在处理数据...",
            "comment": "数据预处理：校验输入格式，清洗无效数据，转换为计算所需的结构",
            "progress": 30,
        })
        time.sleep(1.5)

        # 【步骤3】核心计算循环
        # 说明：执行主要的业务逻辑，分多个子步骤迭代处理
        for i in range(1, 6):
            send_to_client(task_id, {
                "step": 3,
                "status": "running",
                "message": f"正在计算第 {i}/5 步...",
                "comment": f"核心计算（第{i}/5步）：执行子任务 {i}，处理数据分片并汇总中间结果",
                "progress": 30 + i * 10,
                "detail": f"子任务 {i} 完成",
            })
            time.sleep(0.8)

        # 【步骤4】结果聚合与格式化
        # 说明：将所有子步骤的结果合并，按照输出格式要求进行整理
        send_to_client(task_id, {
            "step": 4,
            "status": "running",
            "message": "正在生成最终结果...",
            "comment": "结果聚合：合并所有子任务的中间结果，执行最终的数据格式化",
            "progress": 90,
        })
        time.sleep(1)

        # 【步骤5】任务完成
        # 说明：输出最终结果，释放资源，标记任务状态为完成
        send_to_client(task_id, {
            "step": 5,
            "status": "completed",
            "message": "任务执行完成！",
            "comment": "完成：生成最终报告，释放临时资源，更新任务状态",
            "progress": 100,
            "result": {
                "total_items": random.randint(100, 1000),
                "success_count": random.randint(90, 100),
            },
        })

    except Exception as e:
        # 【异常处理】任务执行失败
        # 说明：捕获执行过程中的任何异常，记录错误信息并通知客户端
        send_to_client(task_id, {
            "step": -1,
            "status": "error",
            "message": f"任务执行失败: {str(e)}",
            "comment": f"异常处理：捕获到错误 [{type(e).__name__}]，详细原因: {str(e)}",
        })
    finally:
        # 【清理】无论成功或失败，都要清理资源
        # 说明：发送结束信号，释放队列占用的内存
        if task_id in task_queues:
            task_queues[task_id].put(None)  # None 表示结束信号


@app.route("/api/task", methods=["POST"])
def create_task():
    """接收 POST 请求，创建并启动后台任务"""
    data = request.json or {}
    task_id = str(uuid.uuid4())

    # 为该任务创建消息队列
    task_queues[task_id] = queue.Queue()

    # 在后台线程中执行任务
    thread = threading.Thread(
        target=run_background_task, args=(task_id, data), daemon=True
    )
    thread.start()

    return jsonify({"task_id": task_id, "message": "任务已创建"})


@app.route("/api/task/<task_id>/sse")
def task_sse(task_id):
    """SSE 端点，订阅指定任务的执行过程"""
    if task_id not in task_queues:
        return jsonify({"error": "任务不存在"}), 404

    def generate():
        q = task_queues[task_id]
        while True:
            try:
                # 阻塞等待消息，超时防止永久阻塞
                message = q.get(timeout=30)
                if message is None:  # 结束信号
                    yield "event: close\ndata: {\"status\": \"closed\"}\n\n"
                    break
                yield f"data: {json.dumps(message)}\n\n"
            except queue.Empty:
                # 超时，发送心跳
                yield ": heartbeat\n\n"
                continue

        # 清理队列
        task_queues.pop(task_id, None)

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


# 保留原有的演示端点
@app.route("/api/sse")
def sse():
    """SSE endpoint that streams data to the client"""

    def generate():
        for i in range(10):
            data = {
                "id": i,
                "message": f"Server sent event #{i}",
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                "value": round(random.uniform(0, 100), 2),
            }
            yield f"data: {json.dumps(data)}\n\n"
            time.sleep(1)

        # Send completion signal
        yield "data: {\"status\": \"completed\", \"message\": \"Stream finished\"}\n\n"
        yield "event: close\n\n"

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


@app.route("/api/sse/continuous")
def sse_continuous():
    """Continuous SSE endpoint that keeps sending data"""

    def generate():
        i = 0
        while True:
            data = {
                "id": i,
                "message": f"Continuous event #{i}",
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                "cpu_usage": round(random.uniform(10, 90), 2),
                "memory_usage": round(random.uniform(30, 80), 2),
            }
            yield f"data: {json.dumps(data)}\n\n"
            i += 1
            time.sleep(2)

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


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=True)