#!/usr/bin/env python3
"""
🦐 皮皮虾记忆管理工具
自动化记忆维护：整理、清理、提炼、归档
"""

import os
import json
import re
from datetime import datetime, timedelta
from pathlib import Path

# 配置
WORKSPACE = Path.home() / ".openclaw" / "workspace"
MEMORY_DIR = WORKSPACE / "memory"
MEMORY_MD = WORKSPACE / "MEMORY.md"
RETENTION_DAYS = 30  # 短期记忆保留天数

# ANSI 颜色
class Colors:
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    BLUE = '\033[94m'
    END = '\033[0m'
    BOLD = '\033[1m'

def print_header(text):
    print(f"\n{Colors.BOLD}{Colors.BLUE}{'='*60}{Colors.END}")
    print(f"{Colors.BOLD}{Colors.BLUE}  {text}{Colors.END}")
    print(f"{Colors.BOLD}{Colors.BLUE}{'='*60}{Colors.END}\n")

def print_success(text):
    print(f"{Colors.GREEN}✅ {text}{Colors.END}")

def print_warning(text):
    print(f"{Colors.YELLOW}⚠️  {text}{Colors.END}")

def print_info(text):
    print(f"{Colors.BLUE}📌 {text}{Colors.END}")

def get_memory_files():
    """获取所有记忆文件"""
    if not MEMORY_DIR.exists():
        return []
    return sorted([f for f in MEMORY_DIR.glob("*.md") if f.is_file()])

def get_daily_logs():
    """获取所有每日日志文件"""
    files = get_memory_files()
    daily_logs = []
    for f in files:
        if re.match(r'^\d{4}-\d{2}-\d{2}\.md$', f.name):
            daily_logs.append(f)
    return daily_logs

def parse_log_date(filename):
    """从文件名解析日期"""
    match = re.match(r'^(\d{4}-\d{2}-\d{2})\.md$', filename)
    if match:
        return datetime.strptime(match.group(1), '%Y-%m-%d')
    return None

def get_file_size_kb(filepath):
    """获取文件大小（KB）"""
    return round(filepath.stat().st_size / 1024, 2)

def read_file_content(filepath):
    """读取文件内容"""
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            return f.read()
    except Exception as e:
        print_warning(f"读取失败 {filepath.name}: {e}")
        return None

def write_file_content(filepath, content):
    """写入文件内容"""
    try:
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(content)
        return True
    except Exception as e:
        print_warning(f"写入失败 {filepath.name}: {e}")
        return False

def extract_important_content(content, filepath):
    """从每日日志中提取重要内容"""
    important_items = []
    filename = filepath.stem  # 不含扩展名的文件名
    
    # 提取会话记录中的关键事件
    session_matches = re.findall(r'###.*?—.*?\n(.*?)(?=###|$)', content, re.DOTALL)
    for session in session_matches[:5]:  # 最多 5 个会话
        # 提取用户请求
        user_req = re.search(r'用户请求：(.*?)(?:\n|$)', session)
        # 提取结果
        result = re.search(r'结果：(✅.*?)(?:\n|$)', session)
        
        if user_req and result:
            important_items.append(f"**{filename}** - {user_req.group(1).strip()} → {result.group(1).strip()}")
        elif user_req:
            important_items.append(f"**{filename}** - {user_req.group(1).strip()}")
    
    # 提取重要性评分记录中的高分项
    rating_matches = re.findall(r'\|\s*([^\|]+)\s*\|\s*([45])\s*\|\s*(✅.*?)(?:\n|$)', content)
    for item, score, action in rating_matches[:3]:
        important_items.append(f"**高分记忆** - {item.strip()} (评分{score}) → {action.strip()}")
    
    return important_items

def update_memory_md(important_items):
    """更新 MEMORY.md 文件"""
    if not MEMORY_MD.exists():
        print_warning("MEMORY.md 不存在，跳过更新")
        return
    
    content = read_file_content(MEMORY_MD)
    if not content:
        return
    
    # 在"核心记忆内容"部分添加新内容
    today = datetime.now().strftime('%Y-%m-%d')
    new_section = f"\n### {today} 自动提炼\n\n"
    for item in important_items[:10]:  # 最多添加 10 条
        new_section += f"{item}\n"
    new_section += "\n"
    
    # 找到"核心记忆内容"部分并插入
    if "## 📝 核心记忆内容" in content:
        parts = content.split("## 📝 核心记忆内容")
        if len(parts) == 2:
            new_content = parts[0] + "## 📝 核心记忆内容" + new_section + parts[1]
            # 更新最后更新时间
            new_content = re.sub(
                r'\*最后更新：\d{4}-\d{2}-\d{2}',
                f'*最后更新：{today}',
                new_content
            )
            if write_file_content(MEMORY_MD, new_content):
                print_success("MEMORY.md 已更新")
    else:
        print_warning("未找到核心记忆内容部分")

