#!/usr/bin/env python3
# ==============================================================================
# 用途：角色記憶（.memory/<角色>/）的儲存引擎。負責 (1) 把每輪對話濃縮結果寫入
#       inbox，(2) 產生 SessionStart 要注入的記憶區塊，(3) 睡眠整理時輸出待整理
#       素材並套用整理結果（分類／去重／標籤／總結／壓縮歸檔），(4) 依使用頻率
#       遺忘日常與其他類記憶。
# 更新時間：2026/07/28 00:00:00
# 相依：Python 3 標準庫。
# 退出碼：0 成功；1 無內容可處理；2 參數錯誤。呼叫端（hook）一律不得因此中斷。
# ==============================================================================

import argparse
import gzip
import hashlib
import json
import os
import re
import shutil
import sys
from datetime import datetime, timedelta, timezone

# ------------------------------------------------------------------------------
# 常數：分類、載入策略、遺忘規則
# ------------------------------------------------------------------------------

# 六個分類的目錄名（英文，跨平台安全）與中文標籤
CATEGORIES = ["important", "interest", "news", "skill", "daily", "other"]
CATEGORY_LABELS = {
    "important": "重要",
    "interest": "興趣",
    "news": "新知",
    "skill": "技能",
    "daily": "日常",
    "other": "其他",
}
# 中文分類名反查（模型可能直接輸出中文）
LABEL_TO_CATEGORY = {label: key for key, label in CATEGORY_LABELS.items()}

# 載入策略：重要與興趣載入全文，其餘僅載入總結與標籤
FULL_CATEGORIES = ["important", "interest"]
# 摘要載入順序：技能 → 新知 → 日常 → 其他
DIGEST_CATEGORIES = ["skill", "news", "daily", "other"]

# 遺忘規則：(未更新天數門檻, 命中次數上限)；只套用於日常與其他
FORGET_RULES = {"daily": (14, 1), "other": (7, 1)}

# 單次睡眠整理最多處理的 inbox 筆數，其餘留待下個睡眠週期
SLEEP_BATCH = 60
# 送進模型的素材字元上限
COLLECT_LIMIT = 40000
# 單則記憶壓縮後的內容字元上限
CONTENT_LIMIT = 1200

TAIPEI = timezone(timedelta(hours=8))
STAMP_FORMAT = "%Y/%m/%d %H:%M:%S"


# ------------------------------------------------------------------------------
# 路徑與時間
# ------------------------------------------------------------------------------


def now_stamp():
    """回傳台灣時區的 yyyy/MM/dd HH:mm:ss 時間字串。"""
    return datetime.now(TAIPEI).strftime(STAMP_FORMAT)


def parse_stamp(value):
    """把 yyyy/MM/dd HH:mm:ss 字串解析成帶時區的 datetime，失敗回 None。"""
    if not isinstance(value, str) or not value.strip():
        return None
    try:
        return datetime.strptime(value.strip(), STAMP_FORMAT).replace(tzinfo=TAIPEI)
    except ValueError:
        return None


def memory_root(role):
    """回傳指定角色的記憶根目錄（可用 ROLE_MEMORY_HOME 覆寫預設 ~/.memory）。"""
    base = os.environ.get("ROLE_MEMORY_HOME") or os.path.join(os.path.expanduser("~"), ".memory")
    return os.path.join(base, role)


def ensure_layout(role):
    """建立角色記憶目錄結構（inbox、六個分類、archive），回傳根目錄。"""
    root = memory_root(role)
    for sub in ["inbox", "archive/raw", "archive/forgotten"] + CATEGORIES:
        os.makedirs(os.path.join(root, sub), exist_ok=True)
    return root


def state_path(role):
    """回傳角色記憶狀態檔（state.json）的路徑。"""
    return os.path.join(memory_root(role), "state.json")


def read_state(role):
    """讀取狀態檔；不存在或損壞時回空 dict。"""
    try:
        with open(state_path(role), encoding="utf-8") as fh:
            data = json.load(fh)
        return data if isinstance(data, dict) else {}
    except (OSError, ValueError):
        return {}


