# 文件路径：personal_assistant.py
from dotenv import load_dotenv
load_dotenv()

import sqlite3
from datetime import datetime
from pydantic import BaseModel, Field

from langchain.tools import tool
from langchain.agents import create_agent
from langchain.agents.middleware import dynamic_prompt
from langchain.chat_models import init_chat_model
from langchain.messages import HumanMessage
from langgraph.checkpoint.sqlite import SqliteSaver


# ========== 1. 模拟数据 ==========

calendar_events = [
    {"id": 1, "title": "Python 学习", "date": "2024-03-25",
     "time": "14:00", "duration": "2小时"},
    {"id": 2, "title": "团队周会", "date": "2024-03-25",
     "time": "10:00", "duration": "1小时"},
    {"id": 3, "title": "代码审查", "date": "2024-03-26",
     "time": "15:00", "duration": "1.5小时"},
]

weather_db = {
    "杭州": {"condition": "晴", "temp": 25, "humidity": 60},
    "北京": {"condition": "多云", "temp": 18, "humidity": 45},
    "上海": {"condition": "小雨", "temp": 22, "humidity": 80},
}


# ========== 2. 定义工具 ==========

@tool
def get_weather(city: str) -> str:
    """查询指定城市的实时天气。

    Args:
        city: 城市名称，如 杭州、北京、上海
    """
    data = weather_db.get(city)
    if not data:
        return f"暂不支持查询 {city} 的天气。支持的城市：{', '.join(weather_db.keys())}"
    return (f"{city}天气：{data['condition']}，"
            f"温度 {data['temp']}°C，湿度 {data['humidity']}%")


@tool
def query_schedule(date: str = None) -> str:
    """查询指定日期的日程安排。不指定日期则查询今天的日程。

    Args:
        date: 日期，格式 YYYY-MM-DD，如 2024-03-25。不传则查询今天
    """
    if date is None:
        date = datetime.now().strftime("%Y-%m-%d")

    events = [e for e in calendar_events if e["date"] == date]
    if not events:
        return f"{date} 没有日程安排。"

    events.sort(key=lambda e: e["time"])
    # 这里必须直接写 emoji 字符，不能写成 📅 这种 HTML 实体——
    # 这段是 Python 字符串，会被原样打印出来，不会被浏览器解析成图形
    lines = [f"📅 {date} 日程安排："]
    for e in events:
        lines.append(f"  - {e['time']} {e['title']}（{e['duration']}）")
    return "\n".join(lines)


@tool
def send_email(to: str, subject: str, body: str) -> str:
    """发送邮件（模拟）。

    Args:
        to: 收件人邮箱
        subject: 邮件主题
        body: 邮件正文
    """
    # 模拟发送
    email_id = f"MSG-{datetime.now().strftime('%Y%m%d%H%M%S')}"
    return f"邮件已发送！收件人：{to}，主题：{subject}，邮件ID：{email_id}"


# ========== 3. 结构化输出模型 ==========

class DailySummary(BaseModel):
    """每日摘要"""
    date: str = Field(description="日期")
    weather_summary: str = Field(description="天气概述")
    event_count: int = Field(description="日程数量")
    key_events: list[str] = Field(description="重要日程列表")
    suggestion: str = Field(description="今日建议")


def to_markdown(summary: DailySummary) -> str:
    """把结构化的 DailySummary 渲染成 Markdown 文本，
    对应系统设计里"结构化输出：日程汇总格式化为 Markdown"这一项"""
    lines = [
        f"## {summary.date} 今日摘要",
        "",
        f"- **天气**：{summary.weather_summary}",
        f"- **日程数量**：{summary.event_count}",
        "",
        "**重要事项：**",
    ]
    if summary.key_events:
        lines += [f"- {event}" for event in summary.key_events]
    else:
        lines.append("- 无")
    lines += ["", f"**今日建议**：{summary.suggestion}"]
    return "\n".join(lines)


# ========== 4. 定义 Middleware ==========

