memory/manager.py

"""
manager.py — Unified Memory Manager
Coordina i 4 layer: Working, Episodic, Semantic, Reflection.

TAM (Token-Aware Memory) — get_context usa algoritmo Waterfall con budget token dinamico.
"""
from .working import WorkingMemory
from .episodic import EpisodicMemory
from .semantic import SemanticMemory
from .reflection import ReflectionMemory

import logging
_logger = logging.getLogger("memory.manager")


class MemoryManager:
    def __init__(self, total_token_budget: int = 4000):
        self.working    = WorkingMemory(max_entries=80)  # QF-5: 40→80 — sessioni lunghe multi-file
        self.episodic   = EpisodicMemory()
        self.semantic   = SemanticMemory()
        self.reflection = ReflectionMemory()

        # TAM: budget totale in token per get_context()
        self.total_token_budget = total_token_budget

        # Distribuzione percentuale iniziale (Waterfall: Reflection → Episodic → Semantic → Working)
        self._budget_distribution = {
            "reflection": 0.10,
            "episodic":   0.20,
            "semantic":   0.30,
            "working":    0.40,
        }

    async def init(self):
        # S274-BUG2: semantic.init() apre connessioni Supabase/ChromaDB — I/O bloccante.
        # asyncio.to_thread scarica sul thread pool per non bloccare l'event loop FastAPI.
        import asyncio as _asyncio
        await _asyncio.to_thread(self.episodic.init)
        await _asyncio.to_thread(self.semantic.init)
        # Auto-restore: se la memoria è vuota, carica l'ultimo snapshot da GitHub
        await self._auto_restore_semantic()

    async def close(self):
        self.episodic.close()

    # ── TAM helpers ────────────────────────────────────────────────────────────

    def _estimate_tokens(self, text: str) -> int:
        """Stima rapida dei token: caratteri / 4."""
        return len(text) // 4

    def _fill_layer(self, header: str, entries: list, budget: int) -> tuple[str, int]:
        """Riempie un layer rispettando il budget token.

        Restituisce (testo_layer, token_usati).
        Se il layer è vuoto o il budget è zero, restituisce ("", 0).
        Se una singola entry supera il budget, la tronca invece di scartarla.
        Il budget residuo non usato viene ceduto al layer successivo tramite il chiamante.
        """
        if not entries or budget <= 0:
            return "", 0

        header_str    = f"--- {header} ---"
        current_text  = header_str
        current_tokens = self._estimate_tokens(header_str)
        used_entries  = 0

        for entry in entries:
            entry_str    = str(entry)
            entry_tokens = self._estimate_tokens(entry_str)

            if current_tokens + entry_tokens + 1 > budget:
                if used_entries == 0:
                    # Prima entry troppo lunga: tronca intelligentemente
                    allowed_chars = (budget - current_tokens - 5) * 4
                    if allowed_chars > 100:
                        truncated     = entry_str[:allowed_chars] + "…"
                        current_text  += "\n" + truncated
                        current_tokens += self._estimate_tokens(truncated)
                # Budget esaurito — passa il residuo al layer successivo
                break

            current_text   += "\n" + entry_str
            current_tokens += entry_tokens + 1
            used_entries   += 1

        return current_text, current_tokens

    # ── get_context — algoritmo Waterfall TAM ──────────────────────────────────

    async def get_context(self, query: str, code_length: int = 0) -> str:
        """Assembla il contesto dai 4 layer con Waterfall Token Budget.

        TAM — Token-Aware Memory:
          Il budget non usato da un layer viene ceduto al successivo.
          Nessun layer può mai sforare il budget totale.

        Budget adattivo per code_length (prompt già grandi su iPhone):
          code_length > 8000 → 2000 token (stringente)
          code_length > 4000 → 3000 token (medio)
          default            → total_token_budget (4000)
        """
        if code_length > 8000:
            effective_budget = 2000
        elif code_length > 4000:
            effective_budget = 3000
        else:
            effective_budget = self.total_token_budget

        remaining_budget = effective_budget
        context_parts    = []

        # ── Layer 1: Reflection (10%) ──────────────────────────────────────────
        reflect_alloc = int(effective_budget * self._budget_distribution["reflection"])
        lessons       = self.reflection.get_relevant_lessons(query, n=5)
        lesson_lines  = []
        for l in lessons:
            if l["type"] == "failure":
                lesson_lines.append(f"EVITA: {l['avoid'][:300]}")
            else:
                lesson_lines.append(f"STRATEGIA: {l['strategy'][:300]}")

        reflect_text, reflect_used = self._fill_layer("Lezioni passate", lesson_lines, reflect_alloc)
        if reflect_text:
            context_parts.append(reflect_text)
        remaining_budget -= reflect_used

        # ── Layer 2: Episodic (20% + residuo reflection) ───────────────────────
        episodic_alloc = int(effective_budget * self._budget_distribution["episodic"]) + (reflect_alloc - reflect_used)
        episodes       = self.episodic.search_text(query, n=5)
        episode_lines  = [
            f"{ep.task} → {ep.output[:300]}"
            for ep in episodes
        ]
        episodic_text, episodic_used = self._fill_layer("Episodi passati", episode_lines, episodic_alloc)
        if episodic_text:
            context_parts.append(episodic_text)
        remaining_budget -= episodic_used

        # ── Layer 3: Semantic (30% + residuo episodic) ─────────────────────────
        semantic_alloc = int(effective_budget * self._budget_distribution["semantic"]) + (episodic_alloc - episodic_used)
        semantic_used  = 0
        if self.semantic.available:
            semantic_hits  = self.semantic.search(query, n_results=8)
            semantic_lines = [
                f"- {h['content'][:300]}"
                for h in semantic_hits
                if h["similarity"] > 0.3
            ]
            semantic_text, semantic_used = self._fill_layer("Conoscenza rilevante", semantic_lines, semantic_alloc)
            if semantic_text:
                context_parts.append(semantic_text)
            remaining_budget -= semantic_used

        # ── Layer 4: Working (tutto il budget residuo — layer più importante) ──
        working_budget = remaining_budget
        working_ctx    = self.working.get_context_string(n=15)
        if working_ctx:
            working_tokens = self._estimate_tokens(working_ctx)
            if working_tokens <= working_budget:
                context_parts.append(working_ctx)
            else:
                # Tronca preservando inizio (più recente = in coda, ma tronco i caratteri extra)
                allowed_chars = working_budget * 4
                context_parts.append(working_ctx[:allowed_chars] + "…")

        final_context = "\n\n".join(context_parts) if context_parts else ""
        _logger.info(
            "[MemoryManager] TAM context: %d/%d token (code_length=%d, layers=%d)",
            self._estimate_tokens(final_context), effective_budget, code_length, len(context_parts),
        )
        return final_context

    # ── Salvataggio dati ───────────────────────────────────────────────────────

    async def save_exchange(self, messages: list, response: str):
        """Salva uno scambio chat nella memoria."""
        user_msg = next((m["content"] for m in reversed(messages) if m["role"] == "user"), "")
        # Working: aggiungi utente + risposta
        if user_msg:
            self.working.add("user", user_msg)
        self.working.add("assistant", response)
        # Episodic: salva la coppia — S571: 500→2000
        self.episodic.add("chat", user_msg[:500], response[:2000], True)
        # Semantic: indicizza per similarity search futura — S571: combined 600→1100 chars
        if self.semantic.available and user_msg and len(response) > 50:
            combined = f"Q: {user_msg[:500]} A: {response[:800]}"
            self.semantic.add(combined, {"type": "chat", "query": user_msg[:300]})

    async def save_episode(self, type_: str, task: str, output: str, success: bool, tags: list | None = None):
        self.episodic.add(type_, task, output, success, tags)
        if self.semantic.available and task:
            self.semantic.add(task, {"type": type_, "success": success})

    async def search(self, query: str, n: int = 5, layer: str | None = None) -> list[dict]:
        results = []
        if layer in (None, "semantic") and self.semantic.available:
            for h in self.semantic.search(query, n_results=n):
                results.append({**h, "layer": "semantic"})
        if layer in (None, "episodic"):
            for ep in self.episodic.search_text(query, n=n):
                results.append({
                    "content": f"{ep.task} → {ep.output[:300]}",
                    "layer":   "episodic",
                    "type":    ep.type,
                    "success": ep.success,
                })
        if layer == "reflection":
            lessons = self.reflection.get_relevant_lessons(query, n=n)
            results.extend([{**l, "layer": "reflection"} for l in lessons])
        return results[:n]

    async def reflect(self, task: str, output: str, success: bool, error: str | None = None) -> dict:
        if success:
            self.reflection.record_success(task, output[:500])
            await self.save_episode("fix", task, output, True)
        else:
            self.reflection.record_failure(task, error or output[:500])
            await self.save_episode("error", task, error or output[:500], False)
        return {
            "recorded":     True,
            "top_patterns": self.reflection.get_top_patterns(5),
            "lessons":      self.reflection.get_relevant_lessons(task, 4),
        }

    # ── Auto-backup semantica cross-restart ─────────────────────────────────────

    async def _auto_restore_semantic(self) -> None:
        """Auto-restore: se la semantic memory è vuota, carica l'ultimo snapshot da GitHub.

        Chiamato dopo init() — garantisce continuità cross-restart (ChromaDB ephemeral + Supabase).
        Non-blocking: fallisce silenziosamente se GitHub non raggiungibile o snapshot assente.
        """
        import asyncio as _asyncio, os
        if not self.semantic.available:
            return
        count = await _asyncio.to_thread(self.semantic.count)
        if count > 0:
            return  # già popolata — Supabase ha i dati persistenti

        token = os.environ.get("GITHUB_TOKEN", "")
        if not token:
            return

        try:
            import urllib.request as _urq, json as _json, base64 as _b64
            req = _urq.Request(
                "https://api.github.com/repos/Baida98/AI/contents/data/semantic_snapshot.json",
                headers={
                    "Authorization": f"Bearer {token}",
                    "Accept":        "application/vnd.github.v3+json",
                    "User-Agent":    "agente-ai-backend",
                },
            )
            with _urq.urlopen(req, timeout=10) as resp:
                meta = _json.loads(resp.read())
            records = _json.loads(_b64.b64decode(meta["content"]).decode("utf-8"))
            if not records:
                return
            result = await _asyncio.to_thread(self.semantic.import_all, records)
            _logger.info(
                "[MemoryManager] ✓ Auto-restore semantica: %d record da GitHub snapshot (skip: %d)",
                result["imported"], result["skipped"],
            )
        except Exception as exc:
            _logger.debug(
                "[MemoryManager] Auto-restore semantica: snapshot non disponibile (%s)",
                exc.__class__.__name__,
            )

    def stats(self) -> dict:
        return {
            "working":    self.working.stats(),
            "episodic":   self.episodic.stats(),
            "semantic":   self.semantic.stats(),
            "reflection": self.reflection.stats(),
        }

    async def clear(self, layer: str | None = None):
        if layer in (None, "working"):
            self.working.clear()
        if layer in (None, "episodic"):
            import sqlite3
            if self.episodic._db:
                self.episodic._db.execute("DELETE FROM episodes")
                self.episodic._db.commit()