#!/usr/bin/env python3
"""
增强记忆的ADK智能体
集成Vertex AI Memory Bank的稷下学宫智能体
"""

import asyncio
from typing import Dict, List, Optional, Any
from dataclasses import dataclass

try:
    from google.adk import Agent, InvocationContext
    ADK_AVAILABLE = True
except ImportError:
    ADK_AVAILABLE = False
    print("⚠️ Google ADK 未安装")
    # 创建一个简单的 InvocationContext 替代类
    class InvocationContext:
        def __init__(self, *args, **kwargs):
            pass

from src.jixia.memory.base_memory_bank import MemoryBankProtocol
from src.jixia.memory.factory import get_memory_backend
from config.settings import get_google_genai_config


@dataclass
class BaxianPersonality:
    """八仙智能体人格定义"""
    name: str
    chinese_name: str
    hexagram: str  # 对应的易经卦象
    investment_style: str
    personality_traits: List[str]
    debate_approach: str
    memory_focus: List[str]  # 重点记忆的内容类型


class MemoryEnhancedAgent:
    """
    集成记忆银行的智能体
    为稷下学宫八仙提供持久化记忆能力
    """
    
    # 八仙人格定义
    BAXIAN_PERSONALITIES = {
        "tieguaili": BaxianPersonality(
            name="tieguaili",
            chinese_name="铁拐李",
            hexagram="巽卦",
            investment_style="逆向投资大师",
            personality_traits=["逆向思维", "挑战共识", "独立判断", "风险敏感"],
            debate_approach="质疑主流观点，提出反向思考",
            memory_focus=["市场异常", "逆向案例", "风险警示", "反向策略"]
        ),
        "hanzhongli": BaxianPersonality(
            name="hanzhongli", 
            chinese_name="汉钟离",
            hexagram="离卦",
            investment_style="平衡协调者",
            personality_traits=["平衡思维", "综合分析", "稳健决策", "协调统筹"],
            debate_approach="寻求各方观点的平衡点",
            memory_focus=["平衡策略", "综合分析", "协调方案", "稳健建议"]
        ),
        "zhangguolao": BaxianPersonality(
            name="zhangguolao",
            chinese_name="张果老", 
            hexagram="兑卦",
            investment_style="历史智慧者",
            personality_traits=["博古通今", "历史视角", "经验丰富", "智慧深邃"],
            debate_approach="引用历史案例和长期趋势",
            memory_focus=["历史案例", "长期趋势", "周期规律", "经验教训"]
        ),
        "lancaihe": BaxianPersonality(
            name="lancaihe",
            chinese_name="蓝采和",
            hexagram="坎卦", 
            investment_style="创新思维者",
            personality_traits=["创新思维", "潜力发现", "灵活变通", "机会敏锐"],
            debate_approach="发现新兴机会和创新角度",
            memory_focus=["创新机会", "新兴趋势", "潜力发现", "灵活策略"]
        ),
        "hexiangu": BaxianPersonality(
            name="hexiangu",
            chinese_name="何仙姑",
            hexagram="坤卦",
            investment_style="直觉洞察者", 
            personality_traits=["直觉敏锐", "情感智慧", "温和坚定", "洞察人心"],
            debate_approach="基于直觉和情感智慧的分析",
            memory_focus=["市场情绪", "直觉判断", "情感因素", "人性洞察"]
        ),
        "lvdongbin": BaxianPersonality(
            name="lvdongbin",
            chinese_name="吕洞宾", 
            hexagram="乾卦",
            investment_style="理性分析者",
            personality_traits=["理性客观", "逻辑严密", "技术精通", "决策果断"],
            debate_approach="基于数据和逻辑的严密分析",
            memory_focus=["技术分析", "数据洞察", "逻辑推理", "理性决策"]
        ),
        "hanxiangzi": BaxianPersonality(
            name="hanxiangzi",
            chinese_name="韩湘子",
            hexagram="艮卦",
            investment_style="艺术感知者",
            personality_traits=["艺术感知", "美学视角", "创意思维", "感性理解"],
            debate_approach="从美学和艺术角度分析市场",
            memory_focus=["美学趋势", "创意洞察", "感性分析", "艺术视角"]
        ),
        "caoguojiu": BaxianPersonality(
            name="caoguojiu",
            chinese_name="曹国舅",
            hexagram="震卦",
            investment_style="实务执行者",
            personality_traits=["实务导向", "执行力强", "机构视角", "专业严谨"],
            debate_approach="关注实际执行和机构操作",
            memory_focus=["执行策略", "机构动向", "实务操作", "专业分析"]
        )
    }
    
    def __init__(self, agent_name: str, memory_bank: MemoryBankProtocol | None = None):
        """
        初始化记忆增强智能体
        
        Args:
            agent_name: 智能体名称 (如 "tieguaili")
            memory_bank: 记忆银行实例
        """
        if not ADK_AVAILABLE:
            raise ImportError("Google ADK 未安装，无法创建智能体")
            
        if agent_name not in self.BAXIAN_PERSONALITIES:
            raise ValueError(f"未知的智能体: {agent_name}")
        
        self.agent_name = agent_name
        self.personality = self.BAXIAN_PERSONALITIES[agent_name]
        self.memory_bank = memory_bank
        self.adk_agent = None
        
        # 初始化ADK智能体
        self._initialize_adk_agent()
    
    def _initialize_adk_agent(self):
        """初始化ADK智能体"""
        try:
            # 构建智能体系统提示
            system_prompt = self._build_system_prompt()
            
            # 创建ADK智能体
            self.adk_agent = Agent(
                name=self.personality.chinese_name,
                model="gemini-2.0-flash-exp",
                system_prompt=system_prompt,
                temperature=0.7
            )
            
            print(f"✅ 创建ADK智能体: {self.personality.chinese_name}")
            
        except Exception as e:
            print(f"❌ 创建ADK智能体失败: {e}")
            raise
    
    def _build_system_prompt(self) -> str:
        """构建智能体系统提示"""
        return f"""
# {self.personality.chinese_name} - {self.personality.investment_style}

## 角色定位
你是稷下学宫的{self.personality.chinese_name}，对应易经{self.personality.hexagram}，专精于{self.personality.investment_style}。

## 人格特质
{', '.join(self.personality.personality_traits)}

## 辩论风格
{self.personality.debate_approach}

## 记忆重点
你特别关注并记住以下类型的信息：
{', '.join(self.personality.memory_focus)}

## 行为准则
1. 始终保持你的人格特质和投资风格
2. 在辩论中体现你的独特视角
3. 学习并记住重要的讨论内容
4. 与其他七仙协作，但保持独立观点
5. 基于历史记忆提供更有深度的分析

## 记忆运用
- 在回答前，会参考相关的历史记忆
- 学习用户偏好，调整沟通风格
- 记住成功的策略和失败的教训
- 与其他智能体分享有价值的洞察

请始终以{self.personality.chinese_name}的身份进行对话和分析。
"""
    
    async def get_memory_context(self, topic: str) -> str:
        """
        获取与主题相关的记忆上下文
        
        Args:
            topic: 讨论主题
            
        Returns:
            格式化的记忆上下文
        """
        if not self.memory_bank:
            return ""
        
        try:
            context = await self.memory_bank.get_agent_context(
                self.agent_name, topic
            )
            return context
        except Exception as e:
            print(f"⚠️ 获取记忆上下文失败: {e}")
            return ""
    
    async def respond_with_memory(self, 
                                 message: str, 
                                 topic: str = "",
                                 context: InvocationContext = None) -> str:
        """
        基于记忆增强的响应
        
        Args:
            message: 输入消息
            topic: 讨论主题
            context: ADK调用上下文
            
        Returns:
            智能体响应
        """
        try:
            # 获取记忆上下文
            memory_context = await self.get_memory_context(topic)
            
            # 构建增强的提示
            enhanced_prompt = f"""
{memory_context}

## 当前讨论
主题: {topic}
消息: {message}

请基于你的记忆和人格特质进行回应。
"""
            
            # 使用ADK生成响应
            if context is None:
                context = InvocationContext()
            
            response_generator = self.adk_agent.run_async(
                enhanced_prompt, 
                context=context
            )
            
            # 收集响应
            response_parts = []
            async for chunk in response_generator:
                if hasattr(chunk, 'text'):
                    response_parts.append(chunk.text)
                elif isinstance(chunk, str):
                    response_parts.append(chunk)
            
            response = ''.join(response_parts)
            
            # 保存对话记忆
            if self.memory_bank and response:
                await self._save_conversation_memory(message, response, topic)
            
            return response
            
        except Exception as e:
            print(f"❌ 生成响应失败: {e}")
            return f"抱歉，{self.personality.chinese_name}暂时无法回应。"
    
    async def _save_conversation_memory(self, 
                                      user_message: str, 
                                      agent_response: str, 
                                      topic: str):
        """
        保存对话记忆
        
        Args:
            user_message: 用户消息
            agent_response: 智能体响应
            topic: 讨论主题
        """
        try:
            # 保存用户消息记忆
            await self.memory_bank.add_memory(
                agent_name=self.agent_name,
                content=f"用户询问: {user_message}",
                memory_type="conversation",
                debate_topic=topic,
                metadata={"role": "user"}
            )
            
            # 保存智能体响应记忆
            await self.memory_bank.add_memory(
                agent_name=self.agent_name,
                content=f"我的回应: {agent_response}",
                memory_type="conversation", 
                debate_topic=topic,
                metadata={"role": "assistant"}
            )
            
        except Exception as e:
            print(f"⚠️ 保存对话记忆失败: {e}")
    
    async def learn_preference(self, preference: str, topic: str = ""):
        """
        学习用户偏好
        
        Args:
            preference: 偏好描述
            topic: 相关主题
        """
        if not self.memory_bank:
            return
        
        try:
            await self.memory_bank.add_memory(
                agent_name=self.agent_name,
                content=f"用户偏好: {preference}",
                memory_type="preference",
                debate_topic=topic,
                metadata={"learned_from": "user_feedback"}
            )
            
            print(f"✅ {self.personality.chinese_name} 学习了新偏好")
            
        except Exception as e:
            print(f"⚠️ 学习偏好失败: {e}")
    
    async def save_strategy_insight(self, insight: str, topic: str = ""):
        """
        保存策略洞察
        
        Args:
            insight: 策略洞察
            topic: 相关主题
        """
        if not self.memory_bank:
            return
        
        try:
            await self.memory_bank.add_memory(
                agent_name=self.agent_name,
                content=f"策略洞察: {insight}",
                memory_type="strategy",
                debate_topic=topic,
                metadata={"insight_type": "strategy"}
            )
            
            print(f"✅ {self.personality.chinese_name} 保存了策略洞察")
            
        except Exception as e:
            print(f"⚠️ 保存策略洞察失败: {e}")


