import { createAction, Property } from '@activepieces/pieces-framework';
import { InvokeModelCommand } from '@aws-sdk/client-bedrock-runtime';
import { ModelModality } from '@aws-sdk/client-bedrock';
import { awsBedrockAuth } from '../auth';
import {
  createBedrockRuntimeClient,
  getBedrockModelOptions,
  formatBedrockError,
} from '../common';

export const generateEmbeddings = createAction({
  audience: 'human',
  auth: awsBedrockAuth,
  name: 'generate_embeddings',
  displayName: 'Generate Embeddings',
  description:
    'Generate vector embeddings from text using Amazon Titan Embed, Cohere Embed, or Amazon Nova Multimodal Embeddings models.',
  props: {
    model: Property.Dropdown({
      displayName: 'Model',
      required: true,
      auth: awsBedrockAuth,
      description: 'The embedding model to use.',
      refreshers: [],
      options: async ({ auth }) => {
        if (!auth) {
          return {
            disabled: true,
            placeholder: 'Connect your AWS account first',
            options: [],
          };
        }
        return getBedrockModelOptions(auth.props, {
          outputModality: ModelModality.EMBEDDING,
        });
      },
    }),
    inputText: Property.LongText({
      displayName: 'Input Text',
      required: true,
      description: 'The text to generate embeddings for.',
    }),
    embeddingPurpose: Property.StaticDropdown({
      displayName: 'Embedding Purpose',
      required: false,
      description:
        'How the embeddings will be used — helps the model optimize the output. Only applies to Nova Multimodal Embeddings models. Leave as "Generic Index" if unsure.',
      defaultValue: 'GENERIC_INDEX',
      options: {
        options: [
          { label: 'Generic Index', value: 'GENERIC_INDEX' },
          { label: 'Generic Retrieval', value: 'GENERIC_RETRIEVAL' },
          { label: 'Text Retrieval', value: 'TEXT_RETRIEVAL' },
          { label: 'Classification', value: 'CLASSIFICATION' },
          { label: 'Clustering', value: 'CLUSTERING' },
        ],
      },
    }),
    dimensions: Property.Number({
      displayName: 'Vector Dimensions',
      required: false,
      description:
        'Size of the output vector. Larger = more accurate but slower/costlier. Titan Embed v2 supports: 256, 512, 1024. Nova Multimodal supports: 256, 384, 1024, 3072. Leave blank to use the model default.',
    }),
    normalize: Property.Checkbox({
      displayName: 'Normalize Output',
      required: false,
      description:
        'Scale the output vector so its values range between -1 and 1. Recommended for similarity comparisons. Only supported by Titan Embed v2.',
      defaultValue: true,
    }),
  },
  async run({ auth, propsValue }) {
    const client = createBedrockRuntimeClient(auth.props);
    const { model, inputText, dimensions, normalize } = propsValue;

    const isTitan = model.startsWith('amazon.titan-embed');
    const isCohere = model.startsWith('cohere.embed');
    const isNova = model.includes('nova') && model.includes('embed');
    const isTwelveLabs = model.startsWith('twelvelabs.');

    let requestBody: Record<string, unknown>;

    if (isNova) {
      requestBody = {
        schemaVersion: 'nova-multimodal-embed-v1',
        taskType: 'SINGLE_EMBEDDING',
        singleEmbeddingParams: {
          embeddingPurpose: propsValue.embeddingPurpose ?? 'GENERIC_INDEX',
          ...(dimensions != null ? { embeddingDimension: dimensions } : {}),
          text: {
            truncationMode: 'END',
            value: inputText,
          },
        },
      };
    } else if (isTitan) {
      requestBody = {
        inputText,
        ...(dimensions != null ? { dimensions } : {}),
        ...(normalize != null ? { normalize } : {}),
      };
    } else if (isCohere) {
      requestBody = {
        texts: [inputText],
        input_type: 'search_document',
        truncate: 'END',
      };
    } else if (isTwelveLabs) {
      requestBody = {
        inputType: 'text',
        inputText,
        textTruncate: 'end',
      };
    } else {
      requestBody = {
        inputText,
      };
    }

    try {
      const response = await client.send(
        new InvokeModelCommand({
          modelId: model,
          body: Buffer.from(JSON.stringify(requestBody)),
          contentType: 'application/json',
          accept: 'application/json',
        })
      );

      const responseBody = JSON.parse(
        new TextDecoder().decode(response.body)
      );

      let embedding: number[] | undefined;

      if (responseBody.embedding) {
        // Titan Embed / TwelveLabs format: { embedding: [floats] }
        embedding = responseBody.embedding;
      } else if (
        responseBody.embeddings &&
        responseBody.embeddings.length > 0
      ) {
        const first = responseBody.embeddings[0];
        if (Array.isArray(first)) {
          // Cohere v3 format: { embeddings: [[floats]] }
          embedding = first;
        } else if (first?.embedding) {
          // Nova format: { embeddings: [{ embedding: [floats] }] }
          embedding = first.embedding;
        }
      }

      if (!embedding) {
        throw new Error(
          'No embedding was returned by the model. The response format may not be supported.'
        );
      }

      return {
        embedding,
        dimensions: embedding.length,
        model,
      };
    } catch (error) {
      if (error instanceof Error && error.message.includes('No embedding')) {
        throw error;
      }
      throw new Error(formatBedrockError(error));
    }
  },
});
