"""
场景 2：多 Agent 共享记忆池

演示多个 Agent 共享同一个 AgentMemory 实例：
1. 使用 agent_id / team_id 实现身份隔离
2. 一个 Agent 写入，另一个 Agent 读取
3. 跨 Agent 知识共享
4. 权限控制（visibility + grant_permission）
"""

from __future__ import annotations

import tempfile
import os


def main():
    # ── 1. 初始化共享记忆池 ──────────────────────────────
    db_dir = tempfile.mkdtemp(prefix="multi_agent_")
    db_path = os.path.join(db_dir, "shared_memory.db")

    from agent_memory import AgentMemory

    # 使用 agent_id + team_id 初始化（启用多 Agent 网络）
    # 先用 researcher 身份初始化
    memory = AgentMemory(
        db_path=db_path,
        agent_id="researcher",
        team_id="team_alpha",
        enable_semantic=False,
        enable_filter=False,
        enable_dedup=False,
    )
    print(f"✅ 共享记忆池已初始化 (agent=researcher, team=team_alpha)")

    # ── 2. 注册多个 Agent ───────────────────────────────
    print("\n" + "=" * 60)
    print("👥 阶段 1：注册 Agent")
    print("=" * 60)

    # AgentMemory 初始化时已自动注册 researcher
    # 手动注册 developer 和 designer
    memory.store.register_agent(
        agent_id="developer",
        agent_name="开发者小王",
        team_id="team_alpha",
        capabilities=["coding", "debugging"],
    )
    print("  ✅ 注册 developer (开发者小王, team_alpha)")

    memory.store.register_agent(
        agent_id="designer",
        agent_name="设计师小李",
        team_id="team_alpha",
        capabilities=["design", "prototype"],
    )
    print("  ✅ 注册 designer (设计师小李, team_alpha)")

    # 注册另一个团队的 Agent
    memory.store.register_agent(
        agent_id="analyst",
        agent_name="分析师小张",
        team_id="team_beta",
        capabilities=["analysis", "reporting"],
    )
    print("  ✅ 注册 analyst (分析师小张, team_beta)")

    # 列出所有 Agent
    agents = memory.store.list_agents()
    print(f"\n  已注册 Agent 数: {len(agents)}")
    for a in agents:
        print(f"    - {a.get('agent_name', a.get('agent_id'))} (team={a.get('team_id')})")

    # ── 3. Agent 写入记忆 ───────────────────────────────
    print("\n" + "=" * 60)
    print("📝 阶段 2：Agent 写入记忆")
    print("=" * 60)

    # researcher 写入项目技术决策
    r1 = memory.remember(
        content="项目技术栈: React 18 + TypeScript + Vite",
        importance="high",
    )
    print(f"  researcher 写入: {r1.get('status')} (id={r1.get('memory_id', '')[:16]}...)")

    r2 = memory.remember(
        content="数据库选型: PostgreSQL + Redis，不用 MongoDB",
        importance="high",
    )
    print(f"  researcher 写入: {r2.get('status')}")

    # developer 写入开发规范（通过 store 直接写入，指定 owner_agent_id）
    import time
    import hashlib
    import uuid

    ts = int(time.time())
    dev_mem_id = f"mem_{uuid.uuid4().hex[:12]}"
    memory.store.insert_memory(
        memory_id=dev_mem_id,
        time_id=f"T{ts}",
        time_ts=ts,
        person_id="developer",
        nature_id="note",
        content="代码规范: 使用 ESLint + Prettier，commit 遵循 Conventional Commits",
        content_hash=hashlib.sha256("代码规范: 使用 ESLint + Prettier".encode()).hexdigest(),
        topics=["dev.standard"],
        importance="medium",
        owner_agent_id="developer",
        visibility="team",  # 团队可见
    )
    print(f"  developer 写入: 代码规范 (team 可见)")

    # designer 写入设计规范
    ts2 = int(time.time())
    design_mem_id = f"mem_{uuid.uuid4().hex[:12]}"
    memory.store.insert_memory(
        memory_id=design_mem_id,
        time_id=f"T{ts2}",
        time_ts=ts2,
        person_id="designer",
        nature_id="note",
        content="设计系统: 使用 Ant Design 5.x 组件库，主色 #1677FF",
        content_hash=hashlib.sha256("设计系统: 使用 Ant Design 5.x".encode()).hexdigest(),
        topics=["design.system"],
        importance="medium",
        owner_agent_id="designer",
        visibility="team",  # 团队可见
    )
    print(f"  designer 写入: 设计规范 (team 可见)")

    # ── 4. 跨 Agent 读取 ────────────────────────────────
    print("\n" + "=" * 60)
    print("🔍 阶段 3：跨 Agent 读取")
    print("=" * 60)

    # developer 检索技术栈信息（能读到 researcher 写的）
    results = memory.recall(
        query="技术栈",
        limit=5,
        query_agent_id="developer",
    )
    print(f"\n  developer 查询「技术栈」: {results.get('total', 0)} 条结果")
    for i, mem in enumerate(results.get("primary", []), 1):
        owner = mem.get("owner_agent_id", "?")
        content = mem.get("content", "")[:50]
        print(f"    {i}. [owner={owner}] {content}...")

    # designer 检索设计相关信息
    results2 = memory.recall(
        query="设计",
        keyword="设计",
        limit=5,
        query_agent_id="designer",
    )
    print(f"\n  designer 查询「设计」: {results2.get('total', 0)} 条结果")
    for i, mem in enumerate(results2.get("primary", []), 1):
        owner = mem.get("owner_agent_id", "?")
        content = mem.get("content", "")[:50]
        print(f"    {i}. [owner={owner}] {content}...")

    # ── 5. 权限控制演示 ─────────────────────────────────
    print("\n" + "=" * 60)
    print("🔐 阶段 4：权限控制")
    print("=" * 60)

    # researcher 写入一条 private 记忆
    ts3 = int(time.time())
    private_mem_id = f"mem_{uuid.uuid4().hex[:12]}"
    memory.store.insert_memory(
        memory_id=private_mem_id,
        time_id=f"T{ts3}",
        time_ts=ts3,
        person_id="researcher",
        nature_id="note",
        content="薪资谈判策略: 目标涨幅 30%，底线 20%",
        content_hash=hashlib.sha256("薪资谈判策略".encode()).hexdigest(),
        importance="high",
        owner_agent_id="researcher",
        visibility="private",  # 仅 owner 可见
    )
    print(f"  researcher 写入 private 记忆: 薪资谈判策略")

    # developer 尝试读取（看不到 private 记忆）
    results_private = memory.store.query_agent_memories(
        agent_id="researcher",
        query_agent_id="developer",
        team_id="team_alpha",
        include_public=True,
        keyword="薪资",
    )
    print(f"  developer 查询「薪资」: {len(results_private)} 条结果（private 不可见）")

    # researcher 自己可以看到
    results_own = memory.store.query_agent_memories(
        agent_id="researcher",
        query_agent_id="researcher",
        team_id="team_alpha",
        keyword="薪资",
    )
    print(f"  researcher 查询「薪资」: {len(results_own)} 条结果（owner 可见）")

    # ── 6. 精细权限授予 ─────────────────────────────────
    print("\n" + "=" * 60)
    print("🔑 阶段 5：精细权限授予")
    print("=" * 60)

    # researcher 授权 developer 读取 private 记忆
    granted = memory.store.grant_permission(
        memory_id=private_mem_id,
        agent_id="developer",
        granted_by="researcher",
        permission="read",
    )
    print(f"  researcher → developer 授权: {granted}")

    # developer 再次查询
    results_after_grant = memory.store.query_agent_memories(
        agent_id="researcher",
        query_agent_id="developer",
        team_id="team_alpha",
        keyword="薪资",
    )
    print(f"  developer 授权后查询「薪资」: {len(results_after_grant)} 条结果")

    # ── 7. 跨 Agent 关联 ────────────────────────────────
    print("\n" + "=" * 60)
    print("🔗 阶段 6：跨 Agent 知识关联")
    print("=" * 60)

    # developer 的记忆与 researcher 的记忆建立关联
    if results.get("primary") and r1.get("memory_id"):
        tech_mem_id = r1["memory_id"]
        memory.store.add_agent_association(
            source_agent="researcher",
            target_agent="developer",
            memory_id=tech_mem_id,
            assoc_type="shares_knowledge",
            reason="技术栈决策影响开发规范",
        )
        print(f"  ✅ 建立关联: researcher → developer (shares_knowledge)")

    # 查看关联
    associations = memory.store.get_agent_associations("developer", direction="incoming")
    print(f"  developer 收到的关联: {len(associations)} 条")
    for assoc in associations:
        print(f"    - {assoc.get('association_type')}: from {assoc.get('source_agent')}")

    # ── 清理 ──────────────────────────────────────────────
    memory.close()
    print(f"\n✅ 示例完成！临时数据库保存在: {db_dir}")


if __name__ == "__main__":
    main()
