/**
 * Copyright 2025 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import {
  ActionMetadata,
  EmbedderAction,
  embedderActionMetadata,
  EmbedderInfo,
  EmbedderReference,
  z,
} from 'genkit';
import { embedderRef } from 'genkit/embedder';
import { embedder as pluginEmbedder } from 'genkit/plugin';
import { toGeminiMessage } from '../common/converters.js';
import { isKnownKey } from '../common/utils.js';
import { embedContent } from './client.js';
import {
  ClientOptions,
  EmbedContentRequest,
  GoogleAIPluginOptions,
  Model,
  TaskTypeSchema,
} from './types.js';
import {
  calculateApiKey,
  checkApiKey,
  checkModelName,
  extractVersion,
} from './utils.js';

export const EmbeddingConfigSchema = z
  .object({
    /** Override the API key provided at plugin initialization. */
    apiKey: z.string().optional(),
    /**
     * The `task_type` parameter is defined as the intended downstream application to help the model
     * produce better quality embeddings.
     **/
    taskType: TaskTypeSchema.optional(),
    title: z.string().optional(),
    version: z.string().optional(),
    /**
     * The `outputDimensionality` parameter allows you to specify the dimensionality of the embedding output.
     * By default, `gemini-embedding-2` and `gemini-embedding-001` generate embeddings with 3072 dimensions.
     * By selecting a smaller output dimensionality, users can save memory and storage space, leading to more efficient computations.
     **/
    outputDimensionality: z.number().min(1).optional(),
  })
  .passthrough();
export type EmbeddingConfigSchemaType = typeof EmbeddingConfigSchema;
export type EmbeddingConfig = z.infer<EmbeddingConfigSchemaType>;

// This contains all the embedder config schema types
type ConfigSchemaType = EmbeddingConfigSchemaType;

function commonRef(
  name: string,
  info?: EmbedderInfo,
  configSchema: ConfigSchemaType = EmbeddingConfigSchema
): EmbedderReference<ConfigSchemaType> {
  return embedderRef({
    name: `googleai/${name}`,
    configSchema,
    info: info ?? {
      dimensions: 3072,
      supports: {
        input: ['text'],
      },
    },
  });
}

const GENERIC_MODEL = commonRef('embedder');

// Models we no longer suggest, but that still work: they still resolve and
// are still registered as actions, but aren't in the KnownModels type, so
// they don't show up in autocomplete. Use this for models that are deprecated
// or < 1 month from shutdown (note the shutdown date).
const HIDDEN_MODELS = {} as const;

const KNOWN_MODELS = {
  'gemini-embedding-001': commonRef('gemini-embedding-001'),
  'gemini-embedding-2': commonRef('gemini-embedding-2', {
    supports: {
      input: ['text', 'image', 'video'],
    },
  }),
} as const;

const ALL_MODELS = {
  ...HIDDEN_MODELS,
  ...KNOWN_MODELS,
};
export type KnownModels = keyof typeof KNOWN_MODELS; // For autocomplete

export type EmbedderModelName = `gemini-embedding-${string}`;
export function isEmbedderName(value: string): value is EmbedderModelName {
  return value.startsWith('gemini-embedding-');
}

export function model(
  version: string,
  config: EmbeddingConfig = {}
): EmbedderReference<ConfigSchemaType> {
  const name = checkModelName(version);

  if (isKnownKey(name, ALL_MODELS)) {
    const known = ALL_MODELS[name];
    return embedderRef({
      name: known.name,
      info: known.info,
      configSchema: known.configSchema,
      config,
    });
  }

  return embedderRef({
    name: `googleai/${name}`,
    info: { ...GENERIC_MODEL.info },
    configSchema: GENERIC_MODEL.configSchema,
    config,
  });
}

export function listActions(models: Model[]): ActionMetadata[] {
  return (
    models
      .filter((m) => m.supportedGenerationMethods.includes('embedContent'))
      // Filter out deprecated
      .filter((m) => !m.description || !m.description.includes('deprecated'))
      .map((m) => {
        const ref = model(m.name);
        return embedderActionMetadata({
          name: ref.name,
          info: ref.info,
          configSchema: ref.configSchema,
        });
      })
  );
}

export function listKnownModels(options?: GoogleAIPluginOptions) {
  return Object.keys(ALL_MODELS).map((name) => defineEmbedder(name, options));
}

export function defineEmbedder(
  name: string,
  pluginOptions?: GoogleAIPluginOptions
): EmbedderAction {
  checkApiKey(pluginOptions?.apiKey);
  const ref = model(name);
  const clientOptions: ClientOptions = {
    apiVersion: pluginOptions?.apiVersion,
    baseUrl: pluginOptions?.baseUrl,
    customHeaders: pluginOptions?.customHeaders,
    experimental_debugTraces: pluginOptions?.experimental_debugTraces,
  };

  return pluginEmbedder(
    {
      name: ref.name,
      configSchema: ref.configSchema,
      info: ref.info,
    },
    async (request, _) => {
      const embedApiKey = calculateApiKey(
        pluginOptions?.apiKey,
        request.options?.apiKey
      );
      const embedVersion = request.options?.version || extractVersion(ref);

      const embeddings = await Promise.all(
        request.input.map(async (doc) => {
          const response = await embedContent(
            embedApiKey,
            embedVersion,
            {
              taskType: request.options?.taskType,
              title: request.options?.title,
              content: toGeminiMessage({ role: 'user', content: doc.content }),
              outputDimensionality: request.options?.outputDimensionality,
            } as EmbedContentRequest,
            clientOptions
          );
          const values = response.embedding.values;
          return { embedding: values };
        })
      );
      return { embeddings };
    }
  );
}

export const TEST_ONLY = { KNOWN_MODELS };