def write_state(role, patch):
    """把 patch 併入狀態檔後寫回（整份覆寫，內容極小）。"""
    state = read_state(role)
    state.update(patch)
    ensure_layout(role)
    with open(state_path(role), "w", encoding="utf-8") as fh:
        json.dump(state, fh, ensure_ascii=False, indent=2)
    return state


# ------------------------------------------------------------------------------
# 記憶檔格式：YAML 風格 frontmatter ＋ 內文
# ------------------------------------------------------------------------------


def normalize_category(value):
    """把模型輸出的分類（英文或中文）正規化為分類鍵；無法判定時回 other。"""
    raw = (value or "").strip().lower()
    if raw in CATEGORIES:
        return raw
    return LABEL_TO_CATEGORY.get((value or "").strip(), "other")


def normalize_tags(value):
    """把標籤（list 或逗號分隔字串）正規化為去重後的小寫標籤 list，最多 6 個。"""
    if isinstance(value, str):
        parts = re.split(r"[,、|]", value)
    elif isinstance(value, list):
        parts = [str(item) for item in value]
    else:
        parts = []
    tags = []
    for part in parts:
        tag = part.strip().strip("[]#").strip()
        if tag and tag.lower() not in [t.lower() for t in tags]:
            tags.append(tag)
    return tags[:6]


def one_line(value, limit=120):
    """把文字壓成單行並截斷，用於 summary 欄位。"""
    text = re.sub(r"\s+", " ", str(value or "")).strip()
    return text[:limit]


def dump_memory(meta, content):
    """把 meta 與內文組成記憶檔全文（frontmatter ＋ 內文）。"""
    lines = ["---"]
    for key in ["id", "category", "summary", "tags", "created", "updated", "hits", "sources"]:
        if key not in meta:
            continue
        value = meta[key]
        if isinstance(value, list):
            value = "[" + ", ".join(str(item) for item in value) + "]"
        lines.append(f"{key}: {value}")
    lines.append("---")
    lines.append("")
    lines.append(content.strip())
    lines.append("")
    return "\n".join(lines)


def load_memory(path):
    """讀取單一記憶檔，回傳 (meta dict, 內文)；讀取失敗回 (None, "")。"""
    try:
        with open(path, encoding="utf-8") as fh:
            raw = fh.read()
    except OSError:
        return None, ""

    meta = {"path": path, "hits": 0, "tags": []}
    body = raw
    if raw.startswith("---"):
        parts = raw.split("---", 2)
        if len(parts) >= 3:
            body = parts[2]
            for line in parts[1].splitlines():
                if ":" not in line:
                    continue
                key, _, value = line.partition(":")
                key = key.strip()
                value = value.strip()
                if key in ("tags", "sources"):
                    meta[key] = normalize_tags(value.strip("[]"))
                elif key == "hits":
                    meta[key] = int(value) if value.isdigit() else 0
                else:
                    meta[key] = value
    meta.setdefault("id", os.path.splitext(os.path.basename(path))[0])
    meta.setdefault("summary", "")
    meta.setdefault("created", "")
    meta.setdefault("updated", meta.get("created", ""))
    return meta, body.strip()


def new_id(seed):
    """以時間與內容雜湊產生記憶 id，確保同一秒多筆也不碰撞。"""
    digest = hashlib.sha1(seed.encode("utf-8", "replace")).hexdigest()[:6]
    return f"{datetime.now(TAIPEI).strftime('%Y%m%d-%H%M%S')}-{digest}"


def list_memories(role, category):
    """列出某分類下的所有記憶（依 updated 新到舊排序）。"""
    directory = os.path.join(memory_root(role), category)
    items = []
    if not os.path.isdir(directory):
        return items
    for name in sorted(os.listdir(directory)):
        if not name.endswith(".md"):
            continue
        meta, content = load_memory(os.path.join(directory, name))
        if meta is None:
            continue
        meta["category"] = category
        items.append((meta, content))
    items.sort(key=lambda item: item[0].get("updated") or "", reverse=True)
    return items


def list_inbox(role):
    """列出 inbox 內尚未整理的記憶（依檔名，即時間先後排序）。"""
    directory = os.path.join(memory_root(role), "inbox")
    items = []
    if not os.path.isdir(directory):
        return items
    for name in sorted(os.listdir(directory)):
        if not name.endswith(".md"):
            continue
        meta, content = load_memory(os.path.join(directory, name))
        if meta is not None:
            items.append((meta, content))
    return items


