# createsonline/agents/memory.py
"""
Memory systems for CREATESONLINE agents.

Provides short-term conversation memory, long-term vector memory,
and a unified MemoryManager.
"""

import time
import json
import hashlib
import math
from typing import Any, Dict, List, Optional, Tuple


class Message:
    """A single message in a conversation."""
    
    __slots__ = ('role', 'content', 'timestamp', 'metadata', 'token_estimate')
    
    def __init__(self, role: str, content: str, metadata: dict = None):
        self.role = role            # 'system', 'user', 'assistant', 'tool'
        self.content = content
        self.timestamp = time.time()
        self.metadata = metadata or {}
        # Rough token estimate (1 token ≈ 4 chars)
        self.token_estimate = max(1, len(content) // 4)
    
    def to_dict(self) -> dict:
        return {
            'role': self.role,
            'content': self.content,
            'timestamp': self.timestamp,
            'metadata': self.metadata,
        }
    
    @classmethod
    def from_dict(cls, d: dict) -> 'Message':
        msg = cls(d['role'], d['content'], d.get('metadata'))
        msg.timestamp = d.get('timestamp', time.time())
        return msg
    
    def __repr__(self):
        preview = self.content[:40] + '...' if len(self.content) > 40 else self.content
        return f"Message({self.role!r}, {preview!r})"


class ConversationMemory:
    """
    Short-term conversation memory with sliding window.
    
    Features:
        - Configurable max turns and max tokens
        - Automatic summarization trigger when memory is full
        - System message always retained
        - Turn-based or token-based eviction
        - Full history preserved separately for persistence
    
    Usage::
    
        memory = ConversationMemory(max_turns=50, max_tokens=4000)
        memory.add_system("You are a helpful assistant.")
        memory.add_user("Hello!")
        memory.add_assistant("Hi! How can I help?")
        
        # Get messages for LLM context
        messages = memory.get_messages()
    """
    
    def __init__(self, max_turns: int = 50, max_tokens: int = 8000,
                 summary_threshold: int = None):
        self.max_turns = max_turns
        self.max_tokens = max_tokens
        self.summary_threshold = summary_threshold or (max_turns * 2 // 3)
        
        self._system_message: Optional[Message] = None
        self._messages: List[Message] = []
        self._full_history: List[Message] = []    # never evicted
        self._summary: str = ''                    # rolling summary of evicted turns
        self._total_messages = 0
    
    # ------------------------------------------------------------------
    # Adding messages
    # ------------------------------------------------------------------
    def add_system(self, content: str, metadata: dict = None):
        """Set or update the system message."""
        self._system_message = Message('system', content, metadata)
    
    def add_user(self, content: str, metadata: dict = None):
        self._add('user', content, metadata)
    
    def add_assistant(self, content: str, metadata: dict = None):
        self._add('assistant', content, metadata)
    
    def add_tool(self, content: str, tool_name: str = '', metadata: dict = None):
        meta = metadata or {}
        meta['tool_name'] = tool_name
        self._add('tool', content, meta)
    
    def add(self, role: str, content: str, metadata: dict = None):
        if role == 'system':
            self.add_system(content, metadata)
        else:
            self._add(role, content, metadata)
    
    def _add(self, role: str, content: str, metadata: dict = None):
        msg = Message(role, content, metadata)
        self._messages.append(msg)
        self._full_history.append(msg)
        self._total_messages += 1
        self._evict_if_needed()
    
    # ------------------------------------------------------------------
    # Eviction
    # ------------------------------------------------------------------
    def _evict_if_needed(self):
        """Remove oldest messages if we exceed limits."""
        # Turn-based eviction
        while len(self._messages) > self.max_turns:
            evicted = self._messages.pop(0)
            self._update_summary(evicted)
        
        # Token-based eviction
        while self._token_count() > self.max_tokens and len(self._messages) > 1:
            evicted = self._messages.pop(0)
            self._update_summary(evicted)
    
    def _token_count(self) -> int:
        total = 0
        if self._system_message:
            total += self._system_message.token_estimate
        if self._summary:
            total += len(self._summary) // 4
        for m in self._messages:
            total += m.token_estimate
        return total
    
    def _update_summary(self, evicted: Message):
        """Append evicted message info to rolling summary."""
        snippet = evicted.content[:200]
        self._summary += f"\n[{evicted.role}]: {snippet}"
        # Keep summary bounded
        if len(self._summary) > 2000:
            self._summary = self._summary[-1500:]
    
    # ------------------------------------------------------------------
    # Retrieval
    # ------------------------------------------------------------------
    def get_messages(self, include_summary: bool = True) -> List[dict]:
        """
        Get the message list formatted for LLM consumption.
        
        Returns list of dicts with 'role' and 'content' keys.
        System message is always first, followed by optional summary,
        then recent conversation history.
        """
        result = []
        
        if self._system_message:
            sys_content = self._system_message.content
            if include_summary and self._summary:
                sys_content += (
                    "\n\n[Previous conversation summary]:\n" + self._summary.strip()
                )
            result.append({'role': 'system', 'content': sys_content})
        elif include_summary and self._summary:
            result.append({
                'role': 'system',
                'content': "[Previous conversation summary]:\n" + self._summary.strip()
            })
        
        for m in self._messages:
            result.append({'role': m.role, 'content': m.content})
        
        return result
    
    def get_last(self, n: int = 1) -> List[Message]:
        return self._messages[-n:]
    
    def get_full_history(self) -> List[Message]:
        return list(self._full_history)
    
    @property
    def turn_count(self) -> int:
        return len(self._messages)
    
    @property
    def total_messages(self) -> int:
        return self._total_messages
    
    # ------------------------------------------------------------------
    # Search within memory
    # ------------------------------------------------------------------
    def search(self, query: str, top_k: int = 5) -> List[Message]:
        """Simple keyword search in conversation history."""
        query_lower = query.lower()
        scored = []
        for msg in self._full_history:
            content_lower = msg.content.lower()
            # Simple relevance: count query word occurrences
            score = sum(1 for w in query_lower.split() if w in content_lower)
            if score > 0:
                scored.append((score, msg))
        scored.sort(key=lambda x: x[0], reverse=True)
        return [m for _, m in scored[:top_k]]
    
    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------
    def save(self, filepath: str):
        """Save conversation to JSON file."""
        data = {
            'system': self._system_message.to_dict() if self._system_message else None,
            'messages': [m.to_dict() for m in self._full_history],
            'summary': self._summary,
            'max_turns': self.max_turns,
            'max_tokens': self.max_tokens,
        }
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, default=str)
    
    def load(self, filepath: str):
        """Load conversation from JSON file."""
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        self.max_turns = data.get('max_turns', self.max_turns)
        self.max_tokens = data.get('max_tokens', self.max_tokens)
        self._summary = data.get('summary', '')
        
        if data.get('system'):
            self._system_message = Message.from_dict(data['system'])
        
        self._full_history = [Message.from_dict(m) for m in data.get('messages', [])]
        # Restore active window (last max_turns messages)
        self._messages = self._full_history[-self.max_turns:]
        self._total_messages = len(self._full_history)
    
    def clear(self):
        """Clear all messages but keep system prompt."""
        self._messages.clear()
        self._full_history.clear()
        self._summary = ''
        self._total_messages = 0
    
    def __len__(self):
        return len(self._messages)
    
    def __repr__(self):
        return f"ConversationMemory(turns={len(self._messages)}/{self.max_turns})"


class VectorMemory:
    """
    Long-term vector memory using numpy-based embeddings.
    
    Stores text chunks with their embeddings for semantic retrieval.
    Uses a simple bag-of-words or TF-IDF approach for embeddings
    (no external dependencies needed).
    
    Usage::
    
        vmem = VectorMemory(embedding_dim=128)
        vmem.add("Python is a programming language", metadata={"source": "wiki"})
        vmem.add("JavaScript runs in browsers")
        
        results = vmem.search("programming", top_k=3)
    """
    
    def __init__(self, embedding_dim: int = 256):
        self.embedding_dim = embedding_dim
        self._entries: List[Dict[str, Any]] = []   # {text, embedding, metadata, timestamp}
        self._vocab: Dict[str, int] = {}            # word -> index (built incrementally)
        self._idf: Dict[str, float] = {}            # word -> IDF score
        self._doc_count = 0
    
    # ------------------------------------------------------------------
    # Embedding (built-in TF-IDF hashing)
    # ------------------------------------------------------------------
    def _tokenize(self, text: str) -> List[str]:
        """Simple whitespace + punctuation tokenizer."""
        import re
        return re.findall(r'\b\w+\b', text.lower())
    
    def _hash_embed(self, text: str) -> 'list':
        """
        Create a fixed-dimension embedding using feature hashing (hashing trick).
        This is dependency-free and works without numpy at the cost of quality.
        """
        tokens = self._tokenize(text)
        if not tokens:
            return [0.0] * self.embedding_dim
        
        vec = [0.0] * self.embedding_dim
        
        # Count term frequencies
        tf: Dict[str, int] = {}
        for t in tokens:
            tf[t] = tf.get(t, 0) + 1
        
        # Hash each token to a dimension
        for token, count in tf.items():
            h = int(hashlib.md5(token.encode()).hexdigest(), 16)
            idx = h % self.embedding_dim
            sign = 1 if (h // self.embedding_dim) % 2 == 0 else -1
            # Weight by TF and IDF
            idf = self._idf.get(token, 1.0)
            vec[idx] += sign * count * idf
            
            # Also add bigram features
            for other in tf:
                if other > token:
                    bigram = f"{token}_{other}"
                    bh = int(hashlib.md5(bigram.encode()).hexdigest(), 16)
                    bidx = bh % self.embedding_dim
                    bsign = 1 if (bh // self.embedding_dim) % 2 == 0 else -1
                    vec[bidx] += bsign * 0.5
        
        # L2 normalize
        norm = math.sqrt(sum(x * x for x in vec))
        if norm > 0:
            vec = [x / norm for x in vec]
        
        return vec
    
    def _update_idf(self):
        """Update IDF scores based on current documents."""
        doc_freq: Dict[str, int] = {}
        for entry in self._entries:
            seen = set(self._tokenize(entry['text']))
            for w in seen:
                doc_freq[w] = doc_freq.get(w, 0) + 1
        
        n = max(len(self._entries), 1)
        self._idf = {
            w: math.log(n / (1 + df)) for w, df in doc_freq.items()
        }
    
    @staticmethod
    def _cosine_similarity(a: list, b: list) -> float:
        """Cosine similarity between two vectors."""
        dot = sum(x * y for x, y in zip(a, b))
        na = math.sqrt(sum(x * x for x in a))
        nb = math.sqrt(sum(x * x for x in b))
        if na == 0 or nb == 0:
            return 0.0
        return dot / (na * nb)
    
    # ------------------------------------------------------------------
    # CRUD
    # ------------------------------------------------------------------
    def add(self, text: str, metadata: dict = None, embedding: list = None) -> int:
        """
        Add a text to vector memory.
        
        Args:
            text: The text content to store.
            metadata: Optional metadata dict.
            embedding: Pre-computed embedding (if None, auto-computed).
        
        Returns:
            Index of the added entry.
        """
        self._doc_count += 1
        
        # Periodically update IDF (every 10 docs)
        if self._doc_count % 10 == 0:
            self._update_idf()
        
        if embedding is None:
            embedding = self._hash_embed(text)
        
        entry = {
            'text': text,
            'embedding': embedding,
            'metadata': metadata or {},
            'timestamp': time.time(),
            'id': len(self._entries),
        }
        self._entries.append(entry)
        return entry['id']
    
    def add_many(self, texts: List[str], metadatas: List[dict] = None) -> List[int]:
        """Add multiple texts at once."""
        if metadatas is None:
            metadatas = [{}] * len(texts)
        ids = []
        for text, meta in zip(texts, metadatas):
            ids.append(self.add(text, meta))
        # Final IDF update
        self._update_idf()
        return ids
    
    def search(self, query: str, top_k: int = 5,
               min_score: float = 0.0,
               filter_metadata: dict = None) -> List[Dict[str, Any]]:
        """
        Semantic search over stored memories.
        
        Args:
            query: Query text.
            top_k: Maximum results to return.
            min_score: Minimum cosine similarity threshold.
            filter_metadata: Only return entries matching these metadata key-values.
        
        Returns:
            List of dicts with 'text', 'score', 'metadata', 'id'.
        """
        if not self._entries:
            return []
        
        q_emb = self._hash_embed(query)
        
        scored = []
        for entry in self._entries:
            # Metadata filter
            if filter_metadata:
                match = all(
                    entry['metadata'].get(k) == v
                    for k, v in filter_metadata.items()
                )
                if not match:
                    continue
            
            score = self._cosine_similarity(q_emb, entry['embedding'])
            if score >= min_score:
                scored.append({
                    'text': entry['text'],
                    'score': score,
                    'metadata': entry['metadata'],
                    'id': entry['id'],
                })
        
        scored.sort(key=lambda x: x['score'], reverse=True)
        return scored[:top_k]
    
    def get(self, entry_id: int) -> Optional[dict]:
        """Get entry by ID."""
        if 0 <= entry_id < len(self._entries):
            return self._entries[entry_id]
        return None
    
    def delete(self, entry_id: int) -> bool:
        """Soft-delete an entry by clearing its content."""
        if 0 <= entry_id < len(self._entries):
            self._entries[entry_id]['text'] = ''
            self._entries[entry_id]['embedding'] = [0.0] * self.embedding_dim
            self._entries[entry_id]['metadata']['_deleted'] = True
            return True
        return False
    
    @property
    def size(self) -> int:
        return sum(1 for e in self._entries if not e.get('metadata', {}).get('_deleted'))
    
    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------
    def save(self, filepath: str):
        """Save vector memory to JSON file."""
        data = {
            'embedding_dim': self.embedding_dim,
            'entries': [
                {
                    'text': e['text'],
                    'metadata': e['metadata'],
                    'timestamp': e['timestamp'],
                    'id': e['id'],
                    # Embeddings are large — store compactly
                    'embedding': [round(x, 6) for x in e['embedding']],
                }
                for e in self._entries
            ],
            'idf': {k: round(v, 6) for k, v in self._idf.items()},
        }
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f)
    
    def load(self, filepath: str):
        """Load vector memory from JSON file."""
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        self.embedding_dim = data.get('embedding_dim', self.embedding_dim)
        self._idf = data.get('idf', {})
        self._entries = []
        for e in data.get('entries', []):
            self._entries.append({
                'text': e['text'],
                'metadata': e.get('metadata', {}),
                'timestamp': e.get('timestamp', time.time()),
                'id': e.get('id', len(self._entries)),
                'embedding': e['embedding'],
            })
        self._doc_count = len(self._entries)
    
    def clear(self):
        self._entries.clear()
        self._idf.clear()
        self._doc_count = 0
    
    def __len__(self):
        return self.size
    
    def __repr__(self):
        return f"VectorMemory(entries={self.size}, dim={self.embedding_dim})"


class MemoryManager:
    """
    Unified memory manager combining conversation and vector memory.
    
    Automatically indexes important messages into long-term vector memory
    for later retrieval (RAG pattern).
    
    Usage::
    
        manager = MemoryManager()
        manager.add_message("user", "Tell me about Python asyncio")
        manager.add_message("assistant", "asyncio is Python's built-in async framework...")
        
        # Later: retrieve relevant past context
        context = manager.retrieve_relevant("How does asyncio work?")
    """
    
    def __init__(self, max_turns: int = 50, max_tokens: int = 8000,
                 embedding_dim: int = 256, auto_index: bool = True,
                 index_min_length: int = 50):
        self.conversation = ConversationMemory(max_turns, max_tokens)
        self.vector = VectorMemory(embedding_dim)
        self.auto_index = auto_index
        self.index_min_length = index_min_length
        
        # Named memory stores for structured data
        self._stores: Dict[str, Dict[str, Any]] = {}
    
    # ------------------------------------------------------------------
    # Conversation
    # ------------------------------------------------------------------
    def add_message(self, role: str, content: str, metadata: dict = None):
        """Add a message and optionally index it for long-term retrieval."""
        self.conversation.add(role, content, metadata)
        
        # Auto-index substantial messages into vector memory
        if self.auto_index and len(content) >= self.index_min_length:
            self.vector.add(content, metadata={
                'role': role,
                'timestamp': time.time(),
                **(metadata or {}),
            })
    
    def get_messages(self) -> List[dict]:
        return self.conversation.get_messages()
    
    # ------------------------------------------------------------------
    # Retrieval-Augmented Generation (RAG)
    # ------------------------------------------------------------------
    def retrieve_relevant(self, query: str, top_k: int = 5,
                          include_conversation: bool = True) -> List[dict]:
        """
        Retrieve relevant context from both conversation and vector memory.
        
        This implements the RAG pattern: given a query, find the most relevant
        past messages and stored knowledge.
        """
        results = []
        
        # Search vector memory
        vec_results = self.vector.search(query, top_k=top_k)
        for vr in vec_results:
            results.append({
                'source': 'vector_memory',
                'text': vr['text'],
                'score': vr['score'],
                'metadata': vr['metadata'],
            })
        
        # Also search conversation history
        if include_conversation:
            conv_results = self.conversation.search(query, top_k=top_k)
            for msg in conv_results:
                results.append({
                    'source': 'conversation',
                    'text': msg.content,
                    'score': 0.5,  # Default score for keyword matches
                    'metadata': {'role': msg.role, **msg.metadata},
                })
        
        # Deduplicate by content hash
        seen = set()
        deduped = []
        for r in results:
            h = hashlib.md5(r['text'].encode()).hexdigest()
            if h not in seen:
                seen.add(h)
                deduped.append(r)
        
        # Sort by score
        deduped.sort(key=lambda x: x['score'], reverse=True)
        return deduped[:top_k]
    
    def build_rag_context(self, query: str, top_k: int = 3,
                          max_context_chars: int = 3000) -> str:
        """
        Build a RAG context string to inject into the prompt.
        
        Returns a formatted string with the most relevant past knowledge.
        """
        results = self.retrieve_relevant(query, top_k=top_k)
        if not results:
            return ''
        
        lines = ["[Relevant context from memory]:"]
        total = 0
        for r in results:
            text = r['text']
            if total + len(text) > max_context_chars:
                text = text[:max_context_chars - total]
            lines.append(f"- {text}")
            total += len(text)
            if total >= max_context_chars:
                break
        
        return '\n'.join(lines)
    
    # ------------------------------------------------------------------
    # Named stores (key-value memory)
    # ------------------------------------------------------------------
    def set(self, store: str, key: str, value: Any):
        """Store a key-value pair in a named store."""
        if store not in self._stores:
            self._stores[store] = {}
        self._stores[store][key] = value
    
    def get(self, store: str, key: str, default: Any = None) -> Any:
        """Get a value from a named store."""
        return self._stores.get(store, {}).get(key, default)
    
    def get_store(self, store: str) -> Dict[str, Any]:
        """Get all key-value pairs in a named store."""
        return dict(self._stores.get(store, {}))
    
    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------
    def save(self, directory: str):
        """Save all memory to a directory."""
        import os
        os.makedirs(directory, exist_ok=True)
        self.conversation.save(os.path.join(directory, 'conversation.json'))
        self.vector.save(os.path.join(directory, 'vectors.json'))
        
        # Save named stores
        stores_path = os.path.join(directory, 'stores.json')
        with open(stores_path, 'w', encoding='utf-8') as f:
            json.dump(self._stores, f, indent=2, default=str)
    
    def load(self, directory: str):
        """Load all memory from a directory."""
        import os
        conv_path = os.path.join(directory, 'conversation.json')
        if os.path.exists(conv_path):
            self.conversation.load(conv_path)
        
        vec_path = os.path.join(directory, 'vectors.json')
        if os.path.exists(vec_path):
            self.vector.load(vec_path)
        
        stores_path = os.path.join(directory, 'stores.json')
        if os.path.exists(stores_path):
            with open(stores_path, 'r', encoding='utf-8') as f:
                self._stores = json.load(f)
    
    def clear(self):
        """Clear all memory."""
        self.conversation.clear()
        self.vector.clear()
        self._stores.clear()
    
    def __repr__(self):
        return (
            f"MemoryManager(conversation={self.conversation.turn_count} turns, "
            f"vector={self.vector.size} entries, "
            f"stores={len(self._stores)})"
        )