@dynamic_prompt
def inject_date_context(request) -> str:
    """动态在系统提示词后追加当前日期信息。

    之前用 @before_model 把日期消息插进 messages 列表的写法有两个问题：
    1. before_model 返回的消息更新走 add_messages reducer，reducer 只按
       "新消息追加到末尾"处理，不认返回列表里的位置，insert(-1, ...)
       想插到用户消息前面的意图其实并不生效；
    2. before_model 在多轮工具调用循环里会反复触发，容易把这条消息插进
       AIMessage(tool_calls=...) 和它对应的 ToolMessage 之间，
       打乱严格的消息顺序要求。
    改用 @dynamic_prompt 直接重写系统提示词字符串，每次模型调用都会
    重新计算一遍，不触碰 messages 列表，没有累积或错位的风险。
    """
    now = datetime.now()
    weekday = ["一", "二", "三", "四", "五", "六", "日"][now.weekday()]
    date_hint = (f"\n\n[系统提示] 当前日期是 {now.strftime('%Y年%m月%d日')}，"
                 f"星期{weekday}。如果用户没有指定日期，默认查询今天。")
    return request.system_prompt + date_hint


# ========== 5. 创建 Agent ==========

# 用 SqliteSaver 持久化对话，实现系统设计里的"对话记忆"。
# 自己建立连接再传给 SqliteSaver 构造函数，而不是用
# SqliteSaver.from_conn_string()（那是个只适合 with 语句、
# 用完即关的上下文管理器，详见《LangChain 智能客服机器人》一篇）。
conn = sqlite3.connect("personal_assistant.db", check_same_thread=False)
checkpointer = SqliteSaver(conn)

model = init_chat_model("deepseek:deepseek-v4-flash", temperature=0.3)
agent = create_agent(
    model=model,
    tools=[get_weather, query_schedule, send_email],
    middleware=[inject_date_context],
    response_format=DailySummary,
    checkpointer=checkpointer,
    system_prompt="""你是个人助手"小助"。你可以查天气、管理日程、发送邮件。

## 工作方式
1. 当用户问"今天怎么样"或类似问题时：
   - 先查询今天的天气（get_weather）
   - 再查询今天的日程（query_schedule）
   - 然后生成每日摘要

2. 当用户要求发邮件时，使用 send_email 工具

3. 当用户只问天气或只问日程时，只调用对应的工具

## 风格
- 语气亲切自然
- 优先使用工具获取实时数据，不要编造""",
)


# ========== 6. 交互函数 ==========

def chat(message: str, thread_id: str = "xiaoming"):
    """与助手对话：
    - 用 stream_mode="values" 流式展示 AI 的思考和处理过程（对应"流式输出"）
    - 传入 thread_id，同一个 thread_id 下的对话会被 SqliteSaver 记住（对应"对话记忆"）
    - 如果本轮生成了结构化摘要，会渲染成 Markdown 一并打印（对应"结构化输出"）
    """
    config = {"configurable": {"thread_id": thread_id}}
    print(f"\n{'='*60}")
    print(f"你: {message}")
    print(f"{'='*60}")

    seen = 0
    final_state = {}
    for state in agent.stream(
        {"messages": [HumanMessage(content=message)]},
        config=config,
        stream_mode="values",
    ):
        final_state = state
        msgs = state.get("messages", [])
        # 每多出几条新消息，就说明 Agent 往前推进了一步，实时打印出来
        for msg in msgs[seen:]:
            if msg.type == "ai" and getattr(msg, "tool_calls", None):
                for call in msg.tool_calls:
                    print(f"🤔 决定调用工具: {call['name']}，参数: {call['args']}")
            elif msg.type == "tool":
                tool_name = getattr(msg, "name", "") or "工具"
                print(f"🔧 {tool_name} 返回: {str(msg.content)[:80]}")
            elif msg.type == "ai" and msg.content:
                print(f"🤖 助手: {msg.content}")
        seen = len(msgs)

    # 注意：response_format 是在 Agent 级别配置的，每一次调用（包括发邮件、
    # 追问这类和"日程摘要"无关的请求）都会强制尝试生成一份 DailySummary，
    # 这是全局 response_format 的已知局限。真要做成多意图助手，更好的做法是
    # 把"生成每日摘要"做成一个单独的工具，让 Agent 自己判断要不要调用，
    # 而不是给整个 Agent 挂一个一直生效的输出 schema。
    if "structured_response" in final_state:
        summary = final_state["structured_response"]
        print("\n--- 结构化摘要（Markdown） ---")
        print(to_markdown(summary))

    return final_state


# ========== 7. 测试 ==========

if __name__ == "__main__":
    chat("杭州今天天气怎么样？看看我的日程，然后给我一个今日总结")
    chat("帮我发一封邮件给 team@runoob.com，主题是'今日总结'，内容是今天日程已确认")
    # 第三轮不带任何新信息，纯粹考验 Agent 是否记得上一轮对话——
    # 因为传的是同一个 thread_id，checkpointer 会把完整历史带回来
    chat("我刚才让你发的那封邮件，主题是什么来着？")