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

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())
        # 暂无会话创建接口，未传会话 ID 时先分配标识，不创建数据库记录
        conversation_id = request.context.conversation_id or str(uuid4())
        payload = request.model_dump(mode="json")
        payload["context"]["conversation_id"] = conversation_id

        return await self._invoke(
            thread_id,
            {
                "request": payload,
                "conversation_id": conversation_id,
                "thread_id": thread_id,
                "message_id": str(uuid4()),
            },
            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, "请选择当前问题的有效选项")

            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 await self._execute(
                thread_id,
                Command(
                    resume=reply.model_dump(mode="json"),
                    update={"message_id": str(uuid4())},
                ),
                conversation_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(
            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"),
        )