"""Multi-layer memory system for Open Agent.

Provides a hierarchical memory architecture:
- Core Identity: Fundamental user traits (rarely changes)
- Long-term Memory: Learned preferences and capabilities (evolves slowly)
- Short-term Memory: Recent context and session state (ephemeral)
"""

import json
from dataclasses import dataclass, field, asdict
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional, Dict, List, Any, Set
from collections import deque
import hashlib


@dataclass
class MemoryEntry:
    """A single memory entry with metadata."""

    content: str
    context: str
    importance: float = 1.0  # 0.0-1.0, higher = more important
    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
    last_accessed: str = field(default_factory=lambda: datetime.now().isoformat())
    access_count: int = 1
    source: str = "user"  # user, inference, system
    tags: List[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict) -> "MemoryEntry":
        return cls(**data)

    def touch(self) -> None:
        """Update last accessed time and count."""
        self.last_accessed = datetime.now().isoformat()
        self.access_count += 1


class CoreIdentity:
    """Core identity layer - fundamental user traits that rarely change.

    This is the deepest memory layer containing:
    - User's role/profession
    - Core values and principles
    - Fundamental working style
    - Long-term goals
    - Key relationships/teams
    """

    MAX_ENTRIES = 20  # Keep core identity focused

    def __init__(self, storage_path: Optional[Path] = None):
        self.entries: Dict[str, MemoryEntry] = {}
        self._storage_path = storage_path or Path(".emdash/memory/core_identity.json")
        self._storage_path.parent.mkdir(parents=True, exist_ok=True)
        self._load()

    def _load(self) -> None:
        """Load core identity from storage."""
        if self._storage_path.exists():
            try:
                with open(self._storage_path, "r") as f:
                    data = json.load(f)
                self.entries = {
                    k: MemoryEntry.from_dict(v) for k, v in data.get("entries", {}).items()
                }
            except Exception as e:
                print(f"Warning: Could not load core identity: {e}")

    def _save(self) -> None:
        """Save core identity to storage."""
        data = {
            "layer": "core_identity",
            "updated_at": datetime.now().isoformat(),
            "entries": {k: v.to_dict() for k, v in self.entries.items()},
        }
        with open(self._storage_path, "w") as f:
            json.dump(data, f, indent=2)

    def store(self, key: str, content: str, importance: float = 1.0, tags: Optional[List[str]] = None) -> None:
        """Store a core identity trait.

        Args:
            key: Identifier for this trait (e.g., "role", "values", "goals")
            content: The identity information
            importance: How fundamental this is (0.0-1.0)
            tags: Optional categorization tags
        """
        # Check if at capacity and this is new
        if key not in self.entries and len(self.entries) >= self.MAX_ENTRIES:
            # Remove least important entry
            least_important = min(self.entries.items(), key=lambda x: x[1].importance)
            if least_important[1].importance < importance:
                del self.entries[least_important[0]]
            else:
                return  # Don't store if not important enough

        self.entries[key] = MemoryEntry(
            content=content,
            context="core_identity",
            importance=importance,
            source="user",
            tags=tags or [],
        )
        self._save()

    def get(self, key: str) -> Optional[MemoryEntry]:
        """Get a core identity entry."""
        entry = self.entries.get(key)
        if entry:
            entry.touch()
            self._save()
        return entry

    def get_all(self) -> Dict[str, MemoryEntry]:
        """Get all core identity entries sorted by importance."""
        return dict(sorted(self.entries.items(), key=lambda x: -x[1].importance))

    def remove(self, key: str) -> bool:
        """Remove a core identity entry."""
        if key in self.entries:
            del self.entries[key]
            self._save()
            return True
        return False

    def build_prompt_section(self) -> str:
        """Build core identity section for system prompt."""
        if not self.entries:
            return ""

        lines = ["## Core Identity", ""]
        for key, entry in self.get_all().items():
            lines.append(f"**{key.replace('_', ' ').title()}**: {entry.content}")
        lines.append("")

        return "\n".join(lines)


