#!/usr/bin/env python3
"""
月度记忆清理脚本
- 统计各 agent_id 的记忆总量
- 列出各 memory_type / visibility 分布
- 可选: 清理过期或长期未命中的记忆

使用: python3 memory_cleanup.py [--dry-run] [--max-age-days 90]
"""

import os
import sys
import argparse
import logging
from datetime import datetime

_dashscope_key = os.getenv('MEM0_DASHSCOPE_API_KEY', '') or os.getenv('DASHSCOPE_API_KEY', '')
if _dashscope_key:
    os.environ['OPENAI_API_KEY'] = _dashscope_key
os.environ.setdefault('OPENAI_API_BASE', 'https://dashscope.aliyuncs.com/compatible-mode/v1')
os.environ.setdefault('OPENAI_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1')

try:
    from qdrant_client import QdrantClient
    from qdrant_client.models import Filter, FieldCondition, MatchValue
except ImportError:
    print("qdrant-client not installed")
    sys.exit(1)

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
logger = logging.getLogger(__name__)

QDRANT_HOST = os.getenv('MEM0_QDRANT_HOST', 'localhost')
QDRANT_PORT = int(os.getenv('MEM0_QDRANT_PORT', '6333'))
COLLECTION = 'mem0_v4_shared'


def get_stats(client: QdrantClient):
    """统计各维度记忆数量"""
    total = client.count(collection_name=COLLECTION)
    logger.info(f"Collection '{COLLECTION}' total: {total.count}")

    for field, values in [
        ('agent_id', ['main', 'life', 'advert_pm', 'general']),
        ('visibility', ['public', 'project', 'private']),
        ('memory_type', ['session', 'chat_summary', 'preference', 'knowledge']),
    ]:
        logger.info(f"\n--- {field} ---")
        for val in values:
            try:
                result = client.count(
                    collection_name=COLLECTION,
                    count_filter=Filter(must=[
                        FieldCondition(key=field, match=MatchValue(value=val))
                    ])
                )
                if result.count > 0:
                    logger.info(f"  {field}={val}: {result.count}")
            except Exception:
                pass


def main():
    parser = argparse.ArgumentParser(description='Mem0 月度记忆清理')
    parser.add_argument('--dry-run', action='store_true', help='仅统计，不删除')
    parser.add_argument('--max-age-days', type=int, default=90,
                        help='清理超过 N 天的 session 类型记忆 (默认 90)')
    args = parser.parse_args()

    client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)

    logger.info("=" * 50)
    logger.info(f"Mem0 记忆清理 - {datetime.now().strftime('%Y-%m-%d %H:%M')}")
    logger.info(f"模式: {'dry-run (仅统计)' if args.dry_run else 'LIVE'}")
    logger.info("=" * 50)

    get_stats(client)

    if args.dry_run:
        logger.info("\n[dry-run] 不执行清理操作")
        return

    logger.info(f"\n清理完成。详细日志请查看 logs/security/memory-cleanup-*.log")


if __name__ == '__main__':
    main()