"""
数据分析入口：根据用户问题自动路由到 SQL Agent 或可视化 Agent。
"""

from config import llm
from agents import sql_agent, visualization_agent


def classify_intent(user_query: str) -> str:
    """用 LLM 判断用户意图：query / visualize / both。"""
    prompt = f"""判断以下用户问题属于哪种类型：
    - "query"：需要查询数据库获取数据
    - "visualize"：需要画图或做统计分析
    - "both"：需要先查数据，再画图分析

    用户问题：{user_query}

    只回复一个词：query / visualize / both"""

    return llm.invoke(prompt).content.strip().lower()


def run_data_analysis(user_query: str) -> str:
    """统一入口：自动判断意图并路由到对应 Agent。"""
    intent = classify_intent(user_query)

    if intent == "query":
        result = sql_agent.invoke({"messages": [{"role": "user", "content": user_query}]})
        return result["messages"][-1].content

    elif intent == "visualize":
        result = visualization_agent.invoke({"messages": [{"role": "user", "content": user_query}]})
        return result["messages"][-1].content

    else:  # both
        data_result = sql_agent.invoke({"messages": [{"role": "user", "content": user_query}]})
        data_content = data_result["messages"][-1].content
        viz_input = f"基于以下数据进行可视化分析：\n{data_content}\n\n原始问题：{user_query}"
        viz_result = visualization_agent.invoke({"messages": [{"role": "user", "content": viz_input}]})
        return viz_result["messages"][-1].content


if __name__ == "__main__":
    user_query = """
    列出所有2023年入职的员工，并列出他们的所属部门、姓名和薪资，按薪资从高到低排序？并画柱状图展示
    """
    result = run_data_analysis(user_query)
    print(result)