#!/usr/bin/env python3
"""
Auto-Memory Recall — query the auto-memory ChromaDB collection.
Usage:
  python3 auto-memory-recall.py "how did we fix nginx"
  python3 auto-memory-recall.py --query "chromadb setup" --limit 3 --topic infrastructure
"""

import sys, json, os, argparse
import requests
import chromadb

CHROMADB_HOST = os.environ.get("CHROMADB_HOST", "192.168.86.25")
CHROMADB_PORT = int(os.environ.get("CHROMADB_PORT", "8000"))
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://192.168.86.40:11434")
EMBED_MODEL = "nomic-embed-text"
COLLECTION = "auto-memory"


def ollama_embed(text: str) -> list[float]:
    r = requests.post(
        f"{OLLAMA_URL}/api/embed",
        json={"model": EMBED_MODEL, "input": text},
        timeout=10,
    )
    r.raise_for_status()
    return r.json()["embeddings"][0]


def recall(query: str, limit: int = 5, topic: str = None, agent_id: str = None) -> list[dict]:
    client = chromadb.HttpClient(host=CHROMADB_HOST, port=CHROMADB_PORT)
    collection = client.get_or_create_collection(name=COLLECTION, metadata={"hnsw:space": "cosine"})

    where = {}
    if topic:
        where["topic"] = topic
    if agent_id:
        where["agent_id"] = agent_id

    try:
        embedding = ollama_embed(query)
        results = collection.query(
            query_embeddings=[embedding],
            n_results=limit,
            where=where if where else None,
        )
    except Exception:
        # Fallback to text query if embedding fails
        results = collection.query(
            query_texts=[query],
            n_results=limit,
            where=where if where else None,
        )

    memories = []
    if results and results["documents"]:
        for i, doc in enumerate(results["documents"][0]):
            meta = results["metadatas"][0][i] if results["metadatas"] else {}
            dist = results["distances"][0][i] if results["distances"] else None
            memories.append({
                "memory": doc,
                "similarity": round(1 - dist, 3) if dist is not None else None,
                "id": results["ids"][0][i] if results["ids"] else None,
                **meta,
            })
    return memories


def main():
    parser = argparse.ArgumentParser(description="Recall memories from auto-memory collection")
    parser.add_argument("query", nargs="?", help="Search query")
    parser.add_argument("--query", "-q", dest="query_flag")
    parser.add_argument("--limit", "-n", type=int, default=5)
    parser.add_argument("--topic", "-t")
    parser.add_argument("--agent-id", "-a")
    args = parser.parse_args()

    query = args.query or args.query_flag
    if not query:
        parser.print_help()
        sys.exit(1)

    results = recall(query, args.limit, args.topic, args.agent_id)
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()