class LongTermMemory:
    """Long-term memory layer - learned preferences and patterns.

    This layer contains:
    - Learned preferences (coding style, communication style)
    - Recurring patterns and workflows
    - Successful strategies that worked for user
    - Important facts about projects/contexts

    Uses importance scoring and access patterns for retention.
    """

    MAX_ENTRIES = 100
    DECAY_DAYS = 90  # Entries decay after 90 days without access

    def __init__(self, storage_path: Optional[Path] = None):
        self.entries: Dict[str, MemoryEntry] = {}
        self._storage_path = storage_path or Path(".emdash/memory/long_term.json")
        self._storage_path.parent.mkdir(parents=True, exist_ok=True)
        self._load()

    def _load(self) -> None:
        """Load long-term memory from storage."""
        if self._storage_path.exists():
            try:
                with open(self._storage_path, "r") as f:
                    data = json.load(f)
                self.entries = {
                    k: MemoryEntry.from_dict(v) for k, v in data.get("entries", {}).items()
                }
                self._decay_entries()
            except Exception as e:
                print(f"Warning: Could not load long-term memory: {e}")

    def _save(self) -> None:
        """Save long-term memory to storage."""
        data = {
            "layer": "long_term",
            "updated_at": datetime.now().isoformat(),
            "entry_count": len(self.entries),
            "entries": {k: v.to_dict() for k, v in self.entries.items()},
        }
        with open(self._storage_path, "w") as f:
            json.dump(data, f, indent=2)

    def _decay_entries(self) -> None:
        """Remove or decay old entries based on last access."""
        now = datetime.now()
        to_remove = []

        for key, entry in self.entries.items():
            last_access = datetime.fromisoformat(entry.last_accessed)
            days_since = (now - last_access).days

            if days_since > self.DECAY_DAYS:
                # Decay importance based on time
                decay_factor = 0.9 ** (days_since / 30)  # Decay monthly
                entry.importance *= decay_factor

                if entry.importance < 0.1:
                    to_remove.append(key)

        for key in to_remove:
            del self.entries[key]

    def store(
        self,
        key: str,
        content: str,
        importance: float = 0.5,
        context: str = "general",
        tags: Optional[List[str]] = None,
        source: str = "inference",
    ) -> None:
        """Store a long-term memory.

        Args:
            key: Memory identifier
            content: The information to store
            importance: Retention priority (0.0-1.0)
            context: Category/domain
            tags: Optional tags
            source: How this was learned (user, inference, system)
        """
        # Check capacity
        if key not in self.entries and len(self.entries) >= self.MAX_ENTRIES:
            # Remove least important and least accessed
            sorted_entries = sorted(
                self.entries.items(),
                key=lambda x: (x[1].importance, x[1].access_count)
            )
            if sorted_entries[0][1].importance < importance:
                del self.entries[sorted_entries[0][0]]
            else:
                return

        self.entries[key] = MemoryEntry(
            content=content,
            context=context,
            importance=importance,
            source=source,
            tags=tags or [],
        )
        self._save()

    def get(self, key: str) -> Optional[MemoryEntry]:
        """Retrieve a memory and update access stats."""
        entry = self.entries.get(key)
        if entry:
            entry.touch()
            self._save()
        return entry

    def query(self, query_text: str, context: Optional[str] = None, top_k: int = 5) -> List[tuple]:
        """Query memories by relevance to text.

        Simple keyword matching - could be enhanced with embeddings.
        """
        query_lower = query_text.lower()
        words = set(query_lower.split())

        scored = []
        for key, entry in self.entries.items():
            score = 0.0

            # Context match
            if context and entry.context == context:
                score += 0.3

            # Content word overlap
            entry_words = set(entry.content.lower().split())
            overlap = len(words & entry_words) / len(words | entry_words)
            score += overlap * 0.5

            # Tag match
            entry_tags = set(t.lower() for t in entry.tags)
            query_tags = words  # Treat query words as potential tags
            tag_overlap = len(entry_tags & query_tags)
            score += tag_overlap * 0.2

            # Importance boost
            score += entry.importance * 0.1

            # Recency boost
            days_since = (datetime.now() - datetime.fromisoformat(entry.last_accessed)).days
            if days_since < 7:
                score += 0.1

            scored.append((key, entry, score))

        # Sort by score and return top_k
        scored.sort(key=lambda x: -x[2])
        return [(k, e) for k, e, _ in scored[:top_k]]

    def get_by_context(self, context: str) -> Dict[str, MemoryEntry]:
        """Get all memories for a specific context."""
        return {
            k: v for k, v in self.entries.items()
            if v.context == context
        }

    def build_prompt_section(self, context: Optional[str] = None, max_entries: int = 10) -> str:
        """Build long-term memory section for system prompt."""
        entries = self.entries.values()

        if context:
            entries = [e for e in entries if e.context == context]

        # Sort by importance and recency
        sorted_entries = sorted(
            entries,
            key=lambda e: (e.importance, e.access_count),
            reverse=True
        )[:max_entries]

        if not sorted_entries:
            return ""

        lines = ["## Learned Preferences & Patterns", ""]

        # Group by context
        by_context: Dict[str, List[MemoryEntry]] = {}
        for entry in sorted_entries:
            ctx = entry.context
            if ctx not in by_context:
                by_context[ctx] = []
            by_context[ctx].append(entry)

        for ctx, ctx_entries in sorted(by_context.items()):
            lines.append(f"### {ctx.title()}")
            for entry in ctx_entries:
                conf_str = ""
                if entry.importance < 0.7:
                    conf_str = f" (confidence: {entry.importance:.0%})"
                lines.append(f"- {entry.content}{conf_str}")
            lines.append("")

        return "\n".join(lines)