def find_memory(role, memory_id):
    """依 id 在六個分類中尋找記憶檔，回傳 (meta, 內文)；找不到回 (None, "")。"""
    for category in CATEGORIES:
        path = os.path.join(memory_root(role), category, f"{memory_id}.md")
        if os.path.isfile(path):
            meta, content = load_memory(path)
            if meta is not None:
                meta["category"] = category
            return meta, content
    return None, ""


def archive_file(path, destination_dir):
    """把檔案 gzip 後搬到歸檔目錄，原檔刪除；失敗時保留原檔。"""
    os.makedirs(destination_dir, exist_ok=True)
    target = os.path.join(destination_dir, os.path.basename(path) + ".gz")
    try:
        with open(path, "rb") as src, gzip.open(target, "wb") as dst:
            shutil.copyfileobj(src, dst)
        os.remove(path)
        return True
    except OSError:
        return False


# ------------------------------------------------------------------------------
# 子命令：write —— 由 Stop hook 寫入一則未整理記憶
# ------------------------------------------------------------------------------

FIELD_PATTERN = re.compile(r"^\s*(CATEGORY|SUMMARY|TAGS|CONTENT)\s*[:：]\s*(.*)$", re.IGNORECASE)


def parse_capture(text):
    """
    解析 Stop hook 濃縮器輸出的四欄格式（CATEGORY／SUMMARY／TAGS／CONTENT）。

    模型可能夾帶前後贅字，故逐行掃描欄位標記，CONTENT 之後的所有內容視為內文。
    """
    category = summary = ""
    tags = []
    content_lines = []
    in_content = False
    for line in text.splitlines():
        match = FIELD_PATTERN.match(line)
        if match and not (in_content and match.group(1).upper() != "CONTENT"):
            field = match.group(1).upper()
            value = match.group(2)
            if field == "CATEGORY":
                category = value
            elif field == "SUMMARY":
                summary = value
            elif field == "TAGS":
                tags = normalize_tags(value)
            elif field == "CONTENT":
                in_content = True
                if value.strip():
                    content_lines.append(value)
            continue
        if in_content:
            content_lines.append(line)
    return category, summary, tags, "\n".join(content_lines).strip()


def cmd_write(args):
    """把 stdin 的濃縮結果寫成一則 inbox 記憶。"""
    raw = sys.stdin.read()
    category, summary, tags, content = parse_capture(raw)
    if not content and not summary:
        return 1
    if not content:
        content = summary
    content = content[:CONTENT_LIMIT]
    stamp = now_stamp()
    meta = {
        "id": new_id(content + stamp),
        "category": normalize_category(category),
        "summary": one_line(summary) or one_line(content),
        "tags": tags,
        "created": stamp,
        "updated": stamp,
        "hits": 1,
    }
    if args.project:
        meta["sources"] = [args.project]
    ensure_layout(args.role)
    path = os.path.join(memory_root(args.role), "inbox", f"{meta['id']}.md")
    with open(path, "w", encoding="utf-8") as fh:
        fh.write(dump_memory(meta, content))
    sys.stdout.write(meta["id"])
    return 0


# ------------------------------------------------------------------------------
# 子命令：load —— 產生 SessionStart 要注入的記憶區塊
# ------------------------------------------------------------------------------


