"""
Helix Memory Store — Persistent swarm memory.

Inspired by FSI-Edge's Helix Memory architecture.
Stores and retrieves memory across the entire nanobot swarm.
"""

import json
import time
import os
from typing import Dict, List, Optional, Any


class HelixMemoryStore:
    def __init__(self, store_path: str = "/tmp/jedi_memory"):
        self.store_path = store_path
        os.makedirs(store_path, exist_ok=True)
        self._memory = {}
        self._episodes = []
    
    def store(self, key: str, value: Any, memory_type: str = "semantic"):
        """Store a memory entry."""
        self._memory[key] = {
            "value": value,
            "type": memory_type,
            "created_at": time.time(),
            "access_count": 0,
        }
    
    def retrieve(self, key: str) -> Optional[Any]:
        """Retrieve a memory entry."""
        entry = self._memory.get(key)
        if entry:
            entry["access_count"] += 1
            entry["last_accessed"] = time.time()
            return entry["value"]
        return None
    
    def record_episode(self, episode: Dict):
        """Record a memory episode."""
        self._episodes.append({
            **episode,
            "timestamp": time.time(),
        })
    
    def search(self, query: str, memory_type: Optional[str] = None) -> List[Dict]:
        """Search memory by query."""
        results = []
        for key, entry in self._memory.items():
            if query.lower() in key.lower() or query.lower() in json.dumps(entry["value"]).lower():
                if memory_type is None or entry["type"] == memory_type:
                    results.append({"key": key, **entry})
        return results
    
    def get_stats(self) -> Dict:
        return {
            "total_memories": len(self._memory),
            "total_episodes": len(self._episodes),
            "memory_types": {},
        }
    
    def save(self):
        """Persist memory to disk."""
        path = os.path.join(self.store_path, "helix_memory.json")
        with open(path, "w") as f:
            json.dump({"memories": self._memory, "episodes": self._episodes}, f, indent=2)
    
    def load(self):
        """Load memory from disk."""
        path = os.path.join(self.store_path, "helix_memory.json")
        if os.path.exists(path):
            with open(path) as f:
                data = json.load(f)
                self._memory = data.get("memories", {})
                self._episodes = data.get("episodes", [])