"""
memory.py — Conversation memory using ChromaDB + sentence-transformers.

Stores past conversation turns as vector embeddings.
Retrieves semantically relevant past context for each new query.
"""

from __future__ import annotations
import time
import os
from typing import List

_EMBED_MODEL_NAME = "all-MiniLM-L6-v2"   # 22 MB, fast CPU embed
_embed  = None
_client = None
_col    = None


def _model():
    global _embed
    if _embed is None:
        from sentence_transformers import SentenceTransformer
        print("[memory] Loading embedding model ...", flush=True)
        _embed = SentenceTransformer(_EMBED_MODEL_NAME)
        print("[memory] Embedding model ready", flush=True)
    return _embed


def _collection():
    global _client, _col
    if _col is None:
        import chromadb
        _client = chromadb.PersistentClient(path="./memory_db")
        _col = _client.get_or_create_collection(
            name="conversations",
            metadata={"hnsw:space": "cosine"},
        )
    return _col


def store(user_msg: str, ai_response: str, session: str = "default") -> None:
    """Store a conversation turn as an embedding."""
    try:
        text = f"User: {user_msg}\nAssistant: {ai_response}"
        emb  = _model().encode(text).tolist()
        doc_id = f"{session}_{int(time.time() * 1000)}"
        _collection().add(
            documents=[text],
            embeddings=[emb],
            ids=[doc_id],
            metadatas=[{
                "session": session,
                "timestamp": time.time(),
                "user_msg": user_msg[:200],
            }],
        )
    except Exception as e:
        print(f"[memory] store error: {e}", flush=True)


def retrieve(query: str, session: str = "default", top_k: int = 3) -> List[str]:
    """Return top_k semantically similar past exchanges."""
    try:
        col = _collection()
        if col.count() == 0:
            return []
        emb     = _model().encode(query).tolist()
        n       = min(top_k, col.count())
        results = col.query(
            query_embeddings=[emb],
            n_results=n,
            where={"session": session},
        )
        return results["documents"][0] if results["documents"] else []
    except Exception as e:
        print(f"[memory] retrieve error: {e}", flush=True)
        return []


def clear(session: str = "default") -> int:
    """Delete all memory entries for a session."""
    try:
        col = _collection()
        ids = col.get(where={"session": session})["ids"]
        if ids:
            col.delete(ids=ids)
        return len(ids)
    except Exception as e:
        print(f"[memory] clear error: {e}", flush=True)
        return 0


def list_recent(session: str = "default", limit: int = 10) -> List[dict]:
    """Return recent conversation entries for a session."""
    try:
        col = _collection()
        res = col.get(where={"session": session}, limit=limit,
                      include=["documents", "metadatas"])
        out = []
        for doc, meta in zip(res["documents"], res["metadatas"]):
            out.append({"text": doc, "timestamp": meta.get("timestamp"),
                        "user_msg": meta.get("user_msg", "")})
        out.sort(key=lambda x: x["timestamp"] or 0, reverse=True)
        return out[:limit]
    except Exception as e:
        print(f"[memory] list error: {e}", flush=True)
        return []


def stats() -> dict:
    """Global memory statistics."""
    try:
        col = _collection()
        return {
            "total_entries": col.count(),
            "embed_model":   _EMBED_MODEL_NAME,
            "db_path":       "./memory_db",
        }
    except Exception:
        return {"total_entries": 0, "embed_model": _EMBED_MODEL_NAME}