"""
Long-term Memory System for Toxo.

Implements persistent long-term memory for knowledge consolidation.
"""

import asyncio
import json
from typing import Dict, Any, List, Optional
from dataclasses import dataclass, asdict
from datetime import datetime
from pathlib import Path

from ..utils.logger import get_logger


@dataclass
class LongTermMemoryItem:
    """A long-term memory item."""
    id: str
    content: str
    knowledge_type: str
    importance: float = 0.5
    consolidation_score: float = 0.0
    access_count: int = 0
    created_at: datetime = None
    last_accessed: Optional[datetime] = None
    
    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.now()
        elif isinstance(self.created_at, str):
            self.created_at = datetime.fromisoformat(self.created_at)
        if isinstance(self.last_accessed, str):
            self.last_accessed = datetime.fromisoformat(self.last_accessed)


class LongTermMemorySystem:
    """Long-term memory system for persistent knowledge storage."""
    
    def __init__(self, storage_path: Optional[Path] = None):
        self.logger = get_logger("toxo.memory.long_term")
        self.storage_path = storage_path or Path("data/long_term_memory.json")
        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
        
        self.memories: Dict[str, LongTermMemoryItem] = {}
        self.knowledge_index: Dict[str, List[str]] = {}
        
        # Load existing memories
        self._load_memories()
        
        self.logger.info(f"Long-term memory system initialized with {len(self.memories)} items")
    
    async def store(self, data: Dict[str, Any]) -> str:
        """Store item in long-term memory."""
        content = data.get("content", "")
        knowledge_type = data.get("knowledge_type", "general")
        importance = data.get("importance", 0.5)
        
        memory_id = f"ltm_{len(self.memories)}_{datetime.now().timestamp()}"
        
        memory = LongTermMemoryItem(
            id=memory_id,
            content=content,
            knowledge_type=knowledge_type,
            importance=importance
        )
        
        self.memories[memory_id] = memory
        
        # Index by knowledge type
        if knowledge_type not in self.knowledge_index:
            self.knowledge_index[knowledge_type] = []
        self.knowledge_index[knowledge_type].append(memory_id)
        
        # Save to disk
        await self._save_memories()
        
        self.logger.debug(f"Stored long-term memory: {memory_id}")
        return memory_id
    
    async def retrieve(self, query: str) -> List[Dict[str, Any]]:
        """Retrieve items from long-term memory."""
        results = []
        for memory in self.memories.values():
            if query.lower() in memory.content.lower():
                # Update access tracking
                memory.access_count += 1
                memory.last_accessed = datetime.now()
                
                results.append({
                    "id": memory.id,
                    "content": memory.content,
                    "knowledge_type": memory.knowledge_type,
                    "importance": memory.importance,
                    "consolidation_score": memory.consolidation_score,
                    "access_count": memory.access_count,
                    "created_at": memory.created_at.isoformat(),
                    "last_accessed": memory.last_accessed.isoformat() if memory.last_accessed else None
                })
        
        # Sort by importance and consolidation score
        results.sort(key=lambda x: (x["importance"], x["consolidation_score"]), reverse=True)
        
        # Save updated access tracking
        await self._save_memories()
        
        return results
    
    async def search(self, query: str) -> List[Dict[str, Any]]:
        """Search long-term memory."""
        return await self.retrieve(query)
    
    async def analyze(self) -> Dict[str, Any]:
        """Analyze long-term memory."""
        return {
            "total_memories": len(self.memories),
            "knowledge_types": list(self.knowledge_index.keys()),
            "avg_importance": sum(m.importance for m in self.memories.values()) / max(1, len(self.memories)),
            "avg_consolidation": sum(m.consolidation_score for m in self.memories.values()) / max(1, len(self.memories)),
            "total_accesses": sum(m.access_count for m in self.memories.values())
        }
    
    async def optimize(self) -> Dict[str, Any]:
        """Optimize long-term memory."""
        # Update consolidation scores based on access patterns
        for memory in self.memories.values():
            # Simple consolidation scoring based on access frequency and recency
            if memory.access_count > 0:
                days_since_creation = (datetime.now() - memory.created_at).days + 1
                memory.consolidation_score = memory.access_count / days_since_creation
        
        await self._save_memories()
        self.logger.info("Optimized long-term memory")
        return {"status": "optimized", "memories_count": len(self.memories)}
    
    async def consolidate(self) -> Dict[str, Any]:
        """Consolidate long-term memory."""
        # Remove low-importance, rarely accessed memories
        threshold_importance = 0.2
        threshold_accesses = 1
        
        to_remove = []
        for memory_id, memory in self.memories.items():
            if memory.importance < threshold_importance and memory.access_count < threshold_accesses:
                to_remove.append(memory_id)
        
        for memory_id in to_remove:
            memory = self.memories[memory_id]
            del self.memories[memory_id]
            # Remove from index
            if memory.knowledge_type in self.knowledge_index:
                self.knowledge_index[memory.knowledge_type].remove(memory_id)
        
        await self._save_memories()
        self.logger.info(f"Consolidated long-term memory, removed {len(to_remove)} items")
        return {"status": "consolidated", "removed_count": len(to_remove), "remaining_count": len(self.memories)}
    
    def _load_memories(self):
        """Load memories from disk."""
        if self.storage_path.exists():
            try:
                with open(self.storage_path, 'r') as f:
                    data = json.load(f)
                
                for memory_data in data.get("memories", []):
                    memory = LongTermMemoryItem(**memory_data)
                    self.memories[memory.id] = memory
                    
                    # Rebuild index
                    if memory.knowledge_type not in self.knowledge_index:
                        self.knowledge_index[memory.knowledge_type] = []
                    self.knowledge_index[memory.knowledge_type].append(memory.id)
                
                self.logger.info(f"Loaded {len(self.memories)} long-term memories from disk")
            except Exception as e:
                self.logger.error(f"Failed to load memories: {e}")
    
    async def _save_memories(self):
        """Save memories to disk."""
        try:
            data = {
                "memories": [asdict(memory) for memory in self.memories.values()],
                "saved_at": datetime.now().isoformat()
            }
            
            # Convert datetime objects to strings for JSON serialization
            for memory_data in data["memories"]:
                if isinstance(memory_data.get("created_at"), datetime):
                    memory_data["created_at"] = memory_data["created_at"].isoformat()
                if isinstance(memory_data.get("last_accessed"), datetime):
                    memory_data["last_accessed"] = memory_data["last_accessed"].isoformat()
            
            with open(self.storage_path, 'w') as f:
                json.dump(data, f, indent=2)
                
        except Exception as e:
            self.logger.error(f"Failed to save memories: {e}")
    
    def get_stats(self) -> Dict[str, Any]:
        """Get memory statistics."""
        return {
            "type": "long_term",
            "memories": len(self.memories),
            "knowledge_types": len(self.knowledge_index),
            "storage_path": str(self.storage_path)
        } 