class ShortTermMemory:
    """Short-term memory layer - recent context and session state.

    This layer contains:
    - Recent conversation snippets
    - Current task/project context
    - Temporary working memory
    - Session-specific facts

    Automatically pruned based on age and capacity.
    """

    MAX_ENTRIES = 50
    MAX_AGE_HOURS = 24  # Entries expire after 24 hours

    def __init__(self, session_id: str, storage_path: Optional[Path] = None):
        self.session_id = session_id
        self.entries: deque = deque(maxlen=self.MAX_ENTRIES)
        self._storage_path = storage_path or Path(f".emdash/memory/short_term/{session_id}.json")
        self._storage_path.parent.mkdir(parents=True, exist_ok=True)
        self._load()

    def _load(self) -> None:
        """Load short-term memory if recent enough."""
        if self._storage_path.exists():
            try:
                with open(self._storage_path, "r") as f:
                    data = json.load(f)

                # Check if session is still fresh
                session_time = datetime.fromisoformat(data.get("session_time", "1970-01-01"))
                age_hours = (datetime.now() - session_time).total_seconds() / 3600

                if age_hours < self.MAX_AGE_HOURS:
                    entries_data = data.get("entries", [])
                    self.entries = deque(
                        [MemoryEntry.from_dict(e) for e in entries_data],
                        maxlen=self.MAX_ENTRIES
                    )
                else:
                    # Session expired, clear it
                    self._storage_path.unlink(missing_ok=True)
            except Exception as e:
                print(f"Warning: Could not load short-term memory: {e}")

    def _save(self) -> None:
        """Save short-term memory to storage."""
        data = {
            "layer": "short_term",
            "session_id": self.session_id,
            "session_time": datetime.now().isoformat(),
            "entries": [e.to_dict() for e in self.entries],
        }
        with open(self._storage_path, "w") as f:
            json.dump(data, f, indent=2)

    def add(self, content: str, context: str = "session", importance: float = 0.3) -> None:
        """Add an entry to short-term memory."""
        entry = MemoryEntry(
            content=content,
            context=context,
            importance=importance,
            source="session",
        )
        self.entries.append(entry)
        self._save()

    def get_recent(self, n: int = 10, context: Optional[str] = None) -> List[MemoryEntry]:
        """Get most recent entries, optionally filtered by context."""
        entries = list(self.entries)
        if context:
            entries = [e for e in entries if e.context == context]
        return entries[-n:]

    def clear(self) -> None:
        """Clear short-term memory."""
        self.entries.clear()
        self._storage_path.unlink(missing_ok=True)

    def build_context_string(self, max_entries: int = 5) -> str:
        """Build recent context string for prompt augmentation."""
        recent = self.get_recent(max_entries)
        if not recent:
            return ""

        lines = ["## Recent Context", ""]
        for entry in recent:
            lines.append(f"- [{entry.context}] {entry.content}")

        return "\n".join(lines)


