"""学习笔记助手：整合 Prompt + Memory + Tool + RAG + Skill (+ 可选 MCP)"""
import json
import os
import re
import sqlite3
import datetime
from openai import OpenAI

client = OpenAI()  # 环境变量 OPENAI_API_KEY

# ════════════════════════════════════════════
# 1. 短期记忆：滑动窗口（第二篇）
# ════════════════════════════════════════════
class ShortMemory:
    def __init__(self, max_messages: int = 12):
        self.history = []
        self.max_messages = max_messages

    def add(self, role: str, content: str):
        self.history.append({"role": role, "content": content})
        self.history = self.history[-self.max_messages:]

# ════════════════════════════════════════════
# 2. 长期记忆：SQLite（第二篇）
# ════════════════════════════════════════════
class LongMemory:
    def __init__(self, db_path: str = "memory.db"):
        self.conn = sqlite3.connect(db_path)
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS facts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                content TEXT NOT NULL,
                kind TEXT DEFAULT "fact",
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
            )
        """)
        self.conn.commit()

    def add(self, content: str, kind: str = "fact"):
        self.conn.execute("INSERT INTO facts (content, kind) VALUES (?, ?)", (content, kind))
        self.conn.commit()

    def all(self) -> list:
        cur = self.conn.execute("SELECT content, kind FROM facts ORDER BY id DESC LIMIT 20")
        return [f"[{kind}] {content}" for content, kind in cur.fetchall()]

# ════════════════════════════════════════════
# 3. RAG：Chroma 检索笔记（第四篇）
# ════════════════════════════════════════════
import chromadb

class NoteRAG:
    def __init__(self, notes_dir: str = "notes"):
        self.notes_dir = notes_dir
        self.client = chromadb.PersistentClient(path="./rag_db")
        self.collection = self.client.get_or_create_collection("notes")

    def index_notes(self):
        """把 notes/ 下的 Markdown 全部切块入库（资料更新后调用）"""
        chunks, ids, metas = [], [], []
        for idx, fname in enumerate(os.listdir(self.notes_dir)):
            if not fname.endswith(".md"):
                continue
            path = os.path.join(self.notes_dir, fname)
            text = open(path, encoding="utf-8").read()
            # 简单按 500 字切块
            for j in range(0, len(text), 500):
                chunks.append(text[j:j + 500])
                ids.append(f"{fname}-{j}")
                metas.append({"source": fname})
        if chunks:
            self.collection.upsert(ids=ids, documents=chunks, metadatas=metas)
        print(f"已索引 {len(chunks)} 个文本块")

    def search(self, query: str, top_k: int = 3) -> str:
        results = self.collection.query(query_texts=[query], n_results=top_k)
        docs = results["documents"][0] if results["documents"] else []
        return "\n\n".join(f"[资料{i + 1}] {d}" for i, d in enumerate(docs))

# ════════════════════════════════════════════
# 4. 工具层（第三篇）
# ════════════════════════════════════════════
def calculator(expr: str) -> str:
    allowed = set("0123456789+-*/(). ")
    if not set(expr).issubset(allowed):
        return "错误：表达式包含非法字符"
    try:
        return str(eval(expr, {"__builtins__": {}}, {}))
    except Exception as e:
        return f"计算失败：{e}"

def save_note(title: str, content: str) -> str:
    """把学习笔记保存到 notes/ 目录，并加入知识库"""
    path = f"notes/{title}.md"
    with open(path, "w", encoding="utf-8") as f:
        f.write(f"# {title}\n\n{content}")
    rag.index_notes()  # 新笔记立即进知识库
    return f"已保存到 {path}"

def web_search(query: str) -> str:
    """演示用搜索（真实项目请替换为搜索 API 或爬虫）"""
    return f"【模拟搜索结果】关于“{query}”：暂无实时数据，建议访问搜索引擎确认。"

def get_time() -> str:
    return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

TOOLS = [
    {"type": "function", "function": {"name": "calculator", "description": "当用户需要数学计算时使用", "parameters": {"type": "object", "properties": {"expr": {"type": "string"}}, "required": ["expr"]}}},
    {"type": "function", "function": {"name": "save_note", "description": "当用户要保存学习笔记时使用", "parameters": {"type": "object", "properties": {"title": {"type": "string"}, "content": {"type": "string"}}, "required": ["title", "content"]}}},
    {"type": "function", "function": {"name": "web_search", "description": "当用户需要实时信息时使用", "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}}},
    {"type": "function", "function": {"name": "get_time", "description": "查询当前日期时间", "parameters": {"type": "object", "properties": {}}}},
]

TOOL_FUNCS = {
    "calculator": calculator,
    "save_note": save_note,
    "web_search": web_search,
    "get_time": get_time,
}

# ════════════════════════════════════════════
# 5. 技能加载器（第五篇）
# ════════════════════════════════════════════
def load_skill(skill_name: str) -> str:
    path = f"skills/{skill_name}/SKILL.md"
    if os.path.isfile(path):
        return open(path, encoding="utf-8").read()
    return ""

# ════════════════════════════════════════════
# 6. 全局状态
# ════════════════════════════════════════════
short_mem = ShortMemory()
long_mem = LongMemory()
rag = NoteRAG()

# ════════════════════════════════════════════
# 7. 系统提示词（第一篇：编排一切）
# ════════════════════════════════════════════
def build_system_prompt(task: str) -> str:
    skill_text = load_skill("blog-writing")
    facts = "；".join(long_mem.all()) if long_mem.all() else "暂无"
    return f"""你是一个学习笔记助手。

