memory/stores.py

"""
In-process and persistent memory stores.

Provides the four memory primitives used by MemoryManager:
  ShortTermMemory  — in-RAM conversation buffer (compressed on overflow)
  WorkingMemory    — per-run structured tool-output cache
  LongTermMemory   — persistent analysis reports + preferences (Supabase / SQLite)
  EpisodicMemory   — timestamped analysis episodes with outcome tracking
"""
from datetime import datetime
from memory.backends import _embed, _get_supabase, _using_supabase, _get_sqlite


# ── 1. SHORT-TERM MEMORY ──────────────────────────────────────────────────────

class ShortTermMemory:
    """
    In-RAM conversation buffer for the current analysis run.

    Problem it solves:
      AutoGen agents have a context window limit. In a long analysis with
      6 agents passing lots of data, earlier messages get pushed out.
      This class keeps a running summary so agents never lose key context.

    How compression works:
      When the buffer exceeds ~3000 tokens, it scans the oldest messages
      for keyword-bearing lines (buy/sell/rsi/confidence etc.) and
      condenses them into a summary string. Only the last 4 messages
      are kept verbatim — everything before becomes a summary.
    """

    def __init__(self):
        self.messages: list[dict] = []
        self.summary:  str        = ""

    def add(self, role: str, content: str):
        self.messages.append({"role": role, "content": str(content)[:2000]})
        if self._token_estimate() > 3000:
            self._compress()

    def _compress(self):
        if len(self.messages) <= 4:
            return
        old  = self.messages[:-4]
        text = "\n".join(f"{m['role']}: {m['content']}" for m in old)

        # Extractive summary — pull out lines with key financial terms
        keywords = {"buy","sell","hold","rsi","pe","signal","confidence",
                    "recommend","error","macd","fusion","score","forecast"}
        key_lines = [
            l for l in text.split("\n")
            if any(kw in l.lower() for kw in keywords)
        ]
        self.summary  = " | ".join(key_lines[:10])
        self.messages = self.messages[-4:]

    def _token_estimate(self) -> int:
        return int(sum(len(m["content"].split()) for m in self.messages) * 1.3)

    def get_context(self) -> str:
        parts = []
        if self.summary:
            parts.append(f"## Earlier in this analysis:\n{self.summary}")
        if self.messages:
            recent = "\n".join(
                f"{m['role']}: {m['content'][:300]}" for m in self.messages[-3:]
            )
            parts.append(f"## Recent messages:\n{recent}")
        return "\n\n".join(parts)

    def clear(self):
        self.messages = []
        self.summary  = ""


# ── 2. WORKING MEMORY ──────────────────────────────────────────────────────────

class WorkingMemory:
    """
    Shared data store for a SINGLE analysis run.

    Problem it solves:
      DataAgent calls 7 tools and gets back price data, fundamentals,
      signals, fusion score etc. TechnicalAnalyst and FundamentalAnalyst
      both need this data. Without working memory they'd each re-fetch
      it (slow, wasteful, costs money).

    How it works:
      DataAgent stores its tool results here via store().
      Other agents read them via get().
      The orchestrator resets it at the start of each run.

    This is different from ShortTermMemory:
      ShortTermMemory = conversation history (text messages)
      WorkingMemory   = structured tool output (dicts / JSON)
    """

    def __init__(self):
        self._store: dict = {}

    def store(self, key: str, value):
        """Store a tool result. Key should be the tool name."""
        self._store[key] = value

    def get(self, key: str, default=None):
        return self._store.get(key, default)

    def get_all(self) -> dict:
        return dict(self._store)

    def has(self, key: str) -> bool:
        return key in self._store

    def to_context_string(self) -> str:
        """
        Format all stored tool results as a readable string
        for injection into agent prompts.
        """
        if not self._store:
            return ""
        lines = ["## Tool results from DataAgent:"]
        for key, val in self._store.items():
            lines.append(f"\n### {key}")
            if isinstance(val, dict):
                # Show summary fields only — not the full data dump
                summary_keys = ["ticker", "signal", "ensemble_signal", "recommendation",
                                "final_score", "confidence", "summary", "error",
                                "overall_sentiment", "trend_direction", "upside_pct",
                                "ml_signal", "current_price"]
                for k in summary_keys:
                    if k in val:
                        lines.append(f"  {k}: {val[k]}")
            else:
                lines.append(f"  {str(val)[:200]}")
        return "\n".join(lines)

    def reset(self):
        self._store = {}


# ── 3. LONG-TERM MEMORY ────────────────────────────────────────────────────────