class HierarchicalMemory:
    """Unified interface to all memory layers.

    Provides:
    - Core identity (persistent, fundamental)
    - Long-term memory (learned patterns, decays slowly)
    - Short-term memory (session context, ephemeral)
    """

    def __init__(
        self,
        session_id: str,
        emdash_dir: Optional[Path] = None,
    ):
        self._emdash_dir = emdash_dir or Path(".emdash")
        self._memory_dir = self._emdash_dir / "memory"
        self._memory_dir.mkdir(parents=True, exist_ok=True)

        # Initialize layers
        self.core = CoreIdentity(self._memory_dir / "core_identity.json")
        self.long_term = LongTermMemory(self._memory_dir / "long_term.json")
        self.short_term = ShortTermMemory(session_id, self._memory_dir / f"short_term/{session_id}.json")

    def store(
        self,
        content: str,
        layer: str = "long_term",
        key: Optional[str] = None,
        importance: float = 0.5,
        context: str = "general",
        tags: Optional[List[str]] = None,
    ) -> str:
        """Store a memory in the appropriate layer.

        Args:
            content: The information to store
            layer: Which layer ("core", "long_term", "short_term")
            key: Optional identifier (auto-generated if not provided)
            importance: How important this memory is
            context: Category/domain
            tags: Optional tags

        Returns:
            The key used to store the memory
        """
        if key is None:
            # Generate key from content hash
            key = hashlib.md5(content.encode()).hexdigest()[:12]

        if layer == "core":
            self.core.store(key, content, importance, tags)
        elif layer == "long_term":
            self.long_term.store(key, content, importance, context, tags)
        elif layer == "short_term":
            self.short_term.add(content, context, importance)

        return key

    def query(
        self,
        query: str,
        layers: Optional[List[str]] = None,
        context: Optional[str] = None,
        top_k: int = 5,
    ) -> Dict[str, List[tuple]]:
        """Query all memory layers.

        Args:
            query: Search query
            layers: Which layers to search (defaults to all)
            context: Optional context filter
            top_k: Max results per layer

        Returns:
            Dict mapping layer names to results
        """
        layers = layers or ["core", "long_term", "short_term"]
        results = {}

        if "core" in layers:
            # Simple keyword match for core
            core_results = []
            query_lower = query.lower()
            for k, entry in self.core.entries.items():
                if query_lower in entry.content.lower():
                    core_results.append((k, entry))
            results["core"] = core_results[:top_k]

        if "long_term" in layers:
            results["long_term"] = self.long_term.query(query, context, top_k)

        if "short_term" in layers:
            # Recent entries from short-term
            recent = self.short_term.get_recent(top_k, context)
            results["short_term"] = [(f"session_{i}", e) for i, e in enumerate(recent)]

        return results

    def build_system_prompt(self, context: Optional[str] = None) -> str:
        """Build complete memory-augmented system prompt."""
        sections = []

        # Core identity (always include if exists)
        core_section = self.core.build_prompt_section()
        if core_section:
            sections.append(core_section)

        # Long-term preferences
        long_term_section = self.long_term.build_prompt_section(context)
        if long_term_section:
            sections.append(long_term_section)

        # Short-term context
        short_term_section = self.short_term.build_context_string()
        if short_term_section:
            sections.append(short_term_section)

        return "\n\n".join(sections)

    def consolidate_short_term(self) -> List[Dict]:
        """Analyze short-term memory and suggest long-term consolidation.

        Returns:
            List of suggested memories to promote to long-term
        """
        suggestions = []

        # Find frequently mentioned topics in short-term
        from collections import Counter

        all_content = " ".join([e.content for e in self.short_term.entries])
        words = [w.lower() for w in all_content.split() if len(w) > 4]
        frequent = Counter(words).most_common(10)

        # Check for patterns that should be remembered
        for entry in self.short_term.entries:
            # High importance short-term memories are candidates
            if entry.importance > 0.7:
                suggestions.append({
                    "content": entry.content,
                    "context": entry.context,
                    "reason": "high_importance",
                    "suggested_layer": "long_term",
                })

        return suggestions

    def export_all(self) -> dict:
        """Export all memory layers."""
        return {
            "core_identity": {
                "entries": {k: v.to_dict() for k, v in self.core.entries.items()},
            },
            "long_term": {
                "entries": {k: v.to_dict() for k, v in self.long_term.entries.items()},
            },
            "short_term": {
                "session_id": self.short_term.session_id,
                "entries": [e.to_dict() for e in self.short_term.entries],
            },
        }

    def reset(self, layer: Optional[str] = None) -> None:
        """Reset memory layer(s).

        Args:
            layer: Which layer to reset ("core", "long_term", "short_term", or None for all)
        """
        if layer is None or layer == "core":
            self.core.entries.clear()
            self.core._save()

        if layer is None or layer == "long_term":
            self.long_term.entries.clear()
            self.long_term._save()

        if layer is None or layer == "short_term":
            self.short_term.clear()

    def get_stats(self) -> dict:
        """Get memory statistics."""
        return {
            "core_identity": {
                "entry_count": len(self.core.entries),
                "max_entries": self.core.MAX_ENTRIES,
            },
            "long_term": {
                "entry_count": len(self.long_term.entries),
                "max_entries": self.long_term.MAX_ENTRIES,
            },
            "short_term": {
                "entry_count": len(self.short_term.entries),
                "max_entries": self.short_term.MAX_ENTRIES,
            },
        }