#!/bin/bash
# Multi-RAG query router
# Usage: rag_query.sh <topic> "query text"
# Topics: bkd, saham, default
#
# Examples:
#   rag_query.sh bkd "syarat Karis Karsu"
#   rag_query.sh saham "GJTL analysis"
#   rag_query.sh default "apa itu LSD"

TOPIC="${1:-default}"
QUERY="$2"

if [ -z "$QUERY" ]; then
    echo "Usage: rag_query.sh <topic> \"query\""
    echo "Topics: bkd, saham, default"
    exit 1
fi

VENV="/home/ubuntu/venv_trading/bin/python3"

case "$TOPIC" in
  bkd)
    RAG_DIR="/home/ubuntu/.hermes/rag_bkd_chroma"
    COLLECTION="bkd_sosialisasi"
    ;;
  dohi)
    RAG_DIR="/home/ubuntu/.hermes/rag_dohi_chroma"
    COLLECTION="dohi_iohi"
    ;;
  saham)
    RAG_DIR="/home/ubuntu/.hermes/rag_saham_chroma"
    COLLECTION="saham"
    ;;
  *)
    echo "[default — no RAG, use general knowledge]"
    exit 0
    ;;
esac

if [ ! -d "$RAG_DIR" ]; then
    echo "RAG not found: $RAG_DIR"
    exit 1
fi

$VENV - <<EOF
import chromadb
from sentence_transformers import SentenceTransformer

query = """${QUERY}"""

chroma_client = chromadb.PersistentClient(path="${RAG_DIR}")
collection = chroma_client.get_or_create_collection("${COLLECTION}")
model = SentenceTransformer("all-MiniLM-L6-v2")

results = collection.query(
    query_embeddings=model.encode([query]).tolist(),
    n_results=3
)

sources = []
for i, (doc, meta) in enumerate(zip(results['documents'][0], results['metadatas'][0])):
    clean = doc.replace('<!-- Slide number: N -->', '').replace('<!-- Nova -->', '')
    clean = clean.replace('![](', '').replace(')', ' ')
    clean = ' '.join(clean.split())
    
    fname = meta.get('filename', 'unknown')
    sources.append((i+1, fname, clean))

for idx, fname, text in sources:
    print(f"[{idx}] 📄 {fname}")
    print(f"    {text[:500]}")
    if len(text) > 500:
        print(f"    ...")
    print()
EOF