import { ChromaClient } from 'chromadb';
import { MODEL_ID, embed, embedMany } from './embed.js';

// ==============================|| Chroma collections ||============================== //
//
// Two collections, deliberately separate:
//
//   intent_examples — ~20 phrasings per intent. Answers "which question is
//                     this?", never "what is the number?".
//   console_docs    — chunked markdown. Genuine document Q&A.
//
// NO OPERATIONAL DATA IS EMBEDDED. No bookings, riders or customers ever enter
// the vector store. A vector store is a snapshot and this data changes by the
// minute; every figure still comes from a live API call through the existing
// typed functions. That is the whole reason this design is safe to build.

export const INTENTS_COLLECTION = 'intent_examples';
export const DOCS_COLLECTION = 'console_docs';

const client = new ChromaClient({ path: process.env.CHROMA_URL || 'http://localhost:8000' });

// Chroma would otherwise call its own default embedding function (which
// downloads a different model). We embed ourselves so queries and documents
// are guaranteed to come from the same model.
const noopEmbeddingFunction = { generate: async (texts) => embedMany(texts) };

export const getCollection = async (name) =>
  client.getOrCreateCollection({
    name,
    embeddingFunction: noopEmbeddingFunction,
    metadata: {
      // Stamped so a model swap is detectable rather than silently degrading
      // every score. Re-seed after changing it.
      'hnsw:space': 'cosine',
      embeddingModel: MODEL_ID
    }
  });

export const resetCollection = async (name) => {
  try {
    await client.deleteCollection({ name });
  } catch {
    // Not present yet — nothing to delete.
  }
  return getCollection(name);
};

// Chroma returns cosine DISTANCE (0 = identical). Similarity is the useful
// direction for a human-facing confidence, so convert once, here, rather than
// leaving every call site to remember which way round it is.
const toSimilarity = (distance) => 1 - Number(distance);

export const query = async (name, text, topK = 5) => {
  const collection = await getCollection(name);
  const vector = await embed(text);
  const res = await collection.query({ queryEmbeddings: [vector], nResults: topK });

  const ids = res.ids?.[0] || [];
  return ids.map((id, i) => ({
    id,
    document: res.documents?.[0]?.[i] || '',
    metadata: res.metadatas?.[0]?.[i] || {},
    score: toSimilarity(res.distances?.[0]?.[i] ?? 1)
  }));
};

export const health = async () => {
  await client.heartbeat();
  const names = (await client.listCollections()).map((c) => c.name ?? c);
  const counts = {};
  for (const name of [INTENTS_COLLECTION, DOCS_COLLECTION]) {
    if (!names.includes(name)) continue;
    // eslint-disable-next-line no-await-in-loop
    counts[name] = await (await getCollection(name)).count();
  }
  return { chroma: 'up', collections: counts };
};