"""异步执行与恢复大屏任务，当前使用单进程内存检查点"""

from contextlib import aclosing
from uuid import uuid4

from langgraph.types import Command
from loguru import logger

from app.core.exceptions import BizException
from app.observability.tracing import astream_screen
from app.schemas.agent import AgentResult, ResumeRequest
from app.schemas.request import DesignRequest


class AgentService:
    def __init__(self, graph):
        self.graph = graph
        self._active: set[str] = set()

    def _prepare_run(self, request: DesignRequest):
        # 新请求始终独立建任务，不用 conversation_id 充当检查点 ID
        thread_id = str(uuid4())
        # 暂无会话创建接口，未传会话 ID 时先分配标识，不创建数据库记录
        conversation_id = request.context.conversation_id or str(uuid4())
        payload = request.model_dump(mode="json")
        payload["context"]["conversation_id"] = conversation_id

        return (
            thread_id,
            {
                "request": payload,
                "conversation_id": conversation_id,
                "thread_id": thread_id,
                "message_id": str(uuid4()),
            },
            conversation_id,
        )

    def stream_run(self, request: DesignRequest):
        execution = self._prepare_run(request)
        self._claim(execution[0])
        return execution

    async def prepare_resume(self, request: ResumeRequest):
        # 在发送流响应头之前完成校验并占用任务
        thread_id = str(request.thread_id)
        # 检查和占用之间没有 await，同一事件循环内不会重复进入
        self._claim(thread_id)
        try:
            snapshot = await self.graph.aget_state(
                {"configurable": {"thread_id": thread_id}}
            )
            if not snapshot.values:
                raise BizException(404, "任务不存在或已过期")
            if not any(task.interrupts for task in snapshot.tasks):
                raise BizException(409, "任务当前未等待用户补充")
            # 恢复前校验选项，避免无效回答消耗一次 interrupt 恢复
            clarification = snapshot.values["clarification"]
            reply = request.reply
            if (
                not reply.cancel
                and reply.option_id is not None
                and reply.option_id not in {item.id for item in clarification.options}
            ):
                raise BizException(422, "请选择当前问题的有效选项")

            conversation_id = snapshot.values["conversation_id"]
            # 恢复同一任务不能切换会话，补充编辑上下文时也保留原会话标识
            if reply.context is not None:
                if reply.context.conversation_id not in (None, conversation_id):
                    raise BizException(422, "恢复任务时不能切换会话")
                reply = reply.model_copy(
                    update={
                        "context": reply.context.model_copy(
                            update={"conversation_id": conversation_id}
                        )
                    }
                )

            # 每次有效补充或取消生成新的输入 ID，任务和会话保持不变
            return (
                thread_id,
                Command(
                    resume=reply.model_dump(mode="json"),
                    update={"message_id": str(uuid4())},
                ),
                conversation_id,
            )
        except BaseException:
            self.release(thread_id)
            raise

    def release(self, thread_id):
        self._active.discard(thread_id)

    def _claim(self, thread_id: str) -> None:
        if thread_id in self._active:
            raise BizException(409, "任务正在执行，请勿重复提交")
        self._active.add(thread_id)

    @staticmethod
    def _result(state, thread_id):
        waiting = bool(state.get("__interrupt__"))

        # 只交付校验通过的完整动作批次，失败草稿留在后端
        return AgentResult(
            conversation_id=state["conversation_id"],
            thread_id=thread_id,
            message_id=state["message_id"],
            status="waiting_for_input" if waiting else state["status"],
            actions=state.get("actions") or []
            if state.get("validation_passed") is True and state["status"] == "completed"
            else [],
            answer=state.get("answer"),
            clarification=state.get("clarification") if waiting else None,
            failure_reason=state.get("failure_reason"),
        )

    async def stream(self, thread_id, payload, session_id):
        """消费图事件，完整状态只用于整理最终结果，不直接暴露给客户端"""
        state = None
        with logger.contextualize(thread_id=thread_id):
            async with aclosing(
                astream_screen(
                    self.graph, payload, thread_id=thread_id, session_id=session_id
                )
            ) as events:
                async for mode, data in events:
                    if mode == "values":
                        state = data
                    elif data.get("type") == "node_progress" and data.get("label"):
                        yield "progress", data
        if state is None:
            raise RuntimeError("图未返回状态")
        yield "final", self._result(state, thread_id).model_dump(mode="json")