def archive_old_logs(old_logs):
    """归档旧日志"""
    if not old_logs:
        return
    
    archive_dir = MEMORY_DIR / "archive"
    archive_dir.mkdir(exist_ok=True)
    
    for log_file in old_logs:
        # 移动到 archive 目录
        dest = archive_dir / log_file.name
        try:
            log_file.rename(dest)
            print_success(f"已归档：{log_file.name} → archive/")
        except Exception as e:
            print_warning(f"归档失败 {log_file.name}: {e}")

def cleanup_empty_files():
    """清理空文件"""
    cleaned = 0
    for f in get_memory_files():
        if f.stat().st_size == 0:
            try:
                f.unlink()
                print_success(f"已删除空文件：{f.name}")
                cleaned += 1
            except Exception as e:
                print_warning(f"删除失败 {f.name}: {e}")
    return cleaned

def generate_report(daily_logs, old_logs, important_items):
    """生成维护报告"""
    today = datetime.now()
    
    print_header("📊 记忆维护报告")
    
    print(f"{Colors.BOLD}记忆文件统计：{Colors.END}")
    print(f"  - 每日日志总数：{len(daily_logs)}")
    print(f"  - 长期记忆文件：{len(get_memory_files()) - len(daily_logs)}")
    
    print(f"\n{Colors.BOLD}清理结果：{Colors.END}")
    if old_logs:
        print(f"  - 超过{RETENTION_DAYS}天的日志：{len(old_logs)} 个")
        print(f"  - 已归档到 memory/archive/")
    else:
        print(f"  - 无需清理（所有日志都在{RETENTION_DAYS}天内）")
    
    print(f"\n{Colors.BOLD}内容提炼：{Colors.END}")
    if important_items:
        print(f"  - 提炼重要内容：{len(important_items)} 条")
        print(f"  - 已添加到 MEMORY.md")
    else:
        print(f"  - 暂无可提炼内容")
    
    print(f"\n{Colors.BOLD}下次维护：{Colors.END}")
    next_check = today + timedelta(days=7)
    print(f"  - 建议下次检查：{next_check.strftime('%Y-%m-%d')}")

def main():
    """主函数"""
    print_header("🦐 皮皮虾记忆维护工具")
    
    # 1. 获取所有每日日志
    daily_logs = get_daily_logs()
    print_info(f"找到 {len(daily_logs)} 个每日日志文件")
    
    # 2. 识别过期日志（超过 RETENTION_DAYS 天）
    today = datetime.now()
    cutoff_date = today - timedelta(days=RETENTION_DAYS)
    old_logs = []
    
    for log_file in daily_logs:
        log_date = parse_log_date(log_file.name)
        if log_date and log_date < cutoff_date:
            old_logs.append(log_file)
    
    if old_logs:
        print_warning(f"发现 {len(old_logs)} 个过期日志（超过{RETENTION_DAYS}天）")
    else:
        print_success(f"所有日志都在{RETENTION_DAYS}天内，无需清理")
    
    # 3. 从最近的日志中提取重要内容
    important_items = []
    recent_logs = sorted(daily_logs, key=lambda x: x.name, reverse=True)[:3]
    
    for log_file in recent_logs:
        content = read_file_content(log_file)
        if content:
            items = extract_important_content(content, log_file)
            important_items.extend(items)
    
    if important_items:
        print_success(f"从最近日志中提取了 {len(important_items)} 条重要内容")
    
    # 4. 更新 MEMORY.md
    if important_items:
        update_memory_md(important_items)
    
    # 5. 归档过期日志（询问用户确认）
    if old_logs:
        print(f"\n{Colors.YELLOW}准备归档 {len(old_logs)} 个过期日志到 memory/archive/{Colors.END}")
        # 自动执行归档
        archive_old_logs(old_logs)
    
    # 6. 清理空文件
    cleaned = cleanup_empty_files()
    if cleaned:
        print_success(f"清理了 {cleaned} 个空文件")
    
    # 7. 生成报告
    generate_report(daily_logs, old_logs, important_items)
    
    print(f"\n{Colors.GREEN}🦐 记忆维护完成！{Colors.END}\n")

if __name__ == "__main__":
    main()