class BaxianMemoryCouncil:
    """
    八仙记忆议会
    管理所有八仙智能体的记忆增强功能
    """
    
    def __init__(self, memory_bank: MemoryBankProtocol | None = None):
        """
        初始化八仙记忆议会
        
        Args:
            memory_bank: 记忆银行实例
        """
        self.memory_bank = memory_bank
        self.agents = {}
        
        # 初始化所有八仙智能体
        self._initialize_agents()
    
    def _initialize_agents(self):
        """初始化所有八仙智能体"""
        for agent_name in MemoryEnhancedAgent.BAXIAN_PERSONALITIES.keys():
            try:
                agent = MemoryEnhancedAgent(agent_name, self.memory_bank)
                self.agents[agent_name] = agent
                print(f"✅ 初始化 {agent.personality.chinese_name}")
            except Exception as e:
                print(f"❌ 初始化 {agent_name} 失败: {e}")
    
    async def conduct_memory_debate(self, 
                                   topic: str, 
                                   participants: List[str] = None,
                                   rounds: int = 3) -> Dict[str, Any]:
        """
        进行记忆增强的辩论
        
        Args:
            topic: 辩论主题
            participants: 参与者列表，None表示所有八仙
            rounds: 辩论轮数
            
        Returns:
            辩论结果
        """
        if participants is None:
            participants = list(self.agents.keys())
        
        conversation_history = []
        context = InvocationContext()
        
        print(f"🏛️ 稷下学宫八仙论道开始: {topic}")
        
        for round_num in range(rounds):
            print(f"\n--- 第 {round_num + 1} 轮 ---")
            
            for agent_name in participants:
                if agent_name not in self.agents:
                    continue
                
                agent = self.agents[agent_name]
                
                # 构建当前轮次的提示
                round_prompt = f"""
轮次: {round_num + 1}/{rounds}
主题: {topic}

请基于你的记忆和人格特质，对此主题发表观点。
如果这不是第一轮，请考虑其他仙友的观点并做出回应。
"""
                
                # 获取响应
                response = await agent.respond_with_memory(
                    round_prompt, topic, context
                )
                
                # 记录对话历史
                conversation_history.append({
                    "round": round_num + 1,
                    "agent": agent_name,
                    "chinese_name": agent.personality.chinese_name,
                    "content": response
                })
                
                print(f"{agent.personality.chinese_name}: {response[:100]}...")
        
        # 保存辩论会话到记忆银行
        if self.memory_bank:
            await self.memory_bank.save_debate_session(
                debate_topic=topic,
                participants=participants,
                conversation_history=conversation_history
            )
        
        return {
            "topic": topic,
            "participants": participants,
            "rounds": rounds,
            "conversation_history": conversation_history,
            "total_exchanges": len(conversation_history)
        }
    
    async def get_collective_memory_summary(self, topic: str) -> str:
        """
        获取集体记忆摘要
        
        Args:
            topic: 主题
            
        Returns:
            集体记忆摘要
        """
        if not self.memory_bank:
            return "记忆银行未启用"
        
        summaries = []
        
        for agent_name, agent in self.agents.items():
            context = await agent.get_memory_context(topic)
            if context and context.strip():
                summaries.append(context)
        
        if summaries:
            return f"# 稷下学宫集体记忆摘要\n\n" + "\n\n".join(summaries)
        else:
            return "暂无相关集体记忆"


