memory.py

import os
import re
from dataclasses import dataclass, asdict
from typing import List, Dict, Optional


@dataclass
class MemoryEntry:
    category: str     
    type: str         
    title: str
    text: str
    turn_created: int
    data_list: Optional[List[str]] = None

MEMORY_SYNTHESIS_PROMPT = MEMORY_SYNTHESIS_PROMPT = """
[ROLE]
You are the Memory Architect. Your goal is to extract ONLY NEW and RELEVANT facts.

[INPUT]
- Location: {location}
- Action: {action}
- Result: {result}
- Known Memories: {existing}

[TASK]
Compare the 'Result' with 'Known Memories'. 
- If the result contains NO NEW information or state changes, return {{"should_remember": false}}.
- If something changed or a new fact appeared, create a memory.
- Use the "supersedes" list to name the titles of old memories that are now obsolete.

[CRITICAL: NO DUPLICATES]
Do NOT extract information already present in 'Known Memories'. 
Focus on state changes (e.g., door opened, item taken).

[OUTPUT FORMAT]
Return ONLY a JSON object. No markdown blocks, no explanations.
{{
    "should_remember": bool,
    "memories": [
        {{
            "title": "Unique Title", 
            "text": "The new fact",
            "category": "DANGER"|"MECHANIC"|"STATE"|"INFO",
            "type": "PERMANENT"|"EPHEMERAL"|"CORE"
        }}
    ],
    "supersedes": ["Old Title"]
}}
"""