def cmd_load(args):
    """
    組出載入用記憶區塊：重要與興趣載入全文，其餘依技能→新知→日常→其他只載總結與標籤。

    超過字元上限時截斷並標明，避免佔滿 context。
    """
    limit = args.limit
    blocks = []
    total_full = 0
    for category in FULL_CATEGORIES:
        items = list_memories(args.role, category)
        if not items:
            continue
        lines = [f"### {CATEGORY_LABELS[category]}記憶（全文）"]
        for meta, content in items:
            tags = "、".join(meta.get("tags") or []) or "無標籤"
            lines.append(f"- **{meta.get('summary') or '（無總結）'}**（標籤：{tags}）")
            for line in content.splitlines():
                if line.strip():
                    lines.append(f"  {line.strip()}")
            total_full += 1
        blocks.append("\n".join(lines))

    digest_lines = []
    digest_count = 0
    for category in DIGEST_CATEGORIES:
        items = list_memories(args.role, category)
        if not items:
            continue
        digest_lines.append(f"### {CATEGORY_LABELS[category]}記憶（總結）")
        for meta, _ in items:
            tags = "、".join(meta.get("tags") or []) or "無標籤"
            digest_lines.append(f"- {meta.get('summary') or '（無總結）'}（標籤：{tags}）")
            digest_count += 1
    if digest_lines:
        blocks.append("\n".join(digest_lines))

    pending = len(list_inbox(args.role))
    if pending:
        blocks.append(f"> 尚有 {pending} 則未整理記憶，將於下次睡眠時段歸檔。")

    if not blocks:
        return 1

    text = "\n\n".join(blocks)
    if len(text) > limit:
        text = text[:limit] + f"\n\n> （記憶內容超過 {limit} 字元已截斷，完整記憶仍保存在磁碟）"
    sys.stdout.write(text)
    return 0


# ------------------------------------------------------------------------------
# 子命令：collect —— 睡眠整理前輸出待整理素材
# ------------------------------------------------------------------------------


def cmd_collect(args):
    """輸出送進模型的整理素材：inbox 待整理項目 ＋ 既有記憶索引（供去重比對）。"""
    inbox = list_inbox(args.role)[:SLEEP_BATCH]
    if not inbox:
        return 1

    lines = ["=== INBOX（待整理，每則以 id 標識）==="]
    for meta, content in inbox:
        lines.append(f"--- id: {meta['id']} | 時間: {meta.get('created', '-')} ---")
        lines.append(f"初判分類: {CATEGORY_LABELS.get(meta.get('category', 'other'), '其他')}")
        lines.append(f"初判總結: {meta.get('summary', '')}")
        lines.append(f"初判標籤: {'、'.join(meta.get('tags') or []) or '無'}")
        lines.append("內容:")
        lines.append(content)
        lines.append("")

    lines.append("=== EXISTING（既有記憶索引，供去重與合併判斷）===")
    existing = 0
    for category in CATEGORIES:
        for meta, _ in list_memories(args.role, category):
            tags = "、".join(meta.get("tags") or []) or "無"
            lines.append(
                f"- id: {meta['id']} | 分類: {CATEGORY_LABELS[category]} | 標籤: {tags} | 總結: {meta.get('summary', '')}"
            )
            existing += 1
    if not existing:
        lines.append("（尚無既有記憶）")

    text = "\n".join(lines)
    if len(text) > COLLECT_LIMIT:
        text = text[:COLLECT_LIMIT] + "\n…（素材過長已截斷，其餘留待下個睡眠週期）…"
    sys.stdout.write(text)
    return 0


# ------------------------------------------------------------------------------
# 子命令：apply —— 套用睡眠整理結果
# ------------------------------------------------------------------------------


def extract_json(text):
    """從模型輸出中取出第一個 JSON 物件（容忍 code fence 與前後贅字）。"""
    stripped = text.strip()
    fence = re.search(r"```(?:json)?\s*(.*?)```", stripped, re.DOTALL)
    if fence:
        stripped = fence.group(1).strip()
    start = stripped.find("{")
    end = stripped.rfind("}")
    if start < 0 or end <= start:
        return None
    try:
        return json.loads(stripped[start : end + 1])
    except ValueError:
        return None


