import { ChromaClient } from 'chromadb';
import { OpenAIEmbeddings } from 'langchain/embeddings/openai';

const chroma = new ChromaClient({
  host: process.env.CHROMA_HOST || 'localhost',
  port: parseInt(process.env.CHROMA_PORT, 10) || 8000,
});

const embeddings = new OpenAIEmbeddings({
  openAIApiKey: process.env.OPENAI_API_KEY,
});

const COLLECTION_NAME = 'rag_collection';

/**
 * Ensure the collection exists in ChromaDB.
 */
async function ensureCollection() {
  const collections = await chroma.listCollections();
  if (!collections.includes(COLLECTION_NAME)) {
    await chroma.createCollection({ name: COLLECTION_NAME });
  }
}

/**
 * Add an array of documents to the vector store.
 * @param {string[]} docs
 */
export async function addDocuments(docs) {
  await ensureCollection();
  const ids = docs.map((_, idx) => `doc-${Date.now()}-${idx}`);
  const embeddingsResult = await embeddings.embedDocuments(docs);
  await chroma.add({
    collection_name: COLLECTION_NAME,
    ids,
    documents: docs,
    embeddings: embeddingsResult,
  });
}

/**
 * Query the vector store for the most relevant documents.
 * @param {string} queryText
 * @param {number} nResults
 * @returns {Promise<{documents: string[]}>}
 */
export async function query(queryText, nResults = 3) {
  await ensureCollection();
  const embedding = await embeddings.embedQuery(queryText);
  const results = await chroma.query({
    collection_name: COLLECTION_NAME,
    query_embeddings: [embedding],
    n_results: nResults,
  });
  // ChromaDB returns an array of objects; extract documents
  const docs = results.documents || [];
  return { documents: docs };
}