class HierarchicalMemoryManager:
    def __init__(self, call_llm_func, filepath="Memories.md"):
        self.call_llm = call_llm_func
        self.filepath = filepath
        self.memories: Dict[str, List[MemoryEntry]] = {} # Key: Location Name
        self.load_from_md() # Chargement au démarrage

    def load_from_md(self):
        """Charge les mémoires depuis le fichier Markdown."""
        if not os.path.exists(self.filepath):
            return

        current_loc = None
        pattern = re.compile(r'- \[(.*?)\] \[(.*?)\] \*\*(.*?)\*\*: (.*)')
        
        with open(self.filepath, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if line.startswith("## Location:"):
                    current_loc = line.split(":", 1)[1].strip()
                    self.memories[current_loc] = []
                elif line.startswith("- [") and current_loc:
                    match = pattern.match(line)
                    if match:
                        mem_type, cat, title, text = match.groups()
                        self.memories[current_loc].append(MemoryEntry(
                            category=cat, type=mem_type, title=title, text=text, turn_created=0
                        ))

    def _is_redundant(self, location: str, new_text: str) -> bool:
        """Vérifie si le texte existe déjà (partiellement ou totalement) dans ce lieu."""
        if location not in self.memories:
            return False
        
        new_text_clean = new_text.lower().strip()
        for m in self.memories[location]:
            existing_text = m.text.lower().strip()
            # 1. Correspondance exacte
            if new_text_clean == existing_text:
                return True
            # 2. Inclusion (si la nouvelle info est déjà contenue dans une ancienne)
            if new_text_clean in existing_text:
                return True
        return False

    def _upsert_memory(self, location: str, new_mem: MemoryEntry):
        """Remplace par titre OU ignore si le contenu est redondant."""
        if location not in self.memories:
            self.memories[location] = []
        
        # Si le contenu est déjà là, on ne fait rien (évite les doublons de sens)
        if self._is_redundant(location, new_mem.text):
            return

        # Remplacement par titre (Update d'état)
        self.memories[location] = [
            m for m in self.memories[location] 
            if m.title != new_mem.title
        ]
        self.memories[location].append(new_mem)
    
    def update_inventory(self, items: List[str], step: int):
        """Met à jour la section spéciale Inventaire dans le Markdown."""
        loc = "GLOBAL_INVENTORY"
        text = ", ".join(items) if items else "Empty"
        self._upsert_memory(loc, MemoryEntry(
            category="PLAYER", type="EPHEMERAL", 
            title="Current Inventory", text=text, turn_created=step
        ))
        self.save_to_md()

    def update_local_state(self, location: str, obs: dict, step: int):
        """Met à jour la mémoire locale à partir d'une StructuredObservation."""
        if location not in self.memories:
            self.memories[location] = []

        if obs.get("takeable_objects"):
            objs_text = ", ".join(obs["takeable_objects"])
            self._upsert_memory(location, MemoryEntry(
                category="ITEMS", type="EPHEMERAL", 
                title="Visible Objects", text=objs_text, turn_created=step
            ))

        if obs.get("visible_exits"):
            exits_text = ", ".join(obs["visible_exits"])
            self._upsert_memory(location, MemoryEntry(
                category="MAP", type="CORE", 
                title="Available Exits", text=exits_text, turn_created=step
            ))
            
        # 3. Sauvegarde
        self.save_to_md()

    def save_to_md(self):
        """Sauvegarde les mémoires dans le fichier Markdown."""
        with open(self.filepath, 'w', encoding='utf-8') as f:
            f.write("# ZorkGPT Agent Memories\n\n")
            
            for loc, entries in sorted(self.memories.items()):
                if not entries: continue
                f.write(f"## Location: {loc}\n")
                
                # On trie pour afficher CORE -> PERMANENT -> EPHEMERAL
                entries.sort(key=lambda x: {"CORE": 0, "PERMANENT": 1, "EPHEMERAL": 2}.get(x.type, 3))
                
                for m in entries:
                    f.write(f"- [{m.type}] [{m.category}] **{m.title}**: {m.text}\n")
                f.write("\n")

    def get_context(self, location: str) -> str:
        """Récupère le contexte formaté pour le LLM."""
        context_lines = []

        if "GLOBAL_INVENTORY" in self.memories:
            inv = self.memories["GLOBAL_INVENTORY"][0] # On sait qu'il n'y a qu'une entrée
            context_lines.append(f"🎒 CURRENT INVENTORY: {inv.text}")
            
        if location in self.memories and self.memories[location]:
            context_lines.append(f"🧠 KNOWLEDGE OF {location.upper()}:")
            # On trie pour avoir CORE (le décor) puis le reste
            sorted_entries = sorted(self.memories[location], key=lambda x: {"CORE": 0, "PERMANENT": 1, "EPHEMERAL": 2}.get(x.type, 3))
            for m in sorted_entries:
                context_lines.append(f"   [{m.type}] {m.title}: {m.text}")
        else:
            context_lines.append(f"📍 You are in {location}. You have no previous memories here.")
            
        return "\n".join(context_lines)

    def synthesize(self, location: str, action: str, result: str, step: int):
        """Synthétise l'action via LLM avec une extraction JSON ultra-robuste."""
        
        # 1. Filtre rapide pour économiser des appels LLM
        if len(result) < 40 and any(k in result.lower() for k in ["nothing", "taken", "dropped", "closed"]):
            return

        existing_txt = self.get_context(location)
        prompt = MEMORY_SYNTHESIS_PROMPT.format(
            location=location, action=action, result=result, existing=existing_txt
        )

        try:
            response = self.call_llm(prompt, "You are a JSON Memory System.", seed=step, max_tokens=1000)
            
            import re
            import json

            # --- EXTRACTION ROBUSTE ---
            # On cherche le premier '{' et le dernier '}' avec des parenthèses () pour créer le groupe 1
            json_match = re.search(r'(\{.*\})', response, re.DOTALL)
            
            if json_match:
                json_str = json_match.group(1) # Maintenant le groupe 1 existe !
            else:
                # Si la regex échoue, on tente un nettoyage manuel des balises Markdown
                json_str = response.strip()
                if "```json" in json_str:
                    json_str = json_str.split("```json")[1].split("```")[0]
                elif "```" in json_str:
                    json_str = json_str.split("```")[1].split("```")[0]

            # --- NETTOYAGE DES GUILLEMETS (Auto-fix fréquent pour LLM) ---
            # Remplace les guillemets "smart" ou simples si le LLM s'est trompé
            json_str = json_str.replace('’', "'").replace('‘', "'")
            
            data = json.loads(json_str)

            # --- LOGIQUE DE MISE À JOUR ---
            if data.get("should_remember"):
                if location not in self.memories:
                    self.memories[location] = []

                # Gestion des Supressions (Supersedes)
                to_delete = data.get("supersedes", [])
                if to_delete:
                    self.memories[location] = [
                        m for m in self.memories[location] 
                        if m.title not in to_delete
                    ]

                # Ajout des nouvelles mémoires
                for item in data.get("memories", []):
                    if "title" not in item or "text" not in item: continue
                    
                    new_mem = MemoryEntry(
                        category=item.get("category", "INFO"),
                        type=item.get("type", "EPHEMERAL"),
                        title=item["title"],
                        text=item["text"],
                        turn_created=step
                    )
                    
                    # On évite d'ajouter deux fois le même titre dans le même lieu
                    if not any(m.title == new_mem.title for m in self.memories[location]):
                        self.memories[location].append(new_mem)
                        print(f"💾 [MEMORY SAVED] [{new_mem.type}] {new_mem.title}")

                self.save_to_md()

        except json.JSONDecodeError as je:
            print(f"response MEMORY {response}")
            print(f"⚠️ JSON Format Error: {je}. Check LLM output.")
        except Exception as e:
            # C'est ici que tu avais l'erreur "no such group"
            print(f"response MEMORY {response}")
            print(f"⚠️ Memory Synthesis Warning: {e}")
            print(f"response MEMORY {response}")