memory/memory.py

from __future__ import annotations
import os
import uuid
import asyncio
import json
from typing import Optional
from dotenv import load_dotenv
import psycopg2
from psycopg2.extras import Json
from sentence_transformers import SentenceTransformer
from server.models import EpisodeMemory, AblationRecord

load_dotenv()

EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"


def get_connection():
    return psycopg2.connect(
        host     = os.getenv("DB_HOST", "127.0.0.1"),
        port     = int(os.getenv("DB_PORT", 5433)),
        dbname   = os.getenv("DB_NAME", "sre_env"),
        user     = os.getenv("DB_USER", "postgres"),
        password = os.getenv("DB_PASSWORD", ""),
    )


class AgenticMemory:
    """
    Episode-level agentic memory using PostgreSQL + pgvector.

    Pattern:
    - Episode start: get_episodic() + search_semantic() → RAM cache
    - Steps 1..N: read from RAM cache (zero DB calls)
    - Episode end: asyncio.create_task(write()) → non-blocking
    """

    def __init__(self, backend: str = "pgvector"):
        self.backend = backend
        self._model: Optional[SentenceTransformer] = None
        if backend == "pgvector":
            self._model = SentenceTransformer(EMBED_MODEL)

    # ------------------------------------------------------------------
    # MCP tools
    # ------------------------------------------------------------------

    def get_episodic(self, task_id: str, limit: int = 5) -> str:
        """
        Retrieves last N episodes for this task_id.
        Called once at episode start — result cached in RAM.
        """
        if self.backend == "disabled":
            return ""
        try:
            conn = get_connection()
            cur  = conn.cursor()
            cur.execute("""
                SELECT task_id, task_success, total_reward,
                       steps_taken, summary
                FROM episodes
                WHERE task_id = %s
                ORDER BY created_at DESC
                LIMIT %s
            """, (task_id, limit))
            rows = cur.fetchall()
            cur.close()
            conn.close()

            if not rows:
                return ""

            lines = ["Past episodes for this task:"]
            for row in rows:
                tid, success, reward, steps, summary = row
                lines.append(
                    f"- task={tid} success={success} "
                    f"reward={reward:.3f} steps={steps} | {summary[:100]}"
                )
            return "\n".join(lines)

        except Exception as e:
            print(f"get_episodic error: {e}")
            return ""

    def search_semantic(self, state_summary: str, top_k: int = 3) -> str:
        """
        Finds top-k semantically similar past episodes via vector search.
        Called once at episode start — result cached in RAM.
        """
        if self.backend == "disabled" or not self._model:
            return ""
        try:
            emb  = self._model.encode(state_summary, normalize_embeddings=True)
            emb_list = emb.tolist()

            conn = get_connection()
            cur  = conn.cursor()
            cur.execute("""
                SELECT task_id, task_success, total_reward,
                       steps_taken, summary,
                       1 - (state_emb <=> %s::vector) AS similarity
                FROM episodes
                ORDER BY state_emb <=> %s::vector
                LIMIT %s
            """, (emb_list, emb_list, top_k))
            rows = cur.fetchall()
            cur.close()
            conn.close()

            if not rows:
                return ""

            lines = ["Semantically similar past episodes:"]
            for row in rows:
                tid, success, reward, steps, summary, sim = row
                lines.append(
                    f"- task={tid} success={success} "
                    f"reward={reward:.3f} steps={steps} "
                    f"similarity={sim:.3f} | {summary[:100]}"
                )
            return "\n".join(lines)

        except Exception as e:
            print(f"search_semantic error: {e}")
            return ""

    def write(self, memory: EpisodeMemory):
        """
        Writes episode to DB. Called non-blocking via asyncio.create_task().
        Stores state embedding for future semantic search.
        """
        if self.backend == "disabled":
            return
        try:
            emb = None
            if self._model and memory.summary:
                emb = self._model.encode(
                    memory.summary, normalize_embeddings=True
                ).tolist()

            conn = get_connection()
            cur  = conn.cursor()
            cur.execute("""
                INSERT INTO episodes
                    (episode_id, task_id, task_success, total_reward,
                     steps_taken, actions, summary, state_emb)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            """, (
                memory.episode_id,
                memory.task_id,
                memory.task_success,
                memory.total_reward,
                memory.steps_taken,
                Json(memory.actions),
                memory.summary,
                emb,
            ))
            conn.commit()
            cur.close()
            conn.close()

        except Exception as e:
            print(f"memory.write error: {e}")

    async def write_async(self, memory: EpisodeMemory):
        """Async wrapper for non-blocking episode end write."""
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, self.write, memory)

    # ------------------------------------------------------------------
    # Ablation tracking (point 7)
    # ------------------------------------------------------------------

    def write_ablation_record(self, record: AblationRecord):
        """Saves ablation metrics after each epoch."""
        if self.backend == "disabled":
            return
        try:
            conn = get_connection()
            cur  = conn.cursor()
            cur.execute("""
                INSERT INTO ablation_records
                    (epoch, task_id, memory_backend, mean_reward,
                     mean_steps_to_resolution, task_success_rate)
                VALUES (%s, %s, %s, %s, %s, %s)
            """, (
                record.epoch,
                record.task_id,
                record.memory_backend,
                record.mean_reward,
                record.mean_steps_to_resolution,
                record.task_success_rate,
            ))
            conn.commit()
            cur.close()
            conn.close()
        except Exception as e:
            print(f"write_ablation_record error: {e}")

    def get_ablation_records(self, task_id: str) -> list:
        """Retrieves ablation records for Gradio dashboard plotting."""
        if self.backend == "disabled":
            return []
        try:
            conn = get_connection()
            cur  = conn.cursor()
            cur.execute("""
                SELECT epoch, memory_backend, mean_reward,
                       mean_steps_to_resolution, task_success_rate
                FROM ablation_records
                WHERE task_id = %s
                ORDER BY epoch ASC
            """, (task_id,))
            rows = cur.fetchall()
            cur.close()
            conn.close()
            return rows
        except Exception as e:
            print(f"get_ablation_records error: {e}")
            return []

    # ------------------------------------------------------------------
    # Cache builder — called once at episode start
    # ------------------------------------------------------------------

    def build_system_context(
        self,
        task_id:       str,
        state_summary: str,
        runbook:       str = "",
    ) -> str:
        """
        Merges episodic memory + semantic search + runbook
        into a single system_context string.
        Cached in RAM for the full episode.
        """
        episodic = self.get_episodic(task_id)
        semantic = self.search_semantic(state_summary)

        parts = []
        if runbook:
            parts.append(f"=== Runbook ===\n{runbook}")
        if episodic:
            parts.append(f"=== Episodic Memory ===\n{episodic}")
        if semantic:
            parts.append(f"=== Similar Episodes ===\n{semantic}")

        return "\n\n".join(parts)