class LongTermMemory:
    """
    Stores analysis reports persistently — across sessions, deployments,
    and devices. Uses Supabase with pgvector for semantic search.

    Falls back to SQLite if Supabase is not configured.

    Two things stored:
      stock_analyses   — the full report for each ticker analysis
      user_preferences — investment style, risk tolerance, sector focus etc.

    The semantic search (pgvector) lets agents ask:
      "Have we ever been bullish on a stock in this sector before?"
      "What did we say last time we saw RSI this low?"
    """

    def store_analysis(self, ticker: str, user_id: str, report: str,
                       recommendation: str, confidence: int):
        """Save a completed analysis. Embeds the report for later semantic search."""
        now = datetime.now().isoformat()

        if _using_supabase():
            try:
                db  = _get_supabase()
                row = {
                    "user_id":        user_id,
                    "ticker":         ticker,
                    "recommendation": recommendation,
                    "confidence":     confidence,
                    "report":         report,
                }
                vec = _embed(report)
                if vec:
                    row["embedding"] = vec
                db.table("stock_analyses").insert(row).execute()
                return
            except Exception as e:
                print(f"[memory] Supabase store_analysis failed: {e} — using SQLite")

        # SQLite fallback
        conn = _get_sqlite()
        conn.execute(
            "INSERT INTO stock_analyses VALUES (NULL,?,?,?,?,?,?)",
            (user_id, ticker, recommendation, confidence, report, now)
        )
        conn.commit()

    def recall_past_analyses(self, ticker: str, user_id: str,
                              n: int = 3) -> list[str]:
        """
        Get the N most recent analyses for this ticker.
        Uses recency — for the same ticker, most recent = most relevant.
        """
        if _using_supabase():
            try:
                db  = _get_supabase()
                res = (
                    db.table("stock_analyses")
                      .select("report, recommendation, confidence, created_at")
                      .eq("user_id", user_id)
                      .eq("ticker",  ticker)
                      .order("created_at", desc=True)
                      .limit(n)
                      .execute()
                )
                return [
                    f"[{r['created_at'][:10]}] {r['recommendation']} "
                    f"({r['confidence']}% conf)\n{r['report'][:300]}"
                    for r in (res.data or [])
                ]
            except Exception as e:
                print(f"[memory] recall_past_analyses failed: {e}")
                return []

        # SQLite fallback
        try:
            conn = _get_sqlite()
            rows = conn.execute(
                """SELECT report, recommendation, confidence, created_at
                   FROM stock_analyses WHERE user_id=? AND ticker=?
                   ORDER BY created_at DESC LIMIT ?""",
                (user_id, ticker, n)
            ).fetchall()
            return [
                f"[{r[3][:10]}] {r[1]} ({r[2]}% conf)\n{r[0][:300]}"
                for r in rows
            ]
        except Exception:
            return []

    def semantic_search(self, query: str, user_id: str,
                         n: int = 3) -> list[str]:
        """
        Find semantically similar past analyses using pgvector.
        Only works when Supabase is configured — skipped in local fallback.

        Example query: "bullish semiconductor stock oversold RSI"
        Returns: snippets of past analyses with similar context.

        Requires the match_analyses() function in Supabase (see supabase_setup.sql).
        """
        if not _using_supabase():
            return []
        try:
            vec = _embed(query)
            if not vec:
                return []
            db  = _get_supabase()
            res = db.rpc("match_analyses", {
                "query_embedding": vec,
                "match_user_id":   user_id,
                "match_count":     n,
            }).execute()
            return [r["report"][:300] for r in (res.data or [])]
        except Exception as e:
            print(f"[memory] semantic_search failed: {e}")
            return []

    def store_preference(self, user_id: str, preference: str):
        """
        Save a user preference.
        Examples:
          "Prefers growth stocks over value"
          "Risk averse — avoids high beta stocks"
          "Focus on UK and European markets"
        """
        now = datetime.now().isoformat()

        if _using_supabase():
            try:
                db  = _get_supabase()
                row = {"user_id": user_id, "preference": preference}
                vec = _embed(preference)
                if vec:
                    row["embedding"] = vec
                db.table("user_preferences").insert(row).execute()
                return
            except Exception as e:
                print(f"[memory] store_preference failed: {e}")

        conn = _get_sqlite()
        conn.execute(
            "INSERT INTO user_preferences VALUES (NULL,?,?,?)",
            (user_id, preference, now)
        )
        conn.commit()

    def get_user_profile(self, user_id: str) -> str:
        """
        Return all stored preferences as a single string.
        Injected into every agent prompt so they know the user's style.
        """
        if _using_supabase():
            try:
                db  = _get_supabase()
                res = (
                    db.table("user_preferences")
                      .select("preference")
                      .eq("user_id", user_id)
                      .order("created_at", desc=True)
                      .limit(10)
                      .execute()
                )
                prefs = [r["preference"] for r in (res.data or [])]
                return "\n".join(prefs) if prefs else ""
            except Exception as e:
                print(f"[memory] get_user_profile failed: {e}")
                return ""

        try:
            conn = _get_sqlite()
            rows = conn.execute(
                "SELECT preference FROM user_preferences WHERE user_id=? "
                "ORDER BY created_at DESC LIMIT 10",
                (user_id,)
            ).fetchall()
            return "\n".join(r[0] for r in rows)
        except Exception:
            return ""


