"""
Marziel OS — Memory Manager.

Three-tier memory system inspired by human cognition:

  Working Memory   → CPU registers  → Current session (fast, volatile)
  Long-Term Memory → RAM            → Key facts, preferences (encrypted, persistent)
  Episodic Memory  → Disk           → Conversation summaries (searchable, permanent)

All persistent storage is encrypted via Sentinel SecureVault.
No plaintext ever touches disk.
"""

import hashlib
import json
import logging
import os
import threading
import time
from datetime import datetime
from pathlib import Path
from typing import Optional

logger = logging.getLogger("marziel.memory")

MARZIEL_HOME = os.path.expanduser("~/.marziel")


# ══════════════════════════════════════════════════════════════
# TIER 1: Working Memory (Session-scoped, volatile)
# ══════════════════════════════════════════════════════════════

class WorkingMemory:
    """
    Fast, session-scoped memory. Cleared on restart.

    Like CPU registers — holds the active conversation context,
    recent tool results, and current task state.
    """

    def __init__(self, max_items: int = 256):
        self._store: dict = {}
        self._access_count: dict = {}
        self._max_items = max_items
        self._lock = threading.Lock()
        self._created_at = time.monotonic()

    def store(self, key: str, value, ttl_seconds: int = 0):
        """Store a value. Optional TTL for auto-expiry."""
        with self._lock:
            self._store[key] = {
                "value": value,
                "stored_at": time.monotonic(),
                "ttl": ttl_seconds,
                "access_count": 0,
            }
            self._access_count[key] = 0

            # LRU eviction if over capacity
            if len(self._store) > self._max_items:
                self._evict_lru()

    def recall(self, key: str, default=None):
        """Recall a value. Returns default if expired or missing."""
        with self._lock:
            entry = self._store.get(key)
            if entry is None:
                return default

            # Check TTL
            if entry["ttl"] > 0:
                age = time.monotonic() - entry["stored_at"]
                if age > entry["ttl"]:
                    del self._store[key]
                    return default

            entry["access_count"] += 1
            self._access_count[key] = entry["access_count"]
            return entry["value"]

    def forget(self, key: str):
        """Remove a specific memory."""
        with self._lock:
            self._store.pop(key, None)
            self._access_count.pop(key, None)

    def _evict_lru(self):
        """Evict least-recently-used entries."""
        if not self._access_count:
            return
        lru_key = min(self._access_count, key=self._access_count.get)
        self._store.pop(lru_key, None)
        self._access_count.pop(lru_key, None)

    def clear(self):
        """Clear all working memory."""
        with self._lock:
            self._store.clear()
            self._access_count.clear()

    @property
    def stats(self) -> dict:
        with self._lock:
            return {
                "tier": "working",
                "items": len(self._store),
                "max_items": self._max_items,
                "uptime_sec": round(time.monotonic() - self._created_at, 1),
            }

    def dump(self) -> dict:
        """Dump all keys and their metadata (not values — for privacy)."""
        with self._lock:
            return {
                k: {
                    "type": type(v["value"]).__name__,
                    "access_count": v["access_count"],
                    "age_sec": round(time.monotonic() - v["stored_at"], 1),
                }
                for k, v in self._store.items()
            }


# ══════════════════════════════════════════════════════════════
# TIER 2: Long-Term Memory (Encrypted, persistent)
# ══════════════════════════════════════════════════════════════

