"""
Contextual Memory Engine for Toxo.

Implements context-aware memory storage and retrieval.
"""

import asyncio
import json
import numpy as np
from typing import Dict, Any, List, Optional
from dataclasses import dataclass
from datetime import datetime

from ..utils.logger import get_logger


@dataclass
class ContextualMemory:
    """A contextual memory item."""
    id: str
    content: str
    context: Dict[str, Any]
    embedding: Optional[np.ndarray] = None
    timestamp: datetime = None
    
    def __post_init__(self):
        if self.timestamp is None:
            self.timestamp = datetime.now()


class ContextualMemorySystem:
    """Contextual memory system for context-aware storage."""
    
    def __init__(self):
        self.logger = get_logger("toxo.memory.contextual")
        self.memories: Dict[str, ContextualMemory] = {}
        self.context_index: Dict[str, List[str]] = {}
        self.logger.info("Contextual memory system initialized")
    
    async def store(self, data: Dict[str, Any]) -> str:
        """Store contextual memory."""
        content = data.get("content", "")
        context = data.get("context", {})
        memory_id = f"contextual_{len(self.memories)}_{datetime.now().timestamp()}"
        
        memory = ContextualMemory(
            id=memory_id,
            content=content,
            context=context
        )
        
        self.memories[memory_id] = memory
        
        # Index by context keys
        for key in context.keys():
            if key not in self.context_index:
                self.context_index[key] = []
            self.context_index[key].append(memory_id)
        
        self.logger.debug(f"Stored contextual memory: {memory_id}")
        return memory_id
    
    async def retrieve(self, query: str) -> List[Dict[str, Any]]:
        """Retrieve contextual memories."""
        results = []
        for memory in self.memories.values():
            if query.lower() in memory.content.lower():
                results.append({
                    "id": memory.id,
                    "content": memory.content,
                    "context": memory.context,
                    "timestamp": memory.timestamp.isoformat()
                })
        return results
    
    async def search(self, query: str) -> List[Dict[str, Any]]:
        """Search contextual memories."""
        return await self.retrieve(query)
    
    async def analyze(self) -> Dict[str, Any]:
        """Analyze contextual memory."""
        return {
            "total_memories": len(self.memories),
            "context_keys": list(self.context_index.keys()),
            "context_diversity": len(self.context_index)
        }
    
    async def optimize(self) -> Dict[str, Any]:
        """Optimize contextual memory."""
        self.logger.info("Optimizing contextual memory")
        return {"status": "optimized", "memories_count": len(self.memories)}
    
    async def consolidate(self) -> Dict[str, Any]:
        """Consolidate contextual memory."""
        self.logger.info("Consolidating contextual memory")
        return {"status": "consolidated", "memories_count": len(self.memories)}
    
    def get_stats(self) -> Dict[str, Any]:
        """Get memory statistics."""
        return {
            "type": "contextual",
            "memories": len(self.memories),
            "context_keys": len(self.context_index)
        } 