#!/usr/bin/env python3
"""
Anchor Decay Engine
移植 Ombre-Brain 遗忘曲线算法，模拟人类自然遗忘。
定期扫描记忆文件，计算活跃度得分，自动归档低活跃记忆。
"""
import os
import math
import yaml
import re
import random
from datetime import datetime, timedelta

VAULT_DIR = "/Users/fyah/Documents/如梦初醒/memory/test/Anchor-Memory"
ARCHIVE_DIR = os.path.join(VAULT_DIR, "Archive")

# 遗忘曲线参数（来自 Ombre-Brain）
DECAY_LAMBDA = 0.05       # 衰减速率
THRESHOLD = 0.3           # 归档阈值
EMOTION_BASE = 1.0        # 情感基础权重
AROUSAL_BOOST = 0.8       # 唤醒度加成
RESURRECTION_CHANCE = 0.1 # 返场概率 (10%)：被遗忘的记忆有几率“诈尸”复活

def parse_yaml_frontmatter(content):
    """极简 YAML 解析"""
    if content.startswith("---"):
        parts = content.split("---", 2)
        if len(parts) >= 2:
            try:
                return yaml.safe_load(parts[1]), parts[2]
            except:
                pass
    return {}, content

def update_yaml_frontmatter(content, metadata):
    """更新 YAML 头"""
    if content.startswith("---"):
        parts = content.split("---", 2)
        if len(parts) >= 2:
            return f"---\n{yaml.dump(metadata, allow_unicode=True)}---{parts[2]}"
    return content

def calc_time_weight(days_since: float) -> float:
    """新鲜度加成：1.0 + e^(-t/36), t 为小时"""
    hours = days_since * 24.0
    return 1.0 + 1.0 * math.exp(-hours / 36.0)

def calculate_score(metadata: dict) -> float:
    """计算记忆活跃度得分"""
    if not isinstance(metadata, dict):
        return 0.0

    # 钉选/永久记忆不衰减
    if metadata.get("pinned") or metadata.get("type") == "permanent":
        return 999.0

    importance = max(1, min(10, int(metadata.get("importance", 5))))
    activation_count = max(1.0, float(metadata.get("activation_count", 1)))

    # 计算天数
    last_active_str = metadata.get("last_active", metadata.get("date", ""))
    try:
        # 兼容多种日期格式
        last_active = datetime.strptime(str(last_active_str), "%Y-%m-%d %H:%M")
        days_since = max(0.0, (datetime.now() - last_active).total_seconds() / 86400)
    except:
        days_since = 30.0

    # 情感权重
    try:
        arousal = max(0.0, min(1.0, float(metadata.get("arousal", 0.3))))
    except:
        arousal = 0.3
    emotion_weight = EMOTION_BASE + arousal * AROUSAL_BOOST

    # 时间权重
    time_weight = calc_time_weight(days_since)

    # 短期/长期分离
    if days_since <= 3.0:
        combined_weight = time_weight * 0.7 + emotion_weight * 0.3
    else:
        combined_weight = emotion_weight * 0.7 + time_weight * 0.3

    # 核心公式
    base_score = (
        importance
        * (activation_count ** 0.3)
        * math.exp(-DECAY_LAMBDA * days_since)
        * combined_weight
    )

    return base_score

def run_decay():
    """执行遗忘扫描"""
    print(f"[*] 开始扫描遗忘曲线: {VAULT_DIR}")
    archived_count = 0

    # 确保归档目录存在
    os.makedirs(ARCHIVE_DIR, exist_ok=True)

    # 遍历所有 .md 文件
    for root, dirs, files in os.walk(VAULT_DIR):
        # 跳过归档目录本身
        if "Archive" in root:
            continue
            
        for file in files:
            if not file.endswith(".md"):
                continue
            
            filepath = os.path.join(root, file)
            try:
                with open(filepath, "r", encoding="utf-8") as f:
                    content = f.read()
                
                metadata, body = parse_yaml_frontmatter(content)
                if not metadata:
                    continue
                
                score = calculate_score(metadata)
                
                # 如果得分低于阈值，且当前状态不是已归档，则归档
                if score < THRESHOLD and metadata.get("status") != "archived":
                    # 随机返场判定
                    if random.random() < RESURRECTION_CHANCE:
                        print(f"  [✨ 返场] {file} (得分: {score:.3f}) - 触发随机复活！")
                        metadata["status"] = "resurrected"
                        metadata["last_resurrected"] = datetime.now().strftime("%Y-%m-%d %H:%M")
                        new_content = update_yaml_frontmatter(content, metadata)
                        with open(filepath, "w", encoding="utf-8") as f:
                            f.write(new_content)
                    else:
                       # 正常归档逻辑
                       metadata["status"] = "archived"
                       metadata["archived_date"] = datetime.now().strftime("%Y-%m-%d %H:%M")
                       metadata["decay_score"] = round(score, 3)
                        metadata["original_path"] = os.path.relpath(filepath, VAULT_DIR)
                       
                       new_content = update_yaml_frontmatter(content, metadata)
                        with open(filepath, "w", encoding="utf-8") as f:
                            f.write(new_content)
                        
                        # 移动文件到归档目录
                        dest_path = os.path.join(ARCHIVE_DIR, file)
                        os.rename(filepath, dest_path)
                        archived_count += 1
                        print(f"  [归档] {file} (得分: {score:.3f})")
                    
            except Exception as e:
                print(f"  [!] 处理失败 {file}: {e}")

    print(f"[+] 遗忘扫描完成。共归档 {archived_count} 个文件。")
   print(f"[+] 遗忘扫描完成。共归档 {archived_count} 个文件。")
    # --- 第二阶段：扫描归档目录，寻找返场机会 ---
    print(f"[*] 扫描归档目录寻找返场机会: {ARCHIVE_DIR}")
    resurrected_count = 0
    
    if os.path.exists(ARCHIVE_DIR):
        for file in os.listdir(ARCHIVE_DIR):
            if not file.endswith(".md"):
                continue
            
            filepath = os.path.join(ARCHIVE_DIR, file)
            try:
                with open(filepath, "r", encoding="utf-8") as f:
                    content = f.read()
                
                metadata, body = parse_yaml_frontmatter(content)
                
                # 只针对已归档的文件
                if metadata.get("status") == "archived":
                    # 掷骰子：10% 概率复活
                    if random.random() < RESURRECTION_CHANCE:
                        print(f"  [✨ 返场] {file} - 从归档中复活！")
                        metadata["status"] = "resurrected"
                        metadata["last_resurrected"] = datetime.now().strftime("%Y-%m-%d %H:%M")
                        
                        new_content = update_yaml_frontmatter(content, metadata)
                        with open(filepath, "w", encoding="utf-8") as f:
                            f.write(new_content)
                        
                        # 恢复原路径
                        orig_path = metadata.get("original_path", file)
                        dest_path = os.path.join(VAULT_DIR, orig_path)
                        os.makedirs(os.path.dirname(dest_path), exist_ok=True)
                        os.rename(filepath, dest_path)
                        resurrected_count += 1
                        
            except Exception as e:
                print(f"  [!] 处理归档文件 {file} 失败: {e}")
                
    print(f"[+] 返场扫描完成。共复活 {resurrected_count} 个文件。")
if __name__ == "__main__":
   run_decay()