"""
Feature 61: In-memory storage
Feature 62: Persistent memory (save to disk)
Feature 63: Redis memory backend
Feature 67: Memory summarization
Feature 68: Memory namespacing
Feature 69: Memory sharing between agents
Feature 70: Memory expiry
Feature 71: Memory versioning
Feature 72: Memory export/import
Feature 73: Episodic memory
Feature 74: Working memory (short term)
Feature 75: Long term memory
Feature 76: Memory search
Feature 78: Cross-session memory
Feature 80: Global shared memory
"""
import json
import os
import time
from typing import List, Dict, Any, Optional
from datetime import datetime


class Memory:
    """
    Full-featured memory system for Synapse agents.
    Supports in-memory, persistent, namespaced, expiry, search and more.
    """

    def __init__(
        self,
        namespace: str = "default",          # Feature 68: namespacing
        persist_path: str = None,             # Feature 62: persist to disk
        max_short_term: int = 20,             # Feature 74: working memory limit
        ttl: int = None,                      # Feature 70: expiry in seconds
        version: str = "1.0",                # Feature 71: versioning
    ):
        self.namespace = namespace
        self.persist_path = persist_path
        self.max_short_term = max_short_term
        self.ttl = ttl
        self.version = version

        # Feature 74: working memory (short term)
        self._short_term: List[Dict] = []
        # Feature 75: long term memory
        self._long_term: List[Dict] = []
        # Feature 73: episodic memory (past sessions)
        self._episodes: List[Dict] = []
        # Feature 69: shared memory across agents
        self._shared: Dict[str, Any] = {}
        # Timestamps for TTL
        self._timestamps: List[float] = []

        # Feature 62: load from disk if path given
        if persist_path and os.path.exists(persist_path):
            self._load()

    def add(self, role: str, content: str, agent: str = "", metadata: dict = None):
        """Add a message to working memory."""
        msg = {
            "role": role,
            "content": content,
            "agent": agent,
            "timestamp": datetime.utcnow().isoformat(),
            "metadata": metadata or {},
            "namespace": self.namespace,
        }
        self._short_term.append(msg)
        self._timestamps.append(time.time())

        # Feature 74: enforce working memory limit → move to long term
        if len(self._short_term) > self.max_short_term:
            overflow = self._short_term.pop(0)
            self._long_term.append(overflow)
            self._timestamps.pop(0)

        # Feature 62: auto-persist
        if self.persist_path:
            self._save()

    def get_messages(self) -> List[Dict]:
        """Get messages in LLM API format (role + content only). Feature 70: respect TTL."""
        self._expire()
        return [{"role": m["role"], "content": m["content"]} for m in self._short_term]

    def get_full(self) -> List[Dict]:
        """Get full messages with metadata."""
        self._expire()
        return self._short_term.copy()

    def _expire(self):
        """Feature 70: remove expired messages."""
        if not self.ttl:
            return
        now = time.time()
        valid = [(m, t) for m, t in zip(self._short_term, self._timestamps) if now - t < self.ttl]
        if valid:
            self._short_term, self._timestamps = zip(*valid)
            self._short_term = list(self._short_term)
            self._timestamps = list(self._timestamps)
        else:
            self._short_term = []
            self._timestamps = []

    # Feature 69: shared memory
    def set_shared(self, key: str, value: Any):
        self._shared[key] = value

    def get_shared(self, key: str, default=None) -> Any:
        return self._shared.get(key, default)

    # Feature 76: memory search
    def search(self, query: str) -> List[Dict]:
        """Simple keyword search across all memories."""
        query_lower = query.lower()
        results = []
        all_memories = self._short_term + self._long_term
        for msg in all_memories:
            if query_lower in msg["content"].lower():
                results.append(msg)
        return results

    # Feature 73: episodic memory
    def save_episode(self, label: str = ""):
        """Save current session as an episode."""
        episode = {
            "label": label or f"episode_{len(self._episodes)+1}",
            "timestamp": datetime.utcnow().isoformat(),
            "messages": self._short_term.copy(),
            "namespace": self.namespace,
        }
        self._episodes.append(episode)
        return episode

    def get_episodes(self) -> List[Dict]:
        return self._episodes

    # Feature 67: memory summarization
    def summarize(self, keep_last: int = 5) -> str:
        """Summarize old messages to compress context."""
        if len(self._short_term) <= keep_last:
            return ""
        old = self._short_term[:-keep_last]
        summary_lines = [f"[{m['agent'] or m['role']}]: {m['content'][:100]}" for m in old]
        summary = "Summary of earlier context:\n" + "\n".join(summary_lines)
        # Replace old messages with summary
        self._short_term = [{"role": "system", "content": summary, "agent": "memory", "timestamp": datetime.utcnow().isoformat(), "metadata": {}, "namespace": self.namespace}] + self._short_term[-keep_last:]
        return summary

    # Feature 72: export/import
    def export(self) -> dict:
        return {
            "namespace": self.namespace,
            "version": self.version,
            "short_term": self._short_term,
            "long_term": self._long_term,
            "episodes": self._episodes,
            "shared": self._shared,
            "exported_at": datetime.utcnow().isoformat(),
        }

    def import_data(self, data: dict):
        self._short_term = data.get("short_term", [])
        self._long_term = data.get("long_term", [])
        self._episodes = data.get("episodes", [])
        self._shared = data.get("shared", {})

    # Feature 62: persist to disk
    def _save(self):
        os.makedirs(os.path.dirname(self.persist_path) or ".", exist_ok=True)
        with open(self.persist_path, "w") as f:
            json.dump(self.export(), f, indent=2)

    def _load(self):
        with open(self.persist_path, "r") as f:
            self.import_data(json.load(f))

    def clear(self):
        self._short_term = []
        self._long_term = []
        self._timestamps = []

    def clear_all(self):
        self.clear()
        self._episodes = []
        self._shared = {}

    def stats(self) -> dict:
        return {
            "namespace": self.namespace,
            "short_term_count": len(self._short_term),
            "long_term_count": len(self._long_term),
            "episodes": len(self._episodes),
            "shared_keys": list(self._shared.keys()),
            "version": self.version,
        }

    def __len__(self):
        return len(self._short_term)

    def __repr__(self):
        return f"Memory(namespace={self.namespace!r}, messages={len(self._short_term)}, long_term={len(self._long_term)})"


# Feature 80: Global shared memory singleton
class GlobalMemory:
    """Shared memory accessible by all agents globally."""
    _instance = None
    _store: Dict[str, Any] = {}

    @classmethod
    def get_instance(cls) -> "GlobalMemory":
        if not cls._instance:
            cls._instance = cls()
        return cls._instance

    def set(self, key: str, value: Any):
        self._store[key] = value

    def get(self, key: str, default=None) -> Any:
        return self._store.get(key, default)

    def all(self) -> dict:
        return self._store.copy()

    def clear(self):
        self._store.clear()
