"""SSE 编码和连接生命周期，不保存全局任务队列"""

import json
from contextlib import aclosing

import anyio
from starlette.responses import StreamingResponse

from app.common.response import error_content, success_response
from app.core.exceptions import BizException, handle_unexpected_exception


def pack_event(event, content):
    return f"event: {event}\ndata: {json.dumps(content, ensure_ascii=False)}\n\n"


async def encode_events(service, execution, request):
    try:
        async with aclosing(service.stream(*execution)) as events:
            async for event, data in events:
                yield pack_event(event, success_response(data).model_dump(mode="json"))
    except BizException as exc:
        yield pack_event("error", error_content(exc.code, exc.msg))
    except Exception as exc:  # noqa: BLE001 流中异常统一封装后终止
        # 流开始后不能修改 HTTP 状态，复用统一异常映射并发送终止事件
        exc._agent_thread_id = execution[0]
        response = await handle_unexpected_exception(request, exc)
        yield pack_event("error", json.loads(response.body))


class AgentEventResponse(StreamingResponse):
    media_type = "text/event-stream"

    def __init__(self, service, execution, request):
        self.service = service
        self.execution = execution
        super().__init__(
            encode_events(service, execution, request),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )

    async def stream_response(self, send):
        try:
            await super().stream_response(send)
        finally:
            # 在消费流的同一协程关闭生成器，避免追踪上下文跨协程退出
            with anyio.CancelScope(shield=True):
                await self.body_iterator.aclose()

    async def __call__(self, scope, receive, send):
        try:
            # 同时监听断开连接，模型尚未产生下一个事件时也能及时取消
            async with anyio.create_task_group() as group:

                async def stream():
                    await self.stream_response(send)
                    group.cancel_scope.cancel()

                group.start_soon(stream)
                await self.listen_for_disconnect(receive)
                group.cancel_scope.cancel()
        finally:
            # 即使连接在生成器开始前断开，也要释放预先占用的任务
            self.service.release(self.execution[0])