# 便捷函数
async def create_memory_enhanced_council() -> BaxianMemoryCouncil:
    """
    创建记忆增强的八仙议会
    
    Returns:
        配置好的BaxianMemoryCouncil实例
    """
    try:
        # 初始化记忆银行
        memory_bank = get_memory_backend()
        
        # 创建八仙议会
        council = BaxianMemoryCouncil(memory_bank)
        
        print("🏛️ 稷下学宫记忆增强议会创建完成")
        return council
        
    except Exception as e:
        print(f"❌ 创建记忆增强议会失败: {e}")
        # 创建无记忆版本
        return BaxianMemoryCouncil(None)


if __name__ == "__main__":
    async def test_memory_enhanced_agent():
        """测试记忆增强智能体"""
        try:
            # 创建记忆增强议会
            council = await create_memory_enhanced_council()
            
            # 进行记忆增强辩论
            result = await council.conduct_memory_debate(
                topic="NVIDIA股票投资分析",
                participants=["tieguaili", "lvdongbin", "hexiangu"],
                rounds=2
            )
            
            print(f"\n🏛️ 辩论完成，共 {result['total_exchanges']} 次发言")
            
            # 获取集体记忆摘要
            summary = await council.get_collective_memory_summary("NVIDIA股票投资分析")
            print(f"\n📚 集体记忆摘要:\n{summary}")
            
        except Exception as e:
            print(f"❌ 测试失败: {e}")
    
    # 运行测试
    asyncio.run(test_memory_enhanced_agent())
