"""
Memory system for Toxo layers.

This module implements hierarchical memory management for storing and retrieving
context, feedback, and domain knowledge across training and inference sessions.
"""

import json
import sqlite3
import asyncio
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, asdict
from datetime import datetime
import numpy as np
from sentence_transformers import SentenceTransformer

from .config import MemoryConfig
from ..utils.logger import get_logger
from ..utils.exceptions import MemoryError


@dataclass
class MemoryItem:
    """Individual memory item with content and metadata."""
    id: str
    content: str
    embedding: Optional[np.ndarray]
    memory_type: str  # 'interaction', 'feedback', 'knowledge', 'context'
    metadata: Dict[str, Any]
    created_at: str
    accessed_count: int = 0
    last_accessed: Optional[str] = None


class ToxoMemory:
    """
    Hierarchical memory system for Toxo layers.
    
    Manages multiple layers of memory:
    - Working memory: Current session context
    - Session memory: Recent interactions within current session
    - Domain memory: Long-term domain-specific knowledge
    - Global memory: Cross-domain patterns and knowledge
    """
    
    def __init__(self, config: MemoryConfig):
        """
        Initialize Toxo memory system.
        
        Args:
            config: Memory configuration
        """
        self.config = config
        self.logger = get_logger("toxo.memory")
        
        # Memory storage
        self.working_memory: List[MemoryItem] = []
        self.session_memory: List[MemoryItem] = []
        self.domain_memory: List[MemoryItem] = []
        
        # Embedding model for semantic search
        self.embedding_model = None
        self._initialize_embedding_model()
        
        # SQLite database for persistent storage
        self.db_path: Optional[Path] = None
        self.connection: Optional[sqlite3.Connection] = None
        
        # Memory management
        self.interaction_count = 0
        
        self.logger.info("ToxoMemory initialized")
    
    def _initialize_embedding_model(self):
        """Initialize the embedding model for semantic search."""
        try:
            # Use a lightweight sentence transformer model
            self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
            self.logger.info("Embedding model loaded successfully")
        except Exception as e:
            self.logger.warning(f"Failed to load embedding model: {str(e)}")
            self.embedding_model = None
    
    def _get_embedding(self, text: str) -> Optional[np.ndarray]:
        """Get embedding for text."""
        if not self.embedding_model:
            return None
        
        try:
            embedding = self.embedding_model.encode(text)
            return embedding
        except Exception as e:
            self.logger.warning(f"Failed to generate embedding: {str(e)}")
            return None
    
    async def store_interaction(
        self,
        input_text: str,
        response: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Store an interaction in memory.
        
        Args:
            input_text: User input
            response: System response
            metadata: Additional metadata
            
        Returns:
            Memory item ID
        """
        interaction_data = {
            "input": input_text,
            "response": response,
            "timestamp": datetime.now().isoformat()
        }
        if metadata:
            interaction_data.update(metadata)
        
        content = f"Input: {input_text}\nResponse: {response}"
        
        memory_item = MemoryItem(
            id=f"interaction_{self.interaction_count}_{datetime.now().timestamp()}",
            content=content,
            embedding=self._get_embedding(content),
            memory_type="interaction",
            metadata=interaction_data,
            created_at=datetime.now().isoformat()
        )
        
        # Add to appropriate memory layers
        self.working_memory.append(memory_item)
        self.session_memory.append(memory_item)
        
        # Manage memory size limits
        await self._manage_memory_limits()
        
        self.interaction_count += 1
        self.logger.debug(f"Stored interaction: {memory_item.id}")
        
        return memory_item.id
    
    async def store_feedback(
        self,
        input_text: str,
        expected_output: str,
        actual_output: Optional[str] = None,
        feedback_type: str = "correction",
        metadata: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Store feedback for learning.
        
        Args:
            input_text: Original input
            expected_output: Desired output
            actual_output: What was actually produced
            feedback_type: Type of feedback
            metadata: Additional metadata
            
        Returns:
            Memory item ID
        """
        feedback_data = {
            "input": input_text,
            "expected_output": expected_output,
            "actual_output": actual_output,
            "feedback_type": feedback_type,
            "timestamp": datetime.now().isoformat()
        }
        if metadata:
            feedback_data.update(metadata)
        
        content = f"Input: {input_text}\nExpected: {expected_output}"
        if actual_output:
            content += f"\nActual: {actual_output}"
        
        memory_item = MemoryItem(
            id=f"feedback_{datetime.now().timestamp()}",
            content=content,
            embedding=self._get_embedding(content),
            memory_type="feedback",
            metadata=feedback_data,
            created_at=datetime.now().isoformat()
        )
        
        # Feedback goes directly to domain memory for learning
        self.domain_memory.append(memory_item)
        
        await self._manage_memory_limits()
        
        self.logger.debug(f"Stored feedback: {memory_item.id}")
        return memory_item.id
    
    async def retrieve(
        self,
        query: str,
        memory_types: Optional[List[str]] = None,
        limit: int = 5,
        similarity_threshold: float = 0.7
    ) -> List[MemoryItem]:
        """
        Retrieve relevant memories based on query.
        
        Args:
            query: Search query
            memory_types: Types of memory to search
            limit: Maximum number of results
            similarity_threshold: Minimum similarity score
            
        Returns:
            List of relevant memory items
        """
        if not self.embedding_model:
            # Fallback to keyword search
            return await self._keyword_search(query, memory_types, limit)
        
        query_embedding = self._get_embedding(query)
        if query_embedding is None:
            return await self._keyword_search(query, memory_types, limit)
        
        # Collect all relevant memories
        all_memories = []
        if not memory_types or "interaction" in memory_types:
            all_memories.extend(self.working_memory)
            all_memories.extend(self.session_memory)
        if not memory_types or "feedback" in memory_types:
            all_memories.extend([m for m in self.domain_memory if m.memory_type == "feedback"])
        if not memory_types or "knowledge" in memory_types:
            all_memories.extend([m for m in self.domain_memory if m.memory_type == "knowledge"])
        
        # Calculate similarities
        scored_memories = []
        for memory in all_memories:
            if memory.embedding is not None:
                similarity = np.dot(query_embedding, memory.embedding) / (
                    np.linalg.norm(query_embedding) * np.linalg.norm(memory.embedding)
                )
                if similarity >= similarity_threshold:
                    scored_memories.append((memory, similarity))
        
        # Sort by similarity and return top results
        scored_memories.sort(key=lambda x: x[1], reverse=True)
        results = [memory for memory, _ in scored_memories[:limit]]
        
        # Update access tracking
        for memory in results:
            memory.accessed_count += 1
            memory.last_accessed = datetime.now().isoformat()
        
        self.logger.debug(f"Retrieved {len(results)} memories for query: {query[:50]}...")
        return results
    
    async def _keyword_search(
        self,
        query: str,
        memory_types: Optional[List[str]] = None,
        limit: int = 5
    ) -> List[MemoryItem]:
        """Fallback keyword-based search when embeddings are not available."""
        query_words = set(query.lower().split())
        
        all_memories = []
        if not memory_types or "interaction" in memory_types:
            all_memories.extend(self.working_memory + self.session_memory)
        if not memory_types or "feedback" in memory_types:
            all_memories.extend([m for m in self.domain_memory if m.memory_type == "feedback"])
        
        scored_memories = []
        for memory in all_memories:
            content_words = set(memory.content.lower().split())
            overlap = len(query_words.intersection(content_words))
            if overlap > 0:
                score = overlap / len(query_words)
                scored_memories.append((memory, score))
        
        scored_memories.sort(key=lambda x: x[1], reverse=True)
        results = [memory for memory, _ in scored_memories[:limit]]
        
        return results
    
    async def _manage_memory_limits(self):
        """Manage memory size limits and consolidation."""
        # Trim working memory
        if len(self.working_memory) > self.config.working_memory_size:
            excess = len(self.working_memory) - self.config.working_memory_size
            self.working_memory = self.working_memory[excess:]
        
        # Trim session memory
        if len(self.session_memory) > self.config.session_memory_size:
            excess = len(self.session_memory) - self.config.session_memory_size
            # Move oldest to domain memory before removing
            oldest = self.session_memory[:excess]
            for memory in oldest:
                if memory.memory_type == "interaction":
                    memory.memory_type = "knowledge"
                    self.domain_memory.append(memory)
            self.session_memory = self.session_memory[excess:]
        
        # Trim domain memory
        if len(self.domain_memory) > self.config.domain_memory_size:
            excess = len(self.domain_memory) - self.config.domain_memory_size
            # Remove least accessed items
            self.domain_memory.sort(key=lambda x: (x.accessed_count, x.created_at))
            self.domain_memory = self.domain_memory[excess:]
    
    def save(self, path: Path):
        """Save memory to disk."""
        try:
            save_data = {
                "working_memory": [asdict(item) for item in self.working_memory],
                "session_memory": [asdict(item) for item in self.session_memory],
                "domain_memory": [asdict(item) for item in self.domain_memory],
                "interaction_count": self.interaction_count,
                "config": asdict(self.config)
            }
            
            # Convert numpy arrays to lists for JSON serialization
            for memory_list in save_data.values():
                if isinstance(memory_list, list):
                    for item in memory_list:
                        if isinstance(item, dict) and "embedding" in item and item["embedding"] is not None:
                            item["embedding"] = item["embedding"].tolist()
            
            with open(path, 'w') as f:
                json.dump(save_data, f, indent=2)
            
            self.logger.info(f"Memory saved to {path}")
            
        except Exception as e:
            self.logger.error(f"Failed to save memory: {str(e)}")
            raise MemoryError(f"Save failed: {str(e)}")
    
    def load(self, path: Path):
        """Load memory from disk."""
        try:
            with open(path, 'r') as f:
                save_data = json.load(f)
            
            # Convert embedding lists back to numpy arrays
            for memory_list_key in ["working_memory", "session_memory", "domain_memory"]:
                for item_data in save_data.get(memory_list_key, []):
                    if item_data.get("embedding") is not None:
                        item_data["embedding"] = np.array(item_data["embedding"])
                    
                    # Reconstruct MemoryItem objects
                    memory_item = MemoryItem(**item_data)
                    getattr(self, memory_list_key).append(memory_item)
            
            self.interaction_count = save_data.get("interaction_count", 0)
            
            self.logger.info(f"Memory loaded from {path}")
            
        except Exception as e:
            self.logger.error(f"Failed to load memory: {str(e)}")
            raise MemoryError(f"Load failed: {str(e)}")
    
    def size(self) -> Dict[str, int]:
        """Get memory sizes."""
        return {
            "working_memory": len(self.working_memory),
            "session_memory": len(self.session_memory),
            "domain_memory": len(self.domain_memory),
            "total": len(self.working_memory) + len(self.session_memory) + len(self.domain_memory)
        }
    
    def clear_working_memory(self):
        """Clear working memory."""
        self.working_memory.clear()
        self.logger.debug("Working memory cleared")
    
    def clear_session_memory(self):
        """Clear session memory."""
        self.session_memory.clear()
        self.logger.debug("Session memory cleared")
    
    def get_memory_stats(self) -> Dict[str, Any]:
        """Get memory statistics."""
        sizes = self.size()
        
        # Calculate access patterns
        all_memories = self.working_memory + self.session_memory + self.domain_memory
        total_accesses = sum(m.accessed_count for m in all_memories)
        
        return {
            "sizes": sizes,
            "total_interactions": self.interaction_count,
            "total_accesses": total_accesses,
            "embedding_model_loaded": self.embedding_model is not None,
            "avg_accesses_per_memory": total_accesses / max(sizes["total"], 1)
        } 