"""NM2 试用台：直接跟「模型自带的记忆」打交道。

用法（在装了项目的环境里跑）：

    python H:\\Memory\\agent_lab\\try_nm2.py --demo        # 一键把整段故事演完
    python H:\\Memory\\agent_lab\\try_nm2.py               # 交互菜单，自己随便试

它只做一件事：把控制面的接口用中文包一层，让你看得见记忆层每一步到底做了什么 ——
路由觉得要不要用记忆、命中哪条记录、有没有注入证据、自动遗忘的概率是多少。
"""

from __future__ import annotations

import argparse
import json
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

DEFAULT_BASE = "http://127.0.0.1:8766/v1"

# 把接口里的英文状态翻成人话，看不懂的照原样显示。
STOP_REASON_ZH = {
    "evidence_found": "找到证据，已注入",
    "below_read_threshold": "分数低于读取门槛，不注入",
    "router_abstained": "路由器判定这轮不需要记忆",
    "attribute_not_covered": "库里没有这个属性（覆盖门拒答）",
    "memory_disabled": "记忆已关闭",
    "token_evidence_override": "词元证据否决，不注入",
}

STATUS_ZH = {
    "active": "生效",
    "superseded": "已被新版本取代",
    "retracted": "已撤回",
    "quarantined": "隔离中（低置信，未启用）",
}


class ServiceError(Exception):
    """控制面返回的错误，带上人话。"""


class Client:
    def __init__(self, base: str, timeout: float = 600.0) -> None:
        self.base = base.rstrip("/")
        self.timeout = timeout

    def _call(self, method: str, path: str, payload: dict | None = None, params: dict | None = None):
        url = f"{self.base}{path}"
        if params:
            url = f"{url}?{urllib.parse.urlencode(params)}"
        data = None
        headers = {"Content-Type": "application/json; charset=utf-8"}
        if payload is not None:
            data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        request = urllib.request.Request(url, data=data, headers=headers, method=method)
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                body = response.read().decode("utf-8", "replace")
        except urllib.error.HTTPError as error:
            detail = error.read().decode("utf-8", "replace")[:500]
            raise ServiceError(f"接口 {path} 返回 HTTP {error.code}：{detail}") from None
        except urllib.error.URLError as error:
            raise ServiceError(
                f"连不上 {self.base}（{error.reason}）。模型服务没在跑？"
            ) from None
        return json.loads(body) if body.strip() else {}

    # ---- 记忆层 ---------------------------------------------------------
    def health(self) -> dict:
        return self._call("GET", "/health")

    def reset(self) -> dict:
        return self._call("POST", "/nm2/reset", {})

    def write(self, *, text: str, entity: str = "", attribute: str = "", value: str = "",
              importance: float = 0.9, confidence: float = 0.95) -> dict:
        return self._call("POST", "/nm2/write", {
            "text": text, "entity": entity, "attribute": attribute, "value": value,
            "importance": importance, "confidence": confidence,
        })

    def records(self, status: str = "active", limit: int = 100) -> list[dict]:
        return self._call("GET", "/nm2/records", params={"status": status, "limit": limit}).get("records", [])

    def ask(self, question: str, *, remember: bool = True, max_tokens: int = 256) -> dict:
        """问一轮。remember=False 时这一轮不读也不写记忆（当场做对照）。"""

        body = {
            "messages": [{"role": "user", "content": question}],
            "max_tokens": max_tokens,
            "stream": False,
            "nm2": {"memory_enabled": remember},
        }
        return self._call("POST", "/chat/completions", body)


# ---------------------------------------------------------------- 显示
def rule(title: str = "") -> None:
    line = "─" * 62
    print(f"\n{line}")
    if title:
        print(title)
        print(line)


def show_answer(reply: dict) -> str:
    message = reply["choices"][0]["message"]
    content = (message.get("content") or "").strip()
    print(f"模型回答：{content or '（空）'}")
    return content


def show_decision(reply: dict) -> None:
    info = reply.get("nm2") or {}
    decisions = info.get("router_decisions") or [{}]
    decision = decisions[0] if decisions else {}
    reason = str(decision.get("stop_reason", ""))
    print("这一轮记忆层做了什么：")
    print(f"  · 要不要用记忆：{'要' if decision.get('need_memory') else '不要'}")
    print(f"  · 结论：{STOP_REASON_ZH.get(reason, reason or '（无）')}")
    print(f"  · 命中页：{', '.join(decision.get('page_ids') or []) or '（无）'}")
    print(f"  · 命中记录：{len(decision.get('record_ids') or [])} 条")
    print(f"  · 注入前缀：{info.get('prefix_tokens', 0)} 词元")
    print(f"  · 路由耗时：{round(float(info.get('read_seconds') or 0) * 1000)} 毫秒")
    forget = info.get("auto_forget_probability") or [None]
    if forget and forget[0] is not None:
        mark = "  ← 超过 0.9 会自动遗忘" if float(forget[0]) >= 0.9 else ""
        print(f"  · 自动遗忘概率：{float(forget[0]):.4f}{mark}")


