import {
  lazySchema,
  zodSchema,
  type InferSchema,
} from '@ai-sdk/provider-utils';
import { z } from 'zod/v4';
import type { GoogleVertexModelIds } from './google-vertex-model-catalog';

// https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api
export type GoogleVertexEmbeddingModelId =
  | GoogleVertexModelIds['embedding']
  | (string & {});

export const googleVertexEmbeddingModelOptions = lazySchema(() =>
  zodSchema(
    z.object({
      /**
       * Optional. Optional reduced dimension for the output embedding.
       * If set, excessive values in the output embedding are truncated from the end.
       */
      outputDimensionality: z.number().optional(),

      /**
       * Optional. Specifies the task type for generating embeddings.
       * Supported task types:
       * - SEMANTIC_SIMILARITY: Optimized for text similarity.
       * - CLASSIFICATION: Optimized for text classification.
       * - CLUSTERING: Optimized for clustering texts based on similarity.
       * - RETRIEVAL_DOCUMENT: Optimized for document retrieval.
       * - RETRIEVAL_QUERY: Optimized for query-based retrieval.
       * - QUESTION_ANSWERING: Optimized for answering questions.
       * - FACT_VERIFICATION: Optimized for verifying factual information.
       * - CODE_RETRIEVAL_QUERY: Optimized for retrieving code blocks based on natural language queries.
       */
      taskType: z
        .enum([
          'SEMANTIC_SIMILARITY',
          'CLASSIFICATION',
          'CLUSTERING',
          'RETRIEVAL_DOCUMENT',
          'RETRIEVAL_QUERY',
          'QUESTION_ANSWERING',
          'FACT_VERIFICATION',
          'CODE_RETRIEVAL_QUERY',
        ])
        .optional(),

      /**
       * Optional. The title of the document being embedded.
       * Only valid when task_type is set to 'RETRIEVAL_DOCUMENT'.
       * Helps the model produce better embeddings by providing additional context.
       */
      title: z.string().optional(),

      /**
       * Optional. When set to true, input text will be truncated. When set to false,
       * an error is returned if the input text is longer than the maximum length supported by the model. Defaults to true.
       */
      autoTruncate: z.boolean().optional(),
    }),
  ),
);

export type GoogleVertexEmbeddingModelOptions = InferSchema<
  typeof googleVertexEmbeddingModelOptions
>;