"""问答 agent：回答关于大屏和上传数据的问题

会话里有数据文件时挂 describe_data、query_data，先查询再回答；没有文件时不挂工具，只调用一次模型
"""

from dataclasses import dataclass, field
from typing import Any

from langchain.agents import create_agent
from langchain.agents.middleware import ModelCallLimitMiddleware

from app.agent.data_tools import DATA_TOOLS, DataContext
from app.agent.tool_args import normalize_tool_args
from app.clients.llm import get_llm
from app.config import settings
from app.prompt.qa import QA_DATA_RULES, QA_SYSTEM_PROMPT


@dataclass(frozen=True)
class QAContext:
    data: DataContext = field(default_factory=DataContext)  # 本会话可用的数据表


# 是否挂数据工具 → agent
_agents: dict[bool, Any] = {}


def get_qa_agent(with_data: bool = False):
    """第一次使用时组装，之后复用；没有数据文件的请求看不到数据工具和相关提示"""
    if with_data not in _agents:
        _agents[with_data] = create_agent(
            model=get_llm(),
            tools=DATA_TOOLS if with_data else [],
            system_prompt=(
                f"{QA_SYSTEM_PROMPT}\n{QA_DATA_RULES}"
                if with_data
                else QA_SYSTEM_PROMPT
            ),
            context_schema=QAContext,
            middleware=[
                # 有工具时才挂参数兼容：挂了 wrap_tool_call 中间件，没有工具也会多出一个工具节点
                *([normalize_tool_args] if with_data else []),
                # 超过上限时抛 ModelCallLimitExceededError，由节点按失败处理；
                # 写成 end 时中间件只插入一条英文说明，节点分不清是正常结束还是用尽了次数
                ModelCallLimitMiddleware(
                    run_limit=settings.max_qa_steps, exit_behavior="error"
                ),
            ],
            # 不继承外层图的检查点：会话历史只保存回答文本，不保存查询过程
            checkpointer=False,
            name="qa_agent",
        )
    return _agents[with_data]