# 用户长期记忆
已知用户信息：{facts}

# 工作规则
1. 需要实时信息或计算时，调用对应工具；
2. 用户问笔记内容时，先检索知识库再回答；
3. 用户要求写博客/整理笔记成文章时，严格按技能规范执行；
4. 回答要简洁，重要信息用列表；
5. 新学到的用户偏好，回答末尾用 <remember>标签包起来，方便程序抽取。

# 写作技能（仅在写博客时生效）
{skill_text if '写' in task else '（本次任务不涉及写作技能）'}"""

def extract_remember(text: str):
    """从回复里抽取 <remember> 标签并存入长期记忆"""
    for m in re.findall(r"<remember>(.*?)</remember>", text):
        long_mem.add(m.strip(), "preference")

# ════════════════════════════════════════════
# 8. Agent 主循环
# ════════════════════════════════════════════
def run(user_input: str):
    # 先检索知识库，把相关资料放进本轮上下文
    rag_context = rag.search(user_input)
    short_mem.add("user", f"[知识库检索]\n{rag_context}\n\n用户问题：{user_input}")

    messages = [
        {"role": "system", "content": build_system_prompt(user_input)},
        *short_mem.history,
    ]

    for _ in range(8):
        resp = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=messages,
            tools=TOOLS,
        )
        msg = resp.choices[0].message

        if not msg.tool_calls:
            extract_remember(msg.content or "")
            short_mem.add("assistant", msg.content)
            return msg.content

        messages.append(msg)
        for call in msg.tool_calls:
            try:
                args = json.loads(call.function.arguments or "{}")
                result = TOOL_FUNCS[call.function.name](**args)
            except Exception as e:
                result = f"工具异常：{e}"
            print(f"  [工具] {call.function.name} → {result[:60]}")
            messages.append({"role": "tool", "tool_call_id": call.id, "content": str(result)})
    return "已达最大循环次数，请简化问题后重试。"

# ════════════════════════════════════════════
# 9. 入口
# ════════════════════════════════════════════
if __name__ == "__main__":
    rag.index_notes()
    print("学习笔记助手已启动，输入 exit 退出。")
    while True:
        q = input("你：")
        if q.strip().lower() == "exit":
            break
        print("助手：", run(q))

# notes_mcp_server.py
from mcp.server.fastmcp import FastMCP

mcp = FastMCP("notes-server")

@mcp.tool()
def search_notes(query: str, top_k: int = 3) -> str:
    """在个人学习笔记中检索相关内容"""
    # 复用上一节 NoteRAG 的逻辑（简化写法）
    import chromadb
    client = chromadb.PersistentClient(path="./rag_db")
    col = client.get_or_create_collection("notes")
    res = col.query(query_texts=[query], n_results=top_k)
    return "\n".join(res["documents"][0]) if res["documents"] else "未找到相关内容"

if __name__ == "__main__":
    mcp.run()