const {QdrantClient} = require('@qdrant/js-client-rest');
const {QDRANT, COLLECTIONS, getEnv} = require('@nexusai/shared');

let client;

// ********** QDrant Client **********
function getClient(){
    if(!client){
        client = new QdrantClient({url: getEnv('QDRANT_URL', QDRANT.DEFAULT_URL)});
    }
    return client;
}

// ********** Collections **********

async function initCollections() {
  const client = getClient();
  for (const name of Object.values(COLLECTIONS)) {
    const exists = await collectionExists(name);
    if (!exists) {
      await client.createCollection(name, {
        vectors: {          
            size: QDRANT.VECTOR_SIZE,
            distance: QDRANT.DISTANCE_METRIC
        }
      });
      console.log(`Created Qdrant collection: ${name}`);
    } else {
      console.log(`Qdrant collection already exists: ${name}`);
    }
  }
}

async function collectionExists(name) {
    try {
        const client = getClient();
        await client.getCollection(name);
        return true;
    } catch (err) {
        return false;
    }
}

// ********* Upsert Vectors *********

//payload is metadata stored alongside the vector in QDrant.
//SQLite ID stored here to be able to link back to original data
async function upsertVector(collection, id, vector, payload = {}) {
    const client = getClient();
    await client.upsertVector(collection, {
        wait: true,  // Wait for the operation to complete before returning
        points: [{id, vector, payload}]
    });
}

// Upsert an episode vector into the EPISODES collection, with the episode ID as the point ID
async function upsertEpisode(id, vector, payload={}) {
    return upsertVector(COLLECTIONS.EPISODES, id, vector, payload);
}

async function upsertEntity(id, vector, payload={}) {
    return upsertVector(COLLECTIONS.ENTITIES, id, vector, payload);
}

async function upsertSummary(id, vector, payload={}) {
    return upsertVector(COLLECTIONS.SUMMARIES, id, vector, payload);
}

//********** Search Vectors ********** */
async function searchCollection(collection, vector, limit = QDRANT.DEFAULT_LIMIT, filter = null){
    const client = getClient();
    const params = {vector,limit, with_payload: true};

    if (filter) params.filter = filter;

    return client.search(collection, params);
}
async function searchEpisodes(vector, limit = QDRANT.DEFAULT_LIMIT, filter = null) {
    return searchCollection(COLLECTIONS.EPISODES, vector, limit, filter);
}
async function searchEntities(vector, limit = QDRANT.DEFAULT_LIMIT, filter = null) {
    return searchCollection(COLLECTIONS.ENTITIES, vector, limit, filter);
}
async function searchSummaries(vector, limit = QDRANT.DEFAULT_LIMIT, filter = null) {
    return searchCollection(COLLECTIONS.SUMMARIES, vector, limit, filter);
}

//********** Delete Vectors ********** */

async function deleteVector(collection, id) {
    const client = getClient();
    await client.delete(collection, {
        wait: true,
        points: [id]
    });
}

module.exports = {
    initCollections,
    upsertEpisode,
    upsertEntity,
    upsertSummary,
    searchEpisodes,
    searchEntities,
    searchSummaries,
    deleteVector
};