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

from langchain_core.messages import HumanMessage, SystemMessage

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:
        intent = state["intent"]
        print(intent)
        print("=================")
        if intent == "unclear":
            llm = get_llm(Clarification)
            clarification = await llm.ainvoke(
                [
                    SystemMessage(content=CLARIFY_SYSTEM_PROMPT),
                    HumanMessage(
                        content=state["input_context"].model_dump_json(
                            exclude_unset=True
                        )
                    ),
                ]
            )
        else:
            # 需求理解已经给出问题，直接复用
            clarification = Clarification(
                question=state["requirement"].clarification_question
            )
        return {**state, "clarification": clarification}


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

    from pydantic import BaseModel

    from app.clients.decision import close_decision_engine
    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 query_app.ainvoke(state)
        finally:
            await close_decision_engine()
        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,
            )
        )

    asyncio.run(main())