#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Memory Bank 简化测试脚本
"""

import os
import asyncio
from datetime import datetime
import json

# Google GenAI 导入
import google.genai as genai

class MemoryBankTest:
    """Memory Bank 测试类"""
    
    def __init__(self):
        self.api_key = os.getenv('GOOGLE_API_KEY')
        if not self.api_key:
            raise ValueError("请设置 GOOGLE_API_KEY 环境变量")
        
        self.client = genai.Client(api_key=self.api_key)
        
        # 八仙人格基线
        self.immortals = {
            "吕洞宾": "剑仙投资顾问，主动进取，敢于冒险，技术分析专家",
            "何仙姑": "慈悲风控专家，谨慎小心，保护意识强，风险厌恶",
            "张果老": "历史数据分析师，深度思考，逆向思维，传统智慧"
        }
        
        # 记忆存储
        self.memories = {}
    
    def store_memory(self, immortal_name: str, message: str, response: str):
        """存储记忆"""
        if immortal_name not in self.memories:
            self.memories[immortal_name] = []
        
        self.memories[immortal_name].append({
            "message": message,
            "response": response,
            "timestamp": datetime.now().isoformat()
        })
    
    def chat_with_immortal(self, immortal_name: str, message: str) -> str:
        """与仙人对话"""
        # 构建上下文
        context = f"你是{immortal_name}，{self.immortals[immortal_name]}。"
        
        # 添加记忆
        if immortal_name in self.memories and self.memories[immortal_name]:
            context += "\n\n你的历史对话："
            for memory in self.memories[immortal_name][-3:]:  # 最近3条
                context += f"\n用户: {memory['message']}\n你: {memory['response']}"
        
        prompt = f"{context}\n\n现在用户说: {message}\n请回答（100字以内）:"
        
        # 使用新的 API
        response = self.client.models.generate_content(
            model="gemini-2.0-flash-exp",
            contents=[{"parts": [{"text": prompt}]}]
        )
        return response.candidates[0].content.parts[0].text
    
    def test_memory_continuity(self):
        """测试记忆连续性"""
        print("🧪 测试记忆连续性")
        print("=" * 50)
        
        # 测试吕洞宾
        print("\n🎭 测试吕洞宾:")
        messages = [
            "你觉得特斯拉股票怎么样？",
            "现在市场波动很大，你怎么看？",
            "你之前不是看好特斯拉吗？现在还是这个观点吗？"
        ]
        
        for message in messages:
            print(f"\n用户: {message}")
            response = self.chat_with_immortal("吕洞宾", message)
            print(f"吕洞宾: {response}")
            self.store_memory("吕洞宾", message, response)
        
        # 测试何仙姑
        print("\n🎭 测试何仙姑:")
        messages = [
            "特斯拉股票风险大吗？",
            "现在适合投资吗？",
            "你一直很谨慎，现在还是建议观望吗？"
        ]
        
        for message in messages:
            print(f"\n用户: {message}")
            response = self.chat_with_immortal("何仙姑", message)
            print(f"何仙姑: {response}")
            self.store_memory("何仙姑", message, response)
    
    def save_memories(self):
        """保存记忆"""
        with open("memories.json", "w", encoding="utf-8") as f:
            json.dump(self.memories, f, ensure_ascii=False, indent=2)
        print("💾 记忆已保存到 memories.json")

def main():
    """主函数"""
    print("🚀 Memory Bank 测试开始")
    
    test = MemoryBankTest()
    test.test_memory_continuity()
    test.save_memories()
    
    print("\n✅ 测试完成！")

if __name__ == "__main__":
    main()