#!/usr/bin/env python3
"""
测试记忆功能
参考工作流调用测试总结.txt的测试方法
"""
import sys
sys.path.insert(0, 'backend')

from app.core.database import SessionLocal
from app.models.agent import Agent
from app.models.execution import Execution
from app.models.execution_log import ExecutionLog
import json
from datetime import datetime

def test_memory_functionality():
    """测试记忆功能"""
    db = SessionLocal()
    try:
        # 获取智能聊天助手Agent
        agent = db.query(Agent).filter(
            Agent.name == '智能聊天助手（完整示例）'
        ).first()
        
        if not agent:
            print("❌ 未找到'智能聊天助手（完整示例）'Agent")
            return
        
        print(f"✅ 找到Agent: {agent.name} (ID: {agent.id})")
        print("="*80)
        
        # 获取最近的两次执行（应该对应两次对话）
        executions = db.query(Execution).filter(
            Execution.agent_id == agent.id
        ).order_by(Execution.created_at.desc()).limit(2).all()
        
        if len(executions) < 2:
            print(f"⚠️  只找到 {len(executions)} 次执行，需要至少2次执行来测试记忆功能")
            print("请先进行两次对话测试")
            return
        
        print(f"\n找到 {len(executions)} 次执行记录")
        print("="*80)
        
        # 分析每次执行
        for i, exec_record in enumerate(reversed(executions), 1):  # 按时间正序
            print(f"\n{'='*80}")
            print(f"执行 {i}: {exec_record.id}")
            print(f"输入: {exec_record.input_data}")
            print(f"时间: {exec_record.created_at}")
            print(f"状态: {exec_record.status}")
            
            # 检查关键节点的数据流转
            nodes_to_check = [
                ('cache-query', '查询记忆'),
                ('transform-merge', '合并上下文'),
                ('llm-question', '问题回答'),
                ('cache-update', '更新记忆'),
                ('llm-format', '格式化回复'),
                ('end-1', '最终输出')
            ]
            
            for node_id, label in nodes_to_check:
                # 查找节点执行完成的日志
                log = db.query(ExecutionLog).filter(
                    ExecutionLog.execution_id == exec_record.id,
                    ExecutionLog.node_id == node_id,
                    ExecutionLog.message.like(f'节点 {node_id}%执行完成')
                ).first()
                
                if not log:
                    # 如果没有执行完成的日志，查找开始执行的日志
                    log = db.query(ExecutionLog).filter(
                        ExecutionLog.execution_id == exec_record.id,
                        ExecutionLog.node_id == node_id,
                        ExecutionLog.message.like(f'节点 {node_id}%开始执行')
                    ).first()
                
                if log and log.data:
                    data_key = 'output' if '执行完成' in log.message else 'input'
                    data = log.data.get(data_key, {})
                    
                    print(f"\n{label} ({node_id}):")
                    
                    if isinstance(data, dict):
                        print(f"  keys: {list(data.keys())}")
                        
                        # 检查memory字段
                        if 'memory' in data:
                            memory = data['memory']
                            if isinstance(memory, dict):
                                print(f"  ✅ memory存在，keys: {list(memory.keys())}")
                                if 'conversation_history' in memory:
                                    history = memory['conversation_history']
                                    if isinstance(history, list):
                                        print(f"  ✅ conversation_history: {len(history)} 条")
                                        if history:
                                            print(f"     最新一条: {history[-1].get('content', '')[:50]}")
                                    else:
                                        print(f"  ❌ conversation_history不是list: {type(history)}")
                        # 检查conversation_history字段（可能在顶层）
                        elif 'conversation_history' in data:
                            history = data['conversation_history']
                            if isinstance(history, list):
                                print(f"  ✅ conversation_history在顶层: {len(history)} 条")
                                if history:
                                    print(f"     最新一条: {history[-1].get('content', '')[:50]}")
                        
                        # 对于end节点，检查最终输出
                        if node_id == 'end-1' and 'output' in data:
                            output = data['output']
                            if isinstance(output, str):
                                print(f"  ✅ 最终输出: {output[:200]}")
                                # 检查是否包含名字
                                if '老七' in output:
                                    print(f"  ✅ 输出中包含名字'老七'")
                                else:
                                    print(f"  ❌ 输出中不包含名字'老七'")
                    elif isinstance(data, str):
                        print(f"  输出类型: str, 内容: {data[:200]}")
                        if '老七' in data:
                            print(f"  ✅ 输出中包含名字'老七'")
                        else:
                            print(f"  ❌ 输出中不包含名字'老七'")
                else:
                    print(f"\n{label} ({node_id}): ❌ 未找到执行日志")
            
            # 检查最终输出
            if exec_record.output_data:
                output_data = exec_record.output_data
                if isinstance(output_data, dict):
                    result = output_data.get('result', '')
                    if isinstance(result, str):
                        print(f"\n最终结果: {result[:200]}")
                        if '老七' in result:
                            print(f"✅ 最终结果中包含名字'老七'")
                        else:
                            print(f"❌ 最终结果中不包含名字'老七'")
        
        # 检查Redis中的记忆数据
        print(f"\n{'='*80}")
        print("检查Redis中的记忆数据:")
        try:
            from app.core.redis_client import get_redis_client
            redis_client = get_redis_client()
            if redis_client:
                keys = redis_client.keys('user_memory_*')
                if keys:
                    for key in keys:
                        value = redis_client.get(key)
                        if value:
                            try:
                                memory_data = json.loads(value)
                                if 'conversation_history' in memory_data:
                                    history = memory_data['conversation_history']
                                    print(f"  ✅ {key}: {len(history)} 条对话记录")
                                    if history:
                                        print(f"     最新一条: {history[-1].get('content', '')[:50]}")
                            except:
                                print(f"  ⚠️  {key}: 无法解析JSON")
                else:
                    print("  ❌ Redis中没有找到记忆数据")
            else:
                print("  ⚠️  Redis客户端不可用")
        except Exception as e:
            print(f"  ⚠️  检查Redis失败: {str(e)}")
        
        print(f"\n{'='*80}")
        print("测试完成")
        
    finally:
        db.close()

if __name__ == '__main__':
    test_memory_functionality()