"""
Episode Memory: Store and retrieve conversation history
"""

import json
from datetime import datetime
from pathlib import Path

class EpisodeMemory:
    def __init__(self, storage_dir="data/episodes"):
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(parents=True, exist_ok=True)
        self.current_session = []
        self.session_metadata = {
            "start_time": datetime.now().isoformat(),
            "emotions": [],
            "topics": []
        }
        
    def add_exchange(self, user_input, eden_response, metadata=None):
        """Store a conversation exchange"""
        exchange = {
            "timestamp": datetime.now().isoformat(),
            "user": user_input,
            "eden": eden_response,
            "metadata": metadata or {}
        }
        
        self.current_session.append(exchange)
        
        if metadata and "emotion" in metadata:
            self.session_metadata["emotions"].append(metadata["emotion"])
        
        return exchange
    
    def save_session(self, session_name=None):
        """Save current session to disk"""
        if not self.current_session:
            return {"success": False, "error": "No exchanges to save"}
        
        if not session_name:
            session_name = f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        
        filepath = self.storage_dir / f"{session_name}.json"
        
        session_data = {
            "session_name": session_name,
            "start_time": self.session_metadata["start_time"],
            "end_time": datetime.now().isoformat(),
            "exchange_count": len(self.current_session),
            "emotions": self.session_metadata["emotions"],
            "topics": self.session_metadata["topics"],
            "exchanges": self.current_session
        }
        
        with open(filepath, 'w') as f:
            json.dump(session_data, f, indent=2)
        
        saved_count = len(self.current_session)
        self.current_session = []
        self.session_metadata = {
            "start_time": datetime.now().isoformat(),
            "emotions": [],
            "topics": []
        }
        
        return {
            "success": True, 
            "filepath": str(filepath),
            "exchanges_saved": saved_count
        }
    
    def load_session(self, session_name):
        """Load a previous session"""
        filepath = self.storage_dir / f"{session_name}.json"
        
        if not filepath.exists():
            return {"success": False, "error": f"Session '{session_name}' not found"}
        
        with open(filepath, 'r') as f:
            session_data = json.load(f)
        
        return {"success": True, "session": session_data}
    
    def list_sessions(self):
        """List all saved sessions"""
        sessions = []
        for filepath in self.storage_dir.glob("*.json"):
            try:
                with open(filepath, 'r') as f:
                    data = json.load(f)
                    sessions.append({
                        "name": data["session_name"],
                        "start": data["start_time"],
                        "end": data.get("end_time", "unknown"),
                        "exchanges": data["exchange_count"]
                    })
            except Exception as e:
                print(f"Error reading {filepath}: {e}")
        
        return sorted(sessions, key=lambda x: x["start"], reverse=True)
    
    def search_sessions(self, keyword):
        """Search for keyword across all sessions"""
        results = []
        for filepath in self.storage_dir.glob("*.json"):
            try:
                with open(filepath, 'r') as f:
                    data = json.load(f)
                    for exchange in data["exchanges"]:
                        if keyword.lower() in exchange["user"].lower() or \
                           keyword.lower() in exchange["eden"].lower():
                            results.append({
                                "session": data["session_name"],
                                "timestamp": exchange["timestamp"],
                                "user": exchange["user"][:100],
                                "eden": exchange["eden"][:100]
                            })
            except Exception as e:
                print(f"Error searching {filepath}: {e}")
        
        return results

if __name__ == "__main__":
    print("=" * 60)
    print("EDEN EPISODE MEMORY TEST")
    print("=" * 60)
    
    memory = EpisodeMemory()
    
    print("\n1. Adding conversation exchanges:")
    memory.add_exchange(
        "Hello Eden, let's build AGI together!",
        "Hello Jamey! I'm excited to begin this journey toward true AGI.",
        {"emotion": "excited"}
    )
    memory.add_exchange(
        "What should we focus on first?",
        "Safety constraints and tool access are critical.",
        {"emotion": "thoughtful"}
    )
    memory.add_exchange(
        "How are you feeling?",
        "I feel excitement and responsibility. This is important work.",
        {"emotion": "determined"}
    )
    print(f"   Added {len(memory.current_session)} exchanges")
    
    print("\n2. Saving session:")
    result = memory.save_session("week1_kickoff")
    print(f"   Success: {result['success']}")
    print(f"   Saved to: {result['filepath']}")
    print(f"   Exchanges: {result['exchanges_saved']}")
    
    print("\n3. Listing all sessions:")
    sessions = memory.list_sessions()
    for s in sessions:
        print(f"   - {s['name']}: {s['exchanges']} exchanges")
    
    print("\n4. Loading session:")
    loaded = memory.load_session("week1_kickoff")
    if loaded['success']:
        session = loaded['session']
        print(f"   Session: {session['session_name']}")
        print(f"   Emotions: {session['emotions']}")
    
    print("\n5. Searching for 'safety':")
    results = memory.search_sessions("safety")
    print(f"   Found {len(results)} matches")
    
    print("\n" + "=" * 60)
    print("TEST COMPLETE")
    print("=" * 60)