class LongTermMemory:
    """
    Encrypted persistent storage using Sentinel SecureVault.

    Stores key facts, user preferences, learned patterns.
    Survives restarts. All data encrypted at rest via AES-256.
    """

    def __init__(self, vault_path: str = None):
        self._vault_path = Path(vault_path or os.path.join(
            MARZIEL_HOME, "memory", "long_term.enc"
        ))
        self._vault_path.parent.mkdir(parents=True, exist_ok=True)
        self._vault = None
        self._lock = threading.Lock()
        self._init_vault()

    def _init_vault(self):
        """Initialize the Sentinel vault for encrypted storage."""
        try:
            from marziel.sentinel import SecureVault
            self._vault = SecureVault(self._vault_path)
            logger.debug("🔐 Long-term memory vault initialized")
        except Exception as e:
            logger.warning(f"Long-term memory vault unavailable: {e}")

    def store(self, key: str, value, category: str = "general"):
        """Store an encrypted memory."""
        with self._lock:
            if not self._vault:
                return False

            memory_entry = {
                "value": value,
                "category": category,
                "stored_at": datetime.utcnow().isoformat() + "Z",
                "updated_at": datetime.utcnow().isoformat() + "Z",
                "access_count": 0,
            }
            self._vault.store(f"mem:{key}", memory_entry)
            logger.debug(f"💾 Long-term memory stored: {key} ({category})")
            return True

    def recall(self, key: str, default=None):
        """Recall a decrypted memory."""
        with self._lock:
            if not self._vault:
                return default

            entry = self._vault.retrieve(f"mem:{key}")
            if entry is None:
                return default

            # Update access count
            entry["access_count"] = entry.get("access_count", 0) + 1
            entry["updated_at"] = datetime.utcnow().isoformat() + "Z"
            self._vault.store(f"mem:{key}", entry)

            return entry.get("value", default)

    def forget(self, key: str):
        """Securely erase a memory."""
        with self._lock:
            if self._vault:
                self._vault.delete(f"mem:{key}")
                logger.debug(f"🗑️ Long-term memory erased: {key}")

    def list_keys(self) -> list:
        """List all stored memory keys."""
        with self._lock:
            if not self._vault:
                return []
            data = self._vault._load_vault()
            return [k.replace("mem:", "") for k in data.keys() if k.startswith("mem:")]

    @property
    def stats(self) -> dict:
        keys = self.list_keys()
        size = self._vault_path.stat().st_size if self._vault_path.exists() else 0
        return {
            "tier": "long_term",
            "items": len(keys),
            "encrypted": True,
            "vault_size_bytes": size,
            "vault_size_kb": round(size / 1024, 1),
        }


# ══════════════════════════════════════════════════════════════
# TIER 3: Episodic Memory (Conversation summaries, searchable)
# ══════════════════════════════════════════════════════════════

class EpisodicMemory:
    """
    Summarized conversation history with semantic search.

    Each "episode" is a compressed summary of a conversation or session.
    Stored as JSON with hashed content for privacy-preserving search.
    """

    def __init__(self, storage_path: str = None, max_episodes: int = 1000):
        self._storage_path = Path(storage_path or os.path.join(
            MARZIEL_HOME, "memory", "episodes.jsonl"
        ))
        self._storage_path.parent.mkdir(parents=True, exist_ok=True)
        self._max_episodes = max_episodes
        self._lock = threading.Lock()
        self._index: list = []
        self._load_index()

    def _load_index(self):
        """Load episode index from disk."""
        if not self._storage_path.exists():
            return
        try:
            with open(self._storage_path, "r") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        try:
                            self._index.append(json.loads(line))
                        except json.JSONDecodeError:
                            pass
            logger.debug(f"📚 Loaded {len(self._index)} episodes from disk")
        except Exception as e:
            logger.warning(f"Episode load failed: {e}")

    def record(self, summary: str, metadata: dict = None):
        """Record a new episode (conversation summary)."""
        episode = {
            "id": hashlib.md5(
                f"{summary}{time.time()}".encode()
            ).hexdigest()[:12],
            "summary": summary,
            "keywords": self._extract_keywords(summary),
            "timestamp": datetime.utcnow().isoformat() + "Z",
            "metadata": metadata or {},
        }

        with self._lock:
            self._index.append(episode)

            # Append to disk
            try:
                with open(self._storage_path, "a") as f:
                    f.write(json.dumps(episode) + "\n")
            except Exception as e:
                logger.warning(f"Episode write failed: {e}")

            # Trim if over limit
            if len(self._index) > self._max_episodes:
                self._index = self._index[-self._max_episodes:]

        logger.debug(f"📝 Episode recorded: {summary[:60]}...")
        return episode["id"]

    def recall(self, query: str, top_k: int = 5) -> list:
        """
        Semantic-ish recall — find episodes matching a query.

        Uses keyword overlap scoring. For production, upgrade to
        cosine similarity on embeddings.
        """
        query_keywords = self._extract_keywords(query)
        if not query_keywords:
            return self._index[-top_k:]

        scored = []
        for episode in self._index:
            ep_keywords = set(episode.get("keywords", []))
            overlap = len(query_keywords & ep_keywords)
            if overlap > 0:
                scored.append((overlap, episode))

        scored.sort(key=lambda x: x[0], reverse=True)
        return [ep for _, ep in scored[:top_k]]

    def _extract_keywords(self, text: str) -> set:
        """Extract significant keywords from text."""
        stopwords = {
            "the", "a", "an", "is", "was", "are", "were", "be", "been",
            "being", "have", "has", "had", "do", "does", "did", "will",
            "would", "could", "should", "may", "might", "shall", "can",
            "to", "of", "in", "for", "on", "with", "at", "by", "from",
            "about", "as", "into", "through", "during", "before", "after",
            "and", "but", "or", "nor", "not", "so", "yet", "both",
            "it", "its", "this", "that", "these", "those", "i", "you",
            "he", "she", "we", "they", "me", "him", "her", "us", "them",
            "my", "your", "his", "our", "their", "what", "which", "who",
            "how", "when", "where", "why", "if", "then", "than",
        }
        words = set(text.lower().split())
        return {
            w.strip(".,!?;:'\"()-[]{}") for w in words
            if len(w) > 2 and w.lower() not in stopwords
        }

    @property
    def stats(self) -> dict:
        size = self._storage_path.stat().st_size if self._storage_path.exists() else 0
        return {
            "tier": "episodic",
            "episodes": len(self._index),
            "max_episodes": self._max_episodes,
            "storage_bytes": size,
            "storage_kb": round(size / 1024, 1),
        }


