#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
稷下学宫 ADK Memory Bank 论道系统
实现带有记忆银行的八仙智能体辩论
"""

import os
import asyncio
from google.adk import Agent, Runner
from google.adk.sessions import InMemorySessionService
from google.adk.memory import VertexAiMemoryBankService
from google.adk.memory.memory_entry import MemoryEntry
from google.genai import types
import json
from datetime import datetime
from typing import Dict, List, Optional

class BaxianMemoryManager:
    """八仙记忆管理器"""
    
    def __init__(self):
        self.memory_services: Dict[str, VertexAiMemoryBankService] = {}
        self.agents: Dict[str, Agent] = {}
        
    async def initialize_baxian_agents(self):
        """初始化八仙智能体及其记忆银行"""
        # 从环境变量获取项目ID和位置
        project_id = os.getenv('GOOGLE_CLOUD_PROJECT_ID')
        location = os.getenv('GOOGLE_CLOUD_LOCATION', 'us-central1')
        
        if not project_id:
            raise ValueError("未设置 GOOGLE_CLOUD_PROJECT_ID 环境变量")
        
        # 八仙角色配置
        baxian_config = {
            "铁拐李": {
                "instruction": "你是铁拐李，八仙中的逆向思维专家。你善于从批判和质疑的角度看问题，总是能发现事物的另一面。你会从你的记忆中回忆相关的逆向投资案例和失败教训。",
                "memory_context": "逆向投资案例、市场泡沫警告、风险识别经验"
            },
            "吕洞宾": {
                "instruction": "你是吕洞宾，八仙中的理性分析者。你善于平衡各方观点，用理性和逻辑来分析问题。你会从记忆中调用技术分析的成功案例和理论知识。",
                "memory_context": "技术分析理论、成功预测案例、市场趋势分析"
            },
            "何仙姑": {
                "instruction": "你是何仙姑，八仙中的风险控制专家。你总是从风险管理的角度思考问题，善于发现潜在危险。你会回忆历史上的重大风险事件。",
                "memory_context": "风险管理案例、黑天鹅事件、危机预警经验"
            },
            "张果老": {
                "instruction": "你是张果老，八仙中的历史智慧者。你善于从历史数据中寻找规律和智慧，总是能提供长期视角。你会从记忆中调用历史数据和长期趋势。",
                "memory_context": "历史市场数据、长期投资趋势、周期性规律"
            }
        }
        
        # 为每个仙人创建智能体和记忆服务
        for name, config in baxian_config.items():
            # 创建记忆服务
            memory_service = VertexAiMemoryBankService(
                project=project_id,
                location=location
            )
            
            # 初始化记忆内容
            await self._initialize_agent_memory(memory_service, name, config['memory_context'])
            
            # 创建智能体
            agent = Agent(
                name=name,
                model="gemini-2.5-flash",
                instruction=f"{config['instruction']} 在回答时，请先从你的记忆银行中检索相关信息，然后结合当前话题给出回应。",
                memory_service=memory_service
            )
            
            self.memory_services[name] = memory_service
            self.agents[name] = agent
            
        print(f"✅ 已初始化 {len(self.agents)} 个八仙智能体及其记忆服务")
        
    async def _initialize_agent_memory(self, memory_service: VertexAiMemoryBankService, agent_name: str, context: str):
        """为智能体初始化记忆内容"""
        
        # 根据角色添加初始记忆
        initial_memories = {
            "铁拐李": [
                "2000年互联网泡沫破裂，许多高估值科技股暴跌90%以上",
                "2008年金融危机前，房地产市场过度繁荣，逆向思维者提前撤离",
                "比特币从2万美元跌到3千美元，提醒我们任何资产都可能大幅回调",
                "巴菲特说过：别人贪婪时我恐惧，别人恐惧时我贪婪"
            ],
            "吕洞宾": [
                "移动平均线交叉是经典的技术分析信号",
                "RSI指标超过70通常表示超买，低于30表示超卖",
                "支撑位和阻力位是技术分析的核心概念",
                "成功的技术分析需要结合多个指标综合判断"
            ],
            "何仙姑": [
                "2008年雷曼兄弟倒闭引发全球金融危机",
                "长期资本管理公司(LTCM)的失败说明了风险管理的重要性",
                "分散投资是降低风险的基本原则",
                "黑天鹅事件虽然罕见但影响巨大，需要提前准备"
            ],
            "张果老": [
                "股市存在7-10年的长期周期",
                "康德拉季耶夫长波理论描述了50-60年的经济周期",
                "历史上每次重大技术革命都带来新的投资机会",
                "长期来看，优质资产总是向上的"
            ]
        }
        
        memories = initial_memories.get(agent_name, [])
        for memory_text in memories:
            memory_entry = MemoryEntry(
                content=memory_text,
                metadata={
                    "agent": agent_name,
                    "type": "historical_knowledge",
                    "timestamp": datetime.now().isoformat()
                }
            )
            # 注意：VertexAiMemoryBankService 的 add_memory 方法可能需要不同的参数
            # 这里假设它有一个类似的方法
            await memory_service.add_memory(memory_entry)
            
    async def add_debate_memory(self, agent_name: str, content: str, topic: str):
        """为智能体添加辩论记忆"""
        if agent_name in self.memory_services:
            memory_entry = MemoryEntry(
                content=content,
                metadata={
                    "agent": agent_name,
                    "type": "debate_history",
                    "topic": topic,
                    "timestamp": datetime.now().isoformat()
                }
            )
            # 注意：VertexAiMemoryBankService 的 add_memory 方法可能需要不同的参数
            # 这里假设它有一个类似的方法
            await self.memory_services[agent_name].add_memory(memory_entry)
            
    async def retrieve_relevant_memories(self, agent_name: str, query: str, limit: int = 3) -> List[str]:
        """检索智能体的相关记忆"""
        if agent_name not in self.memory_services:
            return []
            
        try:
            # 注意：VertexAiMemoryBankService 的 search 方法可能需要不同的参数
            # 这里假设它有一个类似的方法
            memories = await self.memory_services[agent_name].search(query, limit=limit)
            return [memory.content for memory in memories]
        except Exception as e:
            print(f"⚠️ 记忆检索失败 ({agent_name}): {e}")
            return []

class MemoryEnhancedDebate:
    """带记忆增强的辩论系统"""
    
    def __init__(self):
        self.memory_manager = BaxianMemoryManager()
        self.session_service = InMemorySessionService()
        self.runners: Dict[str, Runner] = {}
        
    async def initialize(self):
        """初始化辩论系统"""
        await self.memory_manager.initialize_baxian_agents()
        
        # 创建会话
        self.session = await self.session_service.create_session(
            state={},
            app_name="稷下学宫记忆增强论道系统",
            user_id="memory_debate_user"
        )
        
        # 为每个智能体创建Runner
        for name, agent in self.memory_manager.agents.items():
            runner = Runner(
                app_name="稷下学宫记忆增强论道系统",
                agent=agent,
                session_service=self.session_service
            )
            self.runners[name] = runner
            
    async def conduct_memory_debate(self, topic: str, participants: List[str] = None):
        """进行带记忆的辩论"""
        if participants is None:
            participants = ["铁拐李", "吕洞宾", "何仙姑", "张果老"]
            
        print(f"\n🎭 稷下学宫记忆增强论道开始...")
        print(f"📋 论道主题: {topic}")
        print(f"🎯 参与仙人: {', '.join(participants)}")
        
        debate_history = []
        
        for round_num in range(2):  # 进行2轮辩论
            print(f"\n🔄 第 {round_num + 1} 轮论道:")
            
            for participant in participants:
                if participant not in self.runners:
                    continue
                    
                print(f"\n🗣️  {participant} 发言:")
                
                # 检索相关记忆
                relevant_memories = await self.memory_manager.retrieve_relevant_memories(
                    participant, topic, limit=2
                )
                
                # 构建包含记忆的提示
                memory_context = ""
                if relevant_memories:
                    memory_context = f"\n从你的记忆中回忆到：\n" + "\n".join([f"- {memory}" for memory in relevant_memories])
                
                # 构建辩论历史上下文
                history_context = ""
                if debate_history:
                    recent_history = debate_history[-3:]  # 最近3条发言
                    history_context = f"\n最近的论道内容：\n" + "\n".join([f"- {h}" for h in recent_history])
                
                prompt = f"关于'{topic}'这个话题{memory_context}{history_context}\n\n请结合你的记忆和当前讨论，从你的角色特点出发发表观点。请控制在150字以内。"
                
                # 发送消息并获取回复
                content = types.Content(role='user', parts=[types.Part(text=prompt)])
                response = self.runners[participant].run_async(
                    user_id=self.session.user_id,
                    session_id=self.session.id,
                    new_message=content
                )
                
                # 收集回复
                reply = ""
                async for event in response:
                    if hasattr(event, 'content') and event.content:
                        if hasattr(event.content, 'parts') and event.content.parts:
                            for part in event.content.parts:
                                if hasattr(part, 'text') and part.text:
                                    reply += str(part.text)
                
                if reply.strip():
                    clean_reply = reply.strip()
                    print(f"   {clean_reply}")
                    
                    # 记录到辩论历史
                    debate_entry = f"{participant}: {clean_reply}"
                    debate_history.append(debate_entry)
                    
                    # 添加到记忆银行
                    await self.memory_manager.add_debate_memory(
                        participant, clean_reply, topic
                    )
                    
                await asyncio.sleep(1)  # 避免API调用过快
                
        print(f"\n🎉 记忆增强论道完成！")
        print(f"📝 本次论道共产生 {len(debate_history)} 条发言，已存储到各仙人的记忆银行中。")
        
        return debate_history
        
    async def close(self):
        """关闭资源"""
        for runner in self.runners.values():
            await runner.close()

async def main():
    """主函数"""
    print("🚀 稷下学宫 ADK Memory Bank 论道系统")
    
    # 检查API密钥
    api_key = os.getenv('GOOGLE_API_KEY')
    if not api_key:
        print("❌ 未找到 GOOGLE_API_KEY 环境变量")
        print("请使用: doppler run -- python src/jixia/debates/adk_memory_debate.py")
        return
    
    print(f"✅ API密钥已配置")
    
    # 创建并初始化辩论系统
    debate_system = MemoryEnhancedDebate()
    
    try:
        await debate_system.initialize()
        
        # 进行辩论
        await debate_system.conduct_memory_debate(
            topic="人工智能对投资市场的影响",
            participants=["铁拐李", "吕洞宾", "何仙姑", "张果老"]
        )
        
    except Exception as e:
        print(f"❌ 运行失败: {e}")
        import traceback
        traceback.print_exc()
    finally:
        await debate_system.close()

if __name__ == "__main__":
    asyncio.run(main())