from deepagents.backends import CompositeBackend, StateBackend, StoreBackend


def make_backend(runtime):
    return CompositeBackend(
        default=StateBackend(runtime),
        routes={
            # 只有这个目录跨线程持久化；其他文件仍属于当前线程。
            "/memory-files/": StoreBackend(runtime),
        },
    )

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal
from uuid import uuid4

from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, StateBackend, StoreBackend
from langchain.tools import ToolRuntime, tool
from langgraph.store.memory import InMemoryStore


@dataclass
class AppContext:
    # 三个字段都应由已认证的服务端会话填充，而不是由模型生成。
    tenant_id: str
    user_id: str
    agent_id: str


store = InMemoryStore()


def ns(runtime: ToolRuntime[AppContext], scope: str) -> tuple[str, ...]:
    ctx = runtime.context
    if scope == "agent":
        return ("tenants", ctx.tenant_id, "agents", ctx.agent_id, "memories")
    if scope == "user":
        return ("tenants", ctx.tenant_id, "users", ctx.user_id, "memories")
    if scope == "org_proposals":
        return ("tenants", ctx.tenant_id, "org-memory-proposals")
    raise ValueError(f"unsupported scope: {scope}")


@tool
def remember(
    scope: Literal["agent", "user"],
    key: str,
    value: str,
    runtime: ToolRuntime[AppContext],
) -> str:
    """保存已确认的 Agent 经验或用户偏好；不要保存秘密和一次性参数。"""
    if len(key) > 80 or len(value) > 2000:
        return "拒绝写入：key 或 value 超过长度限制"
    runtime.store.put(
        ns(runtime, scope),
        key,
        {
            "value": value,
            "scope": scope,
            "status": "active",
            "source": "confirmed-conversation",
        },
    )
    return f"已保存 {scope} memory: {key}"


@tool
def recall(
    scope: Literal["agent", "user"],
    query: str,
    runtime: ToolRuntime[AppContext],
) -> str:
    """在当前租户内检索 Agent 或用户长期记忆。"""
    # 本地示例没有配置 embedding，因此先列出再做关键词过滤。
    # 生产环境可给 Store 配置索引，并改成 search(..., query=query)。
    candidates = runtime.store.search(ns(runtime, scope), limit=100)
    query_lower = query.lower()
    items = [
        item for item in candidates
        if query_lower in item.key.lower()
        or query_lower in item.value["value"].lower()
    ][:8]
    if not items:
        return "没有匹配的长期记忆"
    return "\n".join(f"- {item.key}: {item.value['value']}" for item in items)


@tool
def propose_org_memory(
    statement: str,
    evidence_uri: str,
    runtime: ToolRuntime[AppContext],
) -> str:
    """提交组织记忆候选项；此工具不会修改已发布组织记忆。"""
    proposal_id = str(uuid4())
    runtime.store.put(
        ns(runtime, "org_proposals"),
        proposal_id,
        {
            "statement": statement,
            "evidence_uri": evidence_uri,
            "proposed_by": runtime.context.user_id,
            "status": "pending",
        },
    )
    return f"候选项 {proposal_id} 已进入审核队列"


# 为了让示例可以独立运行，这里使用一个只读快照。
# 生产环境可在这个函数内部查询向量/关键词索引，或读取 MCP Resource。
KB_SNAPSHOT = (
    {"id": "leave-policy@2026-07", "text": "年假申请至少提前两个工作日提交。"},
    {"id": "expense-policy@2026-06", "text": "单笔差旅住宿超过标准时需要直属经理审批。"},
)


@tool
def search_kb(query: str) -> str:
    """搜索已发布的只读企业知识；返回稳定文档 ID 以便引用。"""
    # 演示使用固定领域词；生产环境在此接入全文/向量索引或 MCP Resource。
    terms = ("年假", "申请", "差旅", "住宿", "标准", "审批")
    hits = [
        doc for doc in KB_SNAPSHOT
        if any(term in query and term in doc["text"] for term in terms)
    ]
    if not hits:
        return "未找到；不要根据 memory 猜测正式制度"
    return "\n".join(f"[{doc['id']}] {doc['text']}" for doc in hits)


def make_backend(runtime):
    return CompositeBackend(
        default=StateBackend(runtime),
        routes={"/memory-files/": StoreBackend(runtime)},
    )


agent = create_deep_agent(
    model="openai:gpt-5.2",
    context_schema=AppContext,
    store=store,
    backend=make_backend,
    tools=[remember, recall, propose_org_memory, search_kb],
    # 共享 Skill 只描述流程；它不因此获得知识库写权限。
    skills=["./skills/company-assistant/"],
    system_prompt="""
你是企业助理。正式制度必须先调用 search_kb，并在答案中引用文档 ID。
memory 只能用于个性化和补充背景，不能覆盖正式知识。
组织级结论只能调用 propose_org_memory 提议，不能宣称已经发布。
只有用户明确确认长期有效的信息才能调用 remember。
""",
)


result = agent.invoke(
    {"messages": [{"role": "user", "content": "记住我喜欢中文表格，然后查询年假申请要求"}]},
    context=AppContext(
        tenant_id="acme",
        user_id="u-1042",
        agent_id="enterprise-assistant",
    ),
)
print(result["messages"][-1].content)