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

from uuid import uuid4

from langgraph.types import Command
from loguru import logger

from app.core.exceptions import BizException
from app.observability.tracing import ainvoke_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()

    async def run(self, request: DesignRequest) -> AgentResult:
        # 新请求始终独立建任务，不用 conversation_id 充当检查点 ID
        thread_id = str(uuid4())
        return await self._invoke(
            thread_id,
            {"request": request.model_dump(mode="json")},
            request.context.conversation_id,
        )

    async def resume(self, request: ResumeRequest) -> AgentResult:
        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, "请选择当前问题的有效选项")
            session_id = (
                snapshot.values["request"].get("context", {}).get("conversation_id")
            )
            return await self._execute(
                thread_id, Command(resume=reply.model_dump(mode="json")), session_id
            )
        finally:
            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)

    async def _invoke(self, thread_id, payload, session_id) -> AgentResult:
        self._claim(thread_id)
        try:
            return await self._execute(thread_id, payload, session_id)
        finally:
            # 异常和取消也必须释放占用，不阻塞后续请求
            self._active.discard(thread_id)

    async def _execute(self, thread_id, payload, session_id) -> AgentResult:
        # 上下文只作用于当前异步任务，节点自动继承任务标识
        with logger.contextualize(thread_id=thread_id):
            logger.info("任务执行开始")
            try:
                state = await ainvoke_screen(
                    self.graph, payload, thread_id=thread_id, session_id=session_id
                )
            except Exception as exc:
                # API 异常处理发生在上下文退出后，保留任务标识用于定位
                exc._agent_thread_id = thread_id
                raise
            logger.info(
                "任务执行结束，状态 {}",
                "waiting_for_input"
                if state.get("__interrupt__")
                else state.get("status"),
            )

        waiting = bool(state.get("__interrupt__"))

        # 只交付校验通过的完整动作批次，失败草稿留在后端
        return AgentResult(
            thread_id=thread_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"),
        )