# memory_tier_manager.py - 记忆分层管理器
# 自动将记忆分为 core/hot/cold 三层

import sqlite3
from datetime import datetime, timedelta
from vector_memory import VectorMemorySystem


class MemoryTierManager:
    """记忆分层管理器"""
    
    def __init__(self, vector_memory: VectorMemorySystem):
        self.vm = vector_memory
        self.conn = vector_memory.conn
    
    def add_with_tier(self, content: str, importance: int = 3, 
                      tags: list = None, auto_archive: bool = True):
        """自动分层添加记忆"""
        metadata = {
            'tags': tags or [],
            'importance': importance,
            'auto_archive': auto_archive
        }
        
        memory_id = self.vm.add_memory(
            content=content,
            metadata=metadata,
            importance=importance
        )
        
        # 根据重要性自动分层
        if importance >= 4:
            tier = "core"
        elif importance >= 2:
            tier = "hot"
        else:
            tier = "cold"
        
        # 标记层级
        self.conn.execute(
            "UPDATE memories SET tier=? WHERE id=?",
            (tier, memory_id)
        )
        self.conn.commit()
        
        return memory_id
    
    def get_recent_memories(self, days: int = 7, limit: int = 20):
        """获取最近记忆"""
        cursor = self.conn.execute("""
            SELECT id, content, metadata, importance, created_at 
            FROM memories 
            ORDER BY created_at DESC 
            LIMIT ?
        """, (limit,))
        
        return [{
            'id': row[0],
            'content': row[1],
            'metadata': row[2],
            'importance': row[3],
            'created_at': row[4]
        } for row in cursor.fetchall()]
    
    def get_core_memories(self):
        """获取核心记忆（重要性 >= 4）"""
        cursor = self.conn.execute("""
            SELECT id, content, metadata, importance, created_at 
            FROM memories 
            WHERE importance >= 4
            ORDER BY created_at DESC
        """)
        
        return [{
            'id': row[0],
            'content': row[1],
            'metadata': row[2],
            'importance': row[3],
            'created_at': row[4]
        } for row in cursor.fetchall()]
    
    def migrate_old_memories(self, hot_days: int = 30):
        """迁移旧记忆到冷存储"""
        cutoff = datetime.now() - timedelta(days=hot_days)
        
        # 找出需要归档的记忆
        cursor = self.conn.execute("""
            SELECT id, content, metadata 
            FROM memories 
            WHERE importance < 3
            AND created_at < ?
        """, (cutoff,))
        
        archived = 0
        for row in cursor.fetchall():
            # 可以在这里实现归档逻辑（如写入文件、压缩等）
            archived += 1
        
        return archived
