import { QdrantClient } from '@qdrant/js-client-rest';
import { qdrantAuth } from '../..';
import { createAction, Property } from '@activepieces/pieces-framework';
import { randomUUID } from 'crypto';
import { collectionName, decodeEmbeddings } from '../common';

export const addPointsToCollection = createAction({
  auth: qdrantAuth,
  requireAuth: true,
  name: 'add_points_to_collection',
  displayName: 'Add points to collection',
  description:
    'Insert a point (= embedding or vector + other infos) to a specific collection, if the collection does not exist it will be created',
  audience: 'both',
  aiMetadata: {
    description:
      'Upsert one or more points (vectors plus optional payload) into a Qdrant collection, auto-creating the collection with the chosen distance metric if it does not yet exist. Use to load embeddings into a vector store. Embeddings are passed as a file and must all share the same dimensionality. Not idempotent: when point ids are omitted they are randomly generated, so each call appends new points; supply stable ids to overwrite the same points instead.',
    idempotent: false,
  },
  props: {
    collectionName,
    embeddings: Property.File({
      displayName: 'Embeddings',
      description: 'Embeddings (= vectors) for the points',
      required: true,
    }),
    embeddingsIds: Property.Array({
      displayName: 'Embeddings Ids',
      description:
        'The ids of the embeddings for the points. If not provided, the ids will be generated automatically',
      required: false,
    }),
    distance: Property.StaticDropdown({
      displayName: 'Calculation Method of distance',
      description:
        "The calculation method helps to rank vectors when you want to find the closest points, the method to use depends on the model who's created the embeddings, see the documentation of your model",
      defaultValue: 'Cosine',
      options: {
        options: [
          { label: 'Cosine', value: 'Cosine' },
          { label: 'Euclidean', value: 'Euclid' },
          { label: 'Dot', value: 'Dot' },
        ],
      },
      required: true,
    }),
    payload: Property.Json({
      displayName: 'Additional Payload',
      description: `Please follow [payload documentation](https://qdrant.tech/documentation/concepts/payload/) to add additional information to the points.`,
      required: false,
    }),
    storage: Property.StaticDropdown({
      displayName: 'Storage',
      description: 'Define where points will be stored',
      options: {
        options: [
          { label: 'on Disk', value: 'Disk' },
          { label: 'On Memory', value: 'Memory' },
        ],
      },
      defaultValue: 'Disk',
      required: false,
    }),
  },
  run: async ({ auth, propsValue }) => {
    const client = new QdrantClient({
      apiKey: auth.props.key,
      url: auth.props.serverAddress,
    });
    const embeddings = decodeEmbeddings(propsValue.embeddings.data);

    const numberOfEmbeddings = embeddings.length;
    const embeddingsLen = embeddings[0].length;

    for (const embedding of embeddings) {
      if (embedding.length != embeddingsLen)
        throw new Error(
          'Embeddings must have the same length (=number of dimensions)'
        );
    }

    const embeddingsIds = (propsValue.embeddingsIds as string[]) ?? [];

    const autoEmbeddingsIds = embeddingsIds.length === 0;

    if (!autoEmbeddingsIds && embeddingsIds.length !== numberOfEmbeddings)
      throw new Error(
        'The number of embeddings Ids and the number of embeddings must be the same'
      );

    const payload = propsValue.payload ?? {};
    const points = [];

    for (let i = 0; i < numberOfEmbeddings; i++) {
      const localPayload = { ...payload };
      points.push({
        id: autoEmbeddingsIds ? randomUUID() : embeddingsIds[i],
        payload: localPayload,
        vector: Array.from(embeddings[i]),
      });
    }


    const collections = (await client.getCollections()).collections;
    const collectionName = propsValue.collectionName as string;
    if (!collections.find((collection) => collection.name === collectionName)) {
      await client.createCollection(collectionName, {
        vectors: {
          size: embeddingsLen,
          distance: propsValue.distance as 'Dot' | 'Cosine' | 'Euclid',
          on_disk: propsValue.storage === 'Disk',
        },
        on_disk_payload: propsValue.storage === 'Disk',
      });
    }

    const response = await client.upsert(collectionName, {
      points,
      wait: true,
    });

    return response;
  },
});
