"""
STEW — Vector Memory Engine (ChromaDB)
Long-term semantic memory: stores facts, documents, and project context
that can be retrieved by semantic similarity.
"""
import hashlib
from typing import Optional
from config import Config


class VectorMemory:
    """
    Persistent semantic memory using ChromaDB.
    Falls back to in-memory store if ChromaDB is not available.
    """

    def __init__(self, persist_dir: str = "/tmp/stew_vector_db"):
        self._client = None
        self._collection = None
        self._persist_dir = persist_dir
        self._fallback: list[dict] = []

    def _init_chroma(self):
        if self._collection is not None:
            return True
        try:
            import chromadb
            self._client = chromadb.PersistentClient(path=self._persist_dir)
            self._collection = self._client.get_or_create_collection(
                name="stew_memory",
                metadata={"hnsw:space": "cosine"},
            )
            return True
        except Exception:
            return False

    # ── Document storage ─────────────────────────────────────────────────────

    def store(self, text: str, metadata: dict | None = None, doc_id: str | None = None) -> str:
        """Store a text chunk in vector memory."""
        if not doc_id:
            doc_id = hashlib.sha256(text.encode()).hexdigest()[:16]
        meta = metadata or {}

        if self._init_chroma():
            try:
                self._collection.upsert(
                    ids=[doc_id],
                    documents=[text],
                    metadatas=[meta],
                )
                return f"✅ Stored in vector memory (id: {doc_id})"
            except Exception as e:
                return f"Vector store error: {e}"
        else:
            # Fallback: in-memory list
            self._fallback.append({"id": doc_id, "text": text, "meta": meta})
            return f"✅ Stored in fallback memory (id: {doc_id})"

    def search(self, query: str, n_results: int = 5) -> str:
        """Semantic search over stored memory."""
        if self._init_chroma():
            try:
                results = self._collection.query(
                    query_texts=[query],
                    n_results=min(n_results, self._collection.count() or 1),
                )
                docs = results.get("documents", [[]])[0]
                if not docs:
                    return "No relevant memories found."
                lines = [f"[{i+1}] {d[:300]}" for i, d in enumerate(docs)]
                return "🧠 Relevant memories:\n" + "\n\n".join(lines)
            except Exception as e:
                return f"Vector search error: {e}"
        else:
            # Simple substring fallback
            matches = [f["text"] for f in self._fallback if query.lower() in f["text"].lower()]
            return ("🧠 Matches:\n" + "\n".join(m[:300] for m in matches[:5])) if matches else "No matches."

    def store_project(self, project_name: str, description: str) -> str:
        """Store a project description."""
        return self.store(
            f"Project: {project_name}\n{description}",
            metadata={"type": "project", "name": project_name},
            doc_id=f"project_{project_name[:20]}",
        )

    def count(self) -> int:
        if self._init_chroma():
            try:
                return self._collection.count()
            except Exception:
                pass
        return len(self._fallback)

    def clear(self) -> str:
        if self._init_chroma():
            try:
                self._client.delete_collection("stew_memory")
                self._collection = None
                return "🗑️ Vector memory cleared."
            except Exception as e:
                return f"Clear error: {e}"
        self._fallback.clear()
        return "🗑️ Fallback memory cleared."


# Module-level singleton
_vector_memory: Optional[VectorMemory] = None

def get_vector_memory() -> VectorMemory:
    global _vector_memory
    if _vector_memory is None:
        _vector_memory = VectorMemory()
    return _vector_memory