import time

def store_to_memory(client, collection, query, results, user_id=None):
    """
    Save search results into the Chroma collection,
    tag each record with optional user_id for session filtering.
    """
    docs, metadatas, ids = [], [], []
    for i, r in enumerate(results):
        doc_text = f"{r.get('title')} — {r.get('author')} — {r.get('desc')}"
        meta = {
            "query": query,
            "title": r.get("title"),
            "author": r.get("author"),
            "year": r.get("year"),
            "source": "openlibrary"
        }
        if user_id:
            meta["user_id"] = user_id

        docs.append(doc_text)
        metadatas.append(meta)
        ids.append(f"{user_id or 'anon'}_{int(time.time())}_{i}")

    collection.add(documents=docs, metadatas=metadatas, ids=ids)

    try:
        if hasattr(client, "persist"):
            client.persist()
    except Exception:
        pass
    
def retrieve_similar(collection, text, n=3, user_id=None):
    """
    Query the collection for documents semantically similar to text,
    optional user_id filters results to a single user's memory.
    """
    try:
        if user_id:
            res = collection.query(query_texts=[text], n_results=n, where={"user_id": user_id})
        else:
            res = collection.query(query_texts=[text], n_results=n)

        docs = res.get("documents", [[]])[0]
        metas = res.get("metadatas", [[]])[0]
        combined = []

        for doc, meta in zip(docs, metas):
            combined.append({
                "title": meta.get("title"),
                "author": meta.get("author"),
                "doc": doc
            })

        return combined
    except Exception:
        return []

# to test #
    
from ai_agent import init_chromadb, fetch_openlibrary_by_title_or_subject, store_to_memory, retrieve_similar

client, collection = init_chromadb(".chroma_demo_memory")

results = fetch_openlibrary_by_title_or_subject("Atomic Habits", limit=3)
print("Fetched results:", results)

store_to_memory(client, collection, "Atomic Habits search", results, user_id="John")
print("Stored successfully for user ankit")

similar = retrieve_similar(collection, "books about habits", user_id="John")
print("Memory recall:", similar)