from datetime import datetime

from langchain_core.messages import HumanMessage, AIMessageChunk, SystemMessage, BaseMessage

from agent.bill import bill_agent
from modes.agent import QueryRequest
from prompt.mobile import mobile_prompt


async def query_bill_agent(query: QueryRequest):
    try:
        print(query)
        today = datetime.today().strftime("%Y-%m-%d")

        messages: list[BaseMessage] = [
            SystemMessage(content=f"当前系统日期：{today}"),
            HumanMessage(content=query.message),
        ]

        if query.platform == "mobile":
            messages.insert(0, SystemMessage(content=mobile_prompt))

        tool_called = False
        content_generate = False

        # 流式调用Agent
        for chunk, metadata in bill_agent.stream(
                {"messages": messages},
                stream_mode="messages"
        ):
            # print(chunk)
            if isinstance(chunk, AIMessageChunk):
                if chunk.content:
                    if not content_generate:
                        content_generate = True
                        yield "\n 📋 **分析结果**：\n"
                    yield chunk.content

                for tc in chunk.tool_call_chunks:
                    if tc.get("args"):
                        if not tool_called:
                            tool_called = True
                            yield "🛠️ **开始查询数据**：查询日期:"
                        yield tc["args"]

    except Exception as e:
        print(f"\n[错误]: {str(e)}")
        yield "信息检索失败，请重新输入问题提问"