"""整理澄清问题，等待用户补充"""

import json

from langchain_core.exceptions import OutputParserException
from langchain_core.messages import HumanMessage, SystemMessage
from pydantic import ValidationError

from app.clients.llm import get_llm
from app.processor.base import BaseNode
from app.processor.state import QueryState
from app.prompt.clarify import CLARIFY_SYSTEM_PROMPT
from app.schemas.clarification import Clarification


class ClarifyNode(BaseNode):
    name = "clarify_node"

    async def process(self, state: QueryState) -> QueryState:
        updated_state: QueryState = state.copy()
        intent = state["intent"]

        payload = {
            "intent": intent,
            "input_context": state["input_context"].model_dump(mode="json"),
        }

        if intent != "unclear":
            # 以需求节点的问题为依据补齐选项，不重新判断意图
            payload["requirement"] = state["requirement"].model_dump(mode="json")

        llm_client = get_llm(Clarification)
        messages = [
            SystemMessage(content=CLARIFY_SYSTEM_PROMPT),
            HumanMessage(content=json.dumps(payload, ensure_ascii=False)),
        ]

        # 格式错误重试一次
        for attempt in range(2):
            try:
                clarification = await llm_client.ainvoke(messages)
                break
            except (ValidationError, OutputParserException):
                if attempt == 0:
                    # 要求重新生成完整问题，避免截断选项改变原意
                    messages.append(
                        HumanMessage(
                            content=(
                                "上次输出未通过 Clarification 格式校验，请重新生成完整结果："
                                "question 为非空问题，options 为空或含 2–3 项，"
                                "每项包含唯一 id、非空 label 和 description，最多一个 recommended=true"
                            )
                        )
                    )
                    continue
                # 连续格式失败时退回自由回答，需求阶段沿用已确定的问题
                requirement = state.get("requirement") if intent != "unclear" else None
                question = (
                    requirement.clarification_question if requirement else None
                ) or "你希望我对当前大屏做什么？可以描述想修改的内容，或想了解的问题"
                clarification = Clarification(question=question)

        updated_state["clarification"] = clarification
        # 在进入 interrupt 前保存等待状态，暂停中的节点尚不会返回更新
        updated_state["status"] = "waiting_for_input"
        return updated_state


if __name__ == "__main__":
    import asyncio
    import json

    from pydantic import BaseModel

    from app.clients.decision import close_decision_engine
    from app.observability.tracing import ainvoke_screen, flush_tracing
    from app.processor.graph import query_app

    async def main():
        # 修改请求，通过完整流程测试澄清分支
        state: QueryState = {
            "request": {
                "prompt": "这个图不太对",
                "screen": {"canvas": {"width": 1920, "height": 1080}, "nodes": []},
            }
        }
        try:
            final_state = await ainvoke_screen(
                query_app, state, thread_id="clarify-test"
            )
        finally:
            await close_decision_engine()
            flush_tracing()
        print(
            json.dumps(
                {
                    key: value.model_dump(mode="json")
                    if isinstance(value, BaseModel)
                    else value
                    for key, value in final_state.items()
                },
                ensure_ascii=False,
                indent=2,
                default=str,
            )
        )

    asyncio.run(main())