import asyncio
import time
from typing import Dict, List, Any, Optional
import logging
from dataclasses import dataclass
from enum import Enum
import heapq
from collections import OrderedDict

class MemoryLevel(Enum):
    """Niveaux de la mémoire hiérarchique"""
    IMMEDIATE = "immediate"      # Mémoire sensorielle (quelques secondes)
    WORKING = "working"          # Mémoire de travail (quelques minutes)  
    SHORT_TERM = "short_term"    # Mémoire à court terme (heures)
    LONG_TERM = "long_term"      # Mémoire à long terme (jours/années)
    PERMANENT = "permanent"      # Mémoire permanente (cristallisée)

@dataclass
class MemoryItem:
    """Élément de mémoire avec métadonnées"""
    id: str
    content: Any
    level: MemoryLevel
    strength: float  # 0.0 à 1.0
    access_count: int
    last_accessed: float
    created_at: float
    associations: List[str]
    emotional_weight: float

class HierarchicalMemory:
    """
    Système de mémoire hiérarchique avancé
    avec mécanismes de consolidation et d'oubli
    """
    
    def __init__(self):
        self.logger = logging.getLogger("hierarchical_memory")
        self.memory_levels = {
            MemoryLevel.IMMEDIATE: OrderedDict(),
            MemoryLevel.WORKING: OrderedDict(),
            MemoryLevel.SHORT_TERM: OrderedDict(),
            MemoryLevel.LONG_TERM: OrderedDict(),
            MemoryLevel.PERMANENT: OrderedDict()
        }
        self.access_patterns = {}
        self.consolidation_threshold = 0.7
        self.forgetting_curve = {}
        
    async def initialize(self):
        """Initialise le système de mémoire hiérarchique"""
        self.logger.info("🏗️ Initialisation de la mémoire hiérarchique...")
        
        try:
            await self._setup_memory_structures()
            await self._load_forgetting_curves()
            
            # Démarre le processus de consolidation en arrière-plan
            asyncio.create_task(self._consolidation_worker())
            
            self.logger.info("✅ Mémoire hiérarchique initialisée")
            return True
            
        except Exception as e:
            self.logger.error(f"❌ Erreur d'initialisation mémoire: {e}")
            return False
    
    async def store(self, content: Any, level: MemoryLevel, priority: int = 5) -> str:
        """Stocke un élément en mémoire avec niveau spécifié"""
        try:
            memory_id = self._generate_memory_id(content)
            
            item = MemoryItem(
                id=memory_id,
                content=content,
                level=level,
                strength=0.1,  # Force initiale faible
                access_count=0,
                last_accessed=time.time(),
                created_at=time.time(),
                associations=[],
                emotional_weight=0.0
            )
            
            # Stockage selon le niveau
            self.memory_levels[level][memory_id] = item
            
            # Application des mécanismes d'attention basés sur la priorité
            await self._apply_attention_mechanism(item, priority)
            
            self.logger.info(f"💾 Mémoire stockée [{level.value}]: {memory_id}")
            return memory_id
            
        except Exception as e:
            self.logger.error(f"Erreur stockage mémoire: {e}")
            raise
    
    async def retrieve(self, memory_id: str, strengthen: bool = True) -> Optional[Any]:
        """Récupère un élément de mémoire"""
        try:
            # Recherche dans tous les niveaux
            for level_name, level_dict in self.memory_levels.items():
                if memory_id in level_dict:
                    item = level_dict[memory_id]
                    
                    # Mise à jour des métadonnées d'accès
                    item.access_count += 1
                    item.last_accessed = time.time()
                    
                    if strengthen:
                        await self._strengthen_memory(item)
                    
                    # Promotion potentielle vers un niveau supérieur
                    await self._consider_promotion(item)
                    
                    self.logger.debug(f"📖 Mémoire récupérée: {memory_id} (niveau: {level_name.value})")
                    return item.content
            
            self.logger.warning(f"⚠️ Mémoire introuvable: {memory_id}")
            return None
            
        except Exception as e:
            self.logger.error(f"Erreur récupération mémoire: {e}")
            return None
    
    async def associative_recall(self, cue: Any, max_results: int = 5) -> List[MemoryItem]:
        """Rappel associatif basé sur des indices"""
        try:
            results = []
            cue_str = str(cue).lower()
            
            for level_name, level_dict in self.memory_levels.items():
                for item in level_dict.values():
                    content_str = str(item.content).lower()
                    
                    # Vérifie la similarité
                    similarity = await self._calculate_similarity(cue_str, content_str)
                    if similarity > 0.3:  # Seuil de similarité
                        results.append((similarity, item))
            
            # Trie par similarité
            results.sort(key=lambda x: x[0], reverse=True)
            
            # Retourne les meilleurs résultats
            return [item for _, item in results[:max_results]]
            
        except Exception as e:
            self.logger.error(f"Erreur rappel associatif: {e}")
            return []
    
    async def promote(self, memory_id: str) -> bool:
        """Promouvoir un élément de mémoire vers un niveau supérieur"""
        try:
            # Trouve l'élément dans les niveaux
            current_level = None
            item = None
            
            for level_name, level_dict in self.memory_levels.items():
                if memory_id in level_dict:
                    current_level = level_name
                    item = level_dict[memory_id]
                    break
            
            if not item or not current_level:
                return False
            
            # Détermine le niveau suivant
            next_level = await self._get_next_level(current_level)
            if not next_level:
                return False  # Déjà au niveau maximum
            
            # Supprime de l'ancien niveau et ajoute au nouveau
            del self.memory_levels[current_level][memory_id]
            self.memory_levels[next_level][memory_id] = item
            
            # Renforce la mémoire lors de la promotion
            await self._strengthen_memory(item, boost=0.2)
            
            self.logger.info(f"📈 Mémoire promue: {memory_id} ({current_level