# ══════════════════════════════════════════════════════════════
# MEMORY MANAGER — Orchestrator
# ══════════════════════════════════════════════════════════════

class MemoryManager:
    """
    Three-tier memory orchestrator for Marziel OS.

    Provides a unified interface for storing, recalling, and
    managing memories across all three tiers.
    """

    def __init__(self):
        self.working = WorkingMemory(max_items=256)
        self.long_term = LongTermMemory()
        self.episodic = EpisodicMemory()

        logger.info("🧠 Memory Manager initialized (3 tiers)")

    def store(self, key: str, value, tier: str = "working", **kwargs):
        """Store a memory in the specified tier."""
        if tier == "working":
            self.working.store(key, value, **kwargs)
        elif tier == "long_term":
            self.long_term.store(key, value, **kwargs)
        elif tier == "episodic":
            self.episodic.record(value, metadata=kwargs.get("metadata"))
        else:
            raise ValueError(f"Unknown memory tier: {tier}")

    def recall(self, key_or_query: str, tier: str = "auto", top_k: int = 5):
        """
        Recall memories. If tier='auto', search all tiers in order:
        working → long_term → episodic.
        """
        if tier == "working":
            return self.working.recall(key_or_query)
        elif tier == "long_term":
            return self.long_term.recall(key_or_query)
        elif tier == "episodic":
            return self.episodic.recall(key_or_query, top_k=top_k)
        elif tier == "auto":
            # Try working memory first (fastest)
            result = self.working.recall(key_or_query)
            if result is not None:
                return result

            # Then long-term
            result = self.long_term.recall(key_or_query)
            if result is not None:
                return result

            # Finally episodic (semantic search)
            episodes = self.episodic.recall(key_or_query, top_k=top_k)
            if episodes:
                return episodes

            return None

    def forget(self, key: str, tier: str = "all"):
        """Forget a memory. If tier='all', erase from all tiers."""
        if tier in ("working", "all"):
            self.working.forget(key)
        if tier in ("long_term", "all"):
            self.long_term.forget(key)

    def consolidate(self):
        """
        Memory consolidation — like human sleep.

        Moves important working memories to long-term storage,
        and summarizes old long-term memories into episodes.
        """
        consolidated = 0
        working_dump = self.working.dump()
        for key, meta in working_dump.items():
            # If accessed multiple times, it's important → promote
            if meta.get("access_count", 0) >= 3:
                value = self.working.recall(key)
                if value is not None:
                    self.long_term.store(key, value, category="consolidated")
                    consolidated += 1

        if consolidated > 0:
            logger.info(f"🌙 Consolidated {consolidated} memories → long-term")

        return consolidated

    @property
    def stats(self) -> dict:
        return {
            "working": self.working.stats,
            "long_term": self.long_term.stats,
            "episodic": self.episodic.stats,
            "total_items": (
                self.working.stats["items"]
                + self.long_term.stats["items"]
                + self.episodic.stats["episodes"]
            ),
        }