# ── 4. EPISODIC MEMORY ─────────────────────────────────────────────────────────

class EpisodicMemory:
    """
    Records every completed analysis as a timestamped episode.

    The key feature: outcome tracking.
    After you make a BUY call, you can later update the episode with
    what actually happened ("price rose 8.4% in 30 days — correct").
    This builds an honest track record over time.

    Great for the portfolio README: "Model accuracy: 23 correct out of 31 calls"
    """

    def record(self, user_id: str, ticker: str, recommendation: str,
               confidence: int, price: float, report: str):
        """Save a completed analysis episode."""
        now = datetime.now().isoformat()

        if _using_supabase():
            try:
                _get_supabase().table("episodes").insert({
                    "user_id":        user_id,
                    "ticker":         ticker,
                    "recommendation": recommendation,
                    "confidence":     confidence,
                    "price_at_time":  round(price, 4),
                    "report":         report,
                }).execute()
                return
            except Exception as e:
                print(f"[memory] record episode failed: {e}")

        conn = _get_sqlite()
        conn.execute(
            "INSERT INTO episodes VALUES (NULL,?,?,?,?,?,?,NULL,?)",
            (user_id, ticker, recommendation, confidence, price, report, now)
        )
        conn.commit()

    def get_ticker_history(self, user_id: str, ticker: str) -> str:
        """
        Returns a human-readable history of past calls on this ticker.
        Injected into agent prompts so they can say:
          "Last time we rated AAPL as HOLD at $178, it's now $195..."

        The outcome field shows whether the call was correct — agents
        can use this to calibrate their confidence.
        """
        rows = []

        if _using_supabase():
            try:
                res = (
                    _get_supabase().table("episodes")
                      .select("recommendation, confidence, price_at_time, outcome, created_at")
                      .eq("user_id", user_id)
                      .eq("ticker",  ticker)
                      .order("created_at", desc=True)
                      .limit(5)
                      .execute()
                )
                rows = res.data or []
            except Exception as e:
                print(f"[memory] get_ticker_history failed: {e}")

        else:
            try:
                conn = _get_sqlite()
                raw  = conn.execute(
                    """SELECT recommendation, confidence, price_at_time, outcome, created_at
                       FROM episodes WHERE user_id=? AND ticker=?
                       ORDER BY created_at DESC LIMIT 5""",
                    (user_id, ticker)
                ).fetchall()
                rows = [
                    {"recommendation": r[0], "confidence": r[1],
                     "price_at_time": r[2], "outcome": r[3], "created_at": r[4]}
                    for r in raw
                ]
            except Exception:
                pass

        if not rows:
            return ""

        lines = [f"## Past analyses of {ticker}:"]
        for r in rows:
            outcome = f" → {r['outcome']}" if r.get("outcome") else " → outcome pending"
            lines.append(
                f"- {r['created_at'][:10]}: {r['recommendation']} "
                f"({r['confidence']}% conf) at ${r.get('price_at_time', 0):.2f}"
                f"{outcome}"
            )
        return "\n".join(lines)

    def get_all_history(self, user_id: str) -> list[dict]:
        """All episodes for this user — used by the Streamlit sidebar."""
        if _using_supabase():
            try:
                res = (
                    _get_supabase().table("episodes")
                      .select("ticker, recommendation, confidence, "
                              "price_at_time, outcome, created_at")
                      .eq("user_id", user_id)
                      .order("created_at", desc=True)
                      .limit(50)
                      .execute()
                )
                return [
                    {"ticker": r["ticker"], "rec": r["recommendation"],
                     "conf": r["confidence"], "price": r["price_at_time"],
                     "outcome": r["outcome"], "date": r["created_at"][:10]}
                    for r in (res.data or [])
                ]
            except Exception:
                return []

        try:
            conn = _get_sqlite()
            rows = conn.execute(
                """SELECT ticker, recommendation, confidence, price_at_time,
                          outcome, created_at
                   FROM episodes WHERE user_id=?
                   ORDER BY created_at DESC LIMIT 50""",
                (user_id,)
            ).fetchall()
            return [
                {"ticker": r[0], "rec": r[1], "conf": r[2],
                 "price": r[3], "outcome": r[4], "date": r[5][:10]}
                for r in rows
            ]
        except Exception:
            return []

    def update_outcome(self, episode_id: int, outcome: str):
        """
        Record what actually happened after a recommendation.
        Call this ~30 days after a BUY/SELL call to build track record.
        e.g. outcome = "Price rose 8.4% in 30 days — BUY was correct"
        """
        if _using_supabase():
            try:
                _get_supabase().table("episodes")\
                    .update({"outcome": outcome})\
                    .eq("id", episode_id)\
                    .execute()
                return
            except Exception as e:
                print(f"[memory] update_outcome failed: {e}")

        conn = _get_sqlite()
        conn.execute("UPDATE episodes SET outcome=? WHERE id=?", (outcome, episode_id))
        conn.commit()