import { IEmbeddingFunction } from "./IEmbeddingFunction";
import { validateConfigSchema } from "../schemas/schemaUtils";
// Dynamically import module
let TransformersApi: Promise<any>;

type StoredConfig = {
  model: string;
  revision: string;
  quantized: boolean;
};

export class TransformersEmbeddingFunction implements IEmbeddingFunction {
  name = "transformers";

  private pipelinePromise?: Promise<any> | null;
  private transformersApi: any;
  private model: string;
  private revision: string;
  private quantized: boolean;
  private progress_callback: Function | null;

  /**
   * TransformersEmbeddingFunction constructor.
   * @param options The configuration options.
   * @param options.model The model to use to calculate embeddings. Defaults to 'Xenova/all-MiniLM-L6-v2', which is an ONNX port of `sentence-transformers/all-MiniLM-L6-v2`.
   * @param options.revision The specific model version to use (can be a branch, tag name, or commit id). Defaults to 'main'.
   * @param options.quantized Whether to load the 8-bit quantized version of the model. Defaults to `false`.
   * @param options.progress_callback If specified, this function will be called during model construction, to provide the user with progress updates.
   */
  constructor({
    model = "Xenova/all-MiniLM-L6-v2",
    revision = "main",
    quantized = false,
    progress_callback = null,
  }: {
    model?: string;
    revision?: string;
    quantized?: boolean;
    progress_callback?: Function | null;
  } = {}) {
    this.model = model;
    this.revision = revision;
    this.quantized = quantized;
    this.progress_callback = progress_callback;
  }

  public async generate(texts: string[]): Promise<number[][]> {
    await this.loadClient();

    // Store a promise that resolves to the pipeline
    this.pipelinePromise = new Promise(async (resolve, reject) => {
      try {
        const pipeline = this.transformersApi;

        const quantized = this.quantized;
        const revision = this.revision;
        const progress_callback = this.progress_callback;

        resolve(
          await pipeline("feature-extraction", this.model, {
            quantized,
            revision,
            progress_callback,
          }),
        );
      } catch (e) {
        reject(e);
      }
    });

    let pipe = await this.pipelinePromise;
    let output = await pipe(texts, { pooling: "mean", normalize: true });
    return output.tolist();
  }

  private async loadClient() {
    if (this.transformersApi) return;
    try {
      // eslint-disable-next-line global-require,import/no-extraneous-dependencies
      let { pipeline } = await TransformersEmbeddingFunction.import();
      TransformersApi = pipeline;
    } catch (_a) {
      // @ts-ignore
      if (_a.code === "MODULE_NOT_FOUND") {
        throw new Error(
          "Please install the @xenova/transformers package to use the TransformersEmbeddingFunction, `npm install @xenova/transformers`",
        );
      }
      throw _a; // Re-throw other errors
    }
    this.transformersApi = TransformersApi;
  }

  /** @ignore */
  static async import(): Promise<{
    // @ts-ignore
    pipeline: typeof import("@xenova/transformers");
  }> {
    try {
      // @ts-ignore
      const { pipeline } = await import("@xenova/transformers");
      return { pipeline };
    } catch (e) {
      throw new Error(
        "Please install @xenova/transformers as a dependency with, e.g. `npm install @xenova/transformers`",
      );
    }
  }

  buildFromConfig(config: StoredConfig): TransformersEmbeddingFunction {
    return new TransformersEmbeddingFunction({
      model: config.model,
      revision: config.revision,
      quantized: config.quantized,
    });
  }

  getConfig(): StoredConfig {
    return {
      model: this.model,
      revision: this.revision,
      quantized: this.quantized,
    };
  }

  validateConfigUpdate(oldConfig: StoredConfig, newConfig: StoredConfig): void {
    if (oldConfig.model !== newConfig.model) {
      throw new Error("Cannot change the model of the embedding function.");
    }
    if (oldConfig.revision !== newConfig.revision) {
      throw new Error("Cannot change the revision of the embedding function.");
    }
    if (oldConfig.quantized !== newConfig.quantized) {
      throw new Error(
        "Cannot change the quantization of the embedding function.",
      );
    }
  }

  validateConfig(config: StoredConfig): void {
    validateConfigSchema(config, "transformers");
  }
}