def cmd_apply(args):
    """
    讀取 stdin 的整理結果 JSON，寫入分類記憶並歸檔對應的 inbox 原始檔。

    action 支援 new（新建）／merge（併入既有記憶）／drop（判定無保存價值）。
    未被提及的 inbox 檔一律保留，留待下個睡眠週期，避免整理失敗造成記憶遺失。
    """
    data = extract_json(sys.stdin.read())
    if not isinstance(data, dict):
        sys.stderr.write("整理結果非合法 JSON\n")
        return 1
    entries = data.get("memories")
    if not isinstance(entries, list) or not entries:
        sys.stderr.write("整理結果不含 memories\n")
        return 1

    ensure_layout(args.role)
    root = memory_root(args.role)
    stamp = now_stamp()
    counts = {"new": 0, "merge": 0, "drop": 0}
    consumed = []

    for entry in entries:
        if not isinstance(entry, dict):
            continue
        action = str(entry.get("action") or "new").strip().lower()
        sources = [str(item).strip() for item in (entry.get("from") or []) if str(item).strip()]

        if action == "drop":
            consumed.extend(sources)
            counts["drop"] += 1
            continue

        content = str(entry.get("content") or "").strip()[:CONTENT_LIMIT]
        summary = one_line(entry.get("summary"))
        tags = normalize_tags(entry.get("tags"))
        if not content and not summary:
            continue

        if action == "merge":
            target_id = str(entry.get("target") or "").strip()
            meta, old_content = find_memory(args.role, target_id)
            if meta is None:
                action = "new"
            else:
                category = normalize_category(entry.get("category") or meta.get("category"))
                merged_tags = normalize_tags((meta.get("tags") or []) + tags)
                new_meta = {
                    "id": meta["id"],
                    "category": category,
                    "summary": summary or meta.get("summary", ""),
                    "tags": merged_tags,
                    "created": meta.get("created") or stamp,
                    "updated": stamp,
                    "hits": int(meta.get("hits") or 0) + 1,
                }
                old_path = meta["path"]
                new_path = os.path.join(root, category, f"{meta['id']}.md")
                with open(new_path, "w", encoding="utf-8") as fh:
                    fh.write(dump_memory(new_meta, content or old_content))
                if os.path.abspath(old_path) != os.path.abspath(new_path):
                    try:
                        os.remove(old_path)
                    except OSError:
                        pass
                consumed.extend(sources)
                counts["merge"] += 1
                continue

        category = normalize_category(entry.get("category"))
        meta = {
            "id": new_id(content + summary + stamp),
            "category": category,
            "summary": summary or one_line(content),
            "tags": tags,
            "created": stamp,
            "updated": stamp,
            "hits": 1,
        }
        with open(os.path.join(root, category, f"{meta['id']}.md"), "w", encoding="utf-8") as fh:
            fh.write(dump_memory(meta, content or summary))
        consumed.extend(sources)
        counts["new"] += 1

    archived = 0
    month_dir = os.path.join(root, "archive", "raw", datetime.now(TAIPEI).strftime("%Y-%m"))
    for source_id in set(consumed):
        path = os.path.join(root, "inbox", f"{source_id}.md")
        if os.path.isfile(path) and archive_file(path, month_dir):
            archived += 1

    write_state(args.role, {"last_sleep": stamp, "last_sleep_epoch": int(datetime.now(TAIPEI).timestamp())})
    sys.stdout.write(
        f"新增 {counts['new']} 則、合併 {counts['merge']} 則、捨棄 {counts['drop']} 則、歸檔原始記憶 {archived} 則"
    )
    return 0


# ------------------------------------------------------------------------------
# 子命令：forget —— 依使用頻率遺忘日常與其他
# ------------------------------------------------------------------------------


def cmd_forget(args):
    """把日常／其他分類中久未更新且命中次數低的記憶壓縮到 archive/forgotten 後移除。"""
    root = ensure_layout(args.role)
    now = datetime.now(TAIPEI)
    forgotten = []
    for category, (days, max_hits) in FORGET_RULES.items():
        for meta, _ in list_memories(args.role, category):
            updated = parse_stamp(meta.get("updated")) or parse_stamp(meta.get("created"))
            if updated is None:
                continue
            if (now - updated).days < days:
                continue
            if int(meta.get("hits") or 0) > max_hits:
                continue
            if args.dry_run:
                forgotten.append(f"{CATEGORY_LABELS[category]}｜{meta.get('summary', '')}")
                continue
            if archive_file(meta["path"], os.path.join(root, "archive", "forgotten")):
                forgotten.append(f"{CATEGORY_LABELS[category]}｜{meta.get('summary', '')}")

    if not forgotten:
        sys.stdout.write("沒有符合遺忘條件的記憶")
        return 0
    prefix = "（預覽）" if args.dry_run else ""
    sys.stdout.write(f"{prefix}遺忘 {len(forgotten)} 則：\n" + "\n".join(f"- {item}" for item in forgotten))
    if not args.dry_run:
        write_state(args.role, {"last_forget": now_stamp()})
    return 0