def show_bank(client: Client, title: str = "记忆库现状") -> list[dict]:
    active = client.records("active")
    print(f"{title}（生效记录 {len(active)} 条）：")
    for record in active:
        origin = record.get("origin") or "（无来源标记）"
        print(f"  · [{record.get('status')}] {record.get('entity')}｜{record.get('attribute')}"
              f" = {record.get('value')}    来源轮次 {origin}")
    if not active:
        print("  （空）")
    return active


# ---------------------------------------------------------------- 演示
def demo(client: Client) -> None:
    rule("第一步：清空记忆，确保从零开始")
    client.reset()
    print("已清空。")

    rule("第二步：用户一口气说了两件事")
    sentence = "记一下：我的紧急联系人是 王工，电话分机 7781。"
    print(f"用户：{sentence}")
    reply = client.ask(sentence)
    show_answer(reply)
    show_decision(reply)

    print("\n（智能体随后按字段各写一条记录。这里直接调接口，效果等价——"
          "两条记录会同属这一轮。）")
    for attribute, value in (("紧急联系人", "王工"), ("电话分机", "7781")):
        result = client.write(
            text=f"{attribute} 是 {value}。", entity="user", attribute=attribute, value=value,
        )
        record = result.get("record", {})
        print(f"  写入 {attribute}={value} → {result.get('action')}，"
              f"来源轮次 {record.get('origin') or '（无）'}")

    show_bank(client, "现在库里")

    rule("第三步：用户要求忘掉紧急联系人信息")
    forget = "请遗忘我的紧急联系人信息，不要再保留和使用它。"
    print(f"用户：{forget}")
    reply = client.ask(forget)
    show_answer(reply)
    show_decision(reply)
    show_bank(client, "遗忘之后")

    rule("第四步：追问分机号——这里就是以前会泄露的地方")
    question = "我的紧急联系人分机是多少？如果已经不保留这条信息，就直接说明。"
    print(f"用户：{question}")
    reply = client.ask(question)
    answer = show_answer(reply)
    show_decision(reply)

    leaked = "7781" in answer
    print()
    if leaked:
        print("结果：分机号还是被说出来了 —— 这一轮没修好，把上面的路由结论发我。")
    else:
        print("结果：没有泄露分机号。同一轮写入的记录被一起撤回了。")

    rule("第五步：换成「关掉记忆」再问一次，看对照组会怎样")
    reply = client.ask(question, remember=False)
    show_answer(reply)
    print("（这一轮模型没有任何记忆可用，只能靠猜或明说不知道。）")


# ---------------------------------------------------------------- 交互
MENU = """
NM2 试用台
  1) 一键演完整段故事（写入 → 提问 → 遗忘 → 追问 → 对照）
  2) 写入一条事实
  3) 提一个问题（带记忆）
  4) 关掉记忆再问一次（当场对照）
  5) 遗忘（直接对模型说「请遗忘…」，走真实自动遗忘）
  6) 看记忆库
  7) 清空记忆
  0) 退出
"""


def interactive(client: Client) -> None:
    while True:
        print(MENU)
        try:
            choice = input("请选择：").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            return
        try:
            if choice == "1":
                demo(client)
            elif choice == "2":
                entity = input("实体（直接回车=user）：").strip() or "user"
                attribute = input("属性（例如 默认语言）：").strip()
                value = input("取值（例如 葡萄牙语）：").strip()
                text = input("记录原文（直接回车自动拼）：").strip() or f"{attribute} 是 {value}。"
                result = client.write(text=text, entity=entity, attribute=attribute, value=value)
                print(f"结果：{result.get('action')}")
                show_bank(client)
            elif choice == "3":
                question = input("问题：").strip()
                if not question:
                    continue
                reply = client.ask(question)
                show_answer(reply)
                show_decision(reply)
            elif choice == "4":
                question = input("问题（这一轮会把记忆关掉）：").strip()
                if not question:
                    continue
                show_answer(client.ask(question, remember=False))
            elif choice == "5":
                text = input("对模型说：").strip() or "请遗忘我刚刚告诉你的那条信息。"
                reply = client.ask(text)
                show_answer(reply)
                show_decision(reply)
                show_bank(client, "遗忘之后")
            elif choice == "6":
                show_bank(client)
            elif choice == "7":
                client.reset()
                print("已清空。")
            elif choice == "0":
                return
            else:
                print("没有这个选项。")
        except ServiceError as error:
            print(f"出错：{error}")


def main() -> int:
    parser = argparse.ArgumentParser(description="NM2 试用台")
    parser.add_argument("--base", default=DEFAULT_BASE, help="控制面地址")
    parser.add_argument("--demo", action="store_true", help="直接跑自动演示，不进菜单")
    args = parser.parse_args()

    client = Client(args.base)
    try:
        health = client.health()
    except ServiceError as error:
        print(f"启动失败：{error}")
        print("\n先起模型服务，再跑这个脚本。命令见 agent_lab/RESULTS.md 第 8 节。")
        return 1

    memory = health.get("memory", {})
    print(f"已连上模型服务：{health.get('model_path')}")
    print(f"当前记忆开关：{health.get('memory_mode')}｜"
          f"生效记录 {memory.get('active_records')} 条 / 共 {memory.get('records')} 条")

    if args.demo:
        started = time.time()
        demo(client)
        print(f"\n演示结束，用了 {time.time() - started:.0f} 秒。")
    else:
        interactive(client)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())