# ------------------------------------------------------------------------------
# 子命令：stats —— 供 skill 與診斷顯示記憶概況
# ------------------------------------------------------------------------------


def cmd_stats(args):
    """輸出記憶統計（各分類筆數、待整理筆數、上次整理時間）。"""
    state = read_state(args.role)
    rows = [f"| 分類 | 筆數 |", "| --- | --- |"]
    for category in CATEGORIES:
        rows.append(f"| {CATEGORY_LABELS[category]} | {len(list_memories(args.role, category))} |")
    rows.append(f"| 待整理（inbox） | {len(list_inbox(args.role))} |")
    rows.append("")
    rows.append(f"- 記憶目錄：{memory_root(args.role)}")
    rows.append(f"- 上次睡眠整理：{state.get('last_sleep', '尚未整理')}")
    rows.append(f"- 上次遺忘：{state.get('last_forget', '尚未執行')}")
    sys.stdout.write("\n".join(rows))
    return 0


# ------------------------------------------------------------------------------
# 子命令：need-sleep —— 判斷是否需要補跑睡眠整理
# ------------------------------------------------------------------------------


def cmd_mark_sleep(args):
    """把本次睡眠週期標記為已整理（inbox 為空、無素材可整理時使用）。"""
    write_state(args.role, {"last_sleep": now_stamp(), "last_sleep_epoch": int(datetime.now(TAIPEI).timestamp())})
    sys.stdout.write("已更新上次整理時間")
    return 0


def cmd_need_sleep(args):
    """
    判斷是否需要補跑整理：距上次整理超過門檻小時數且 inbox 有內容。

    輸出 yes／no，供 shell 直接判斷（不用解析 JSON）。
    """
    if not list_inbox(args.role):
        sys.stdout.write("no")
        return 0
    state = read_state(args.role)
    last = parse_stamp(state.get("last_sleep"))
    if last is None:
        sys.stdout.write("yes")
        return 0
    hours = (datetime.now(TAIPEI) - last).total_seconds() / 3600
    sys.stdout.write("yes" if hours >= args.hours else "no")
    return 0


# ------------------------------------------------------------------------------
# CLI
# ------------------------------------------------------------------------------


def build_parser():
    """建立子命令解析器。"""
    parser = argparse.ArgumentParser(description="角色記憶儲存引擎")
    sub = parser.add_subparsers(dest="command", required=True)

    write = sub.add_parser("write", help="自 stdin 讀濃縮結果寫入 inbox")
    write.add_argument("--role", required=True)
    write.add_argument("--project", default="")
    write.set_defaults(func=cmd_write)

    load = sub.add_parser("load", help="輸出 SessionStart 要注入的記憶區塊")
    load.add_argument("--role", required=True)
    load.add_argument("--limit", type=int, default=int(os.environ.get("ROLE_LOAD_LIMIT", "8000")))
    load.set_defaults(func=cmd_load)

    collect = sub.add_parser("collect", help="輸出睡眠整理素材")
    collect.add_argument("--role", required=True)
    collect.set_defaults(func=cmd_collect)

    apply_cmd = sub.add_parser("apply", help="自 stdin 讀整理結果 JSON 並套用")
    apply_cmd.add_argument("--role", required=True)
    apply_cmd.set_defaults(func=cmd_apply)

    forget = sub.add_parser("forget", help="依使用頻率遺忘日常與其他記憶")
    forget.add_argument("--role", required=True)
    forget.add_argument("--dry-run", action="store_true")
    forget.set_defaults(func=cmd_forget)

    stats = sub.add_parser("stats", help="輸出記憶統計")
    stats.add_argument("--role", required=True)
    stats.set_defaults(func=cmd_stats)

    mark = sub.add_parser("mark-sleep", help="標記本次睡眠週期已整理")
    mark.add_argument("--role", required=True)
    mark.set_defaults(func=cmd_mark_sleep)

    need = sub.add_parser("need-sleep", help="判斷是否需要補跑睡眠整理")
    need.add_argument("--role", required=True)
    need.add_argument("--hours", type=float, default=20.0)
    need.set_defaults(func=cmd_need_sleep)

    return parser


def main(argv):
    """CLI 進入點。"""
    args = build_parser().parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))