import { z, ActionMetadata, MediaPart } from 'genkit';
import { ModelReference, ModelAction, GenerateRequest } from 'genkit/model';
import { Model, GoogleAIPluginOptions } from './types.js';
import { ImagenParameters, ImagenPrediction } from '../common/types.js';

/**
 * Copyright 2024 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.
 */

/**
 * See https://ai.google.dev/gemini-api/docs/image-generation#imagen-model
 */
declare const ImagenConfigSchema: z.ZodObject<{
    apiKey: z.ZodOptional<z.ZodString>;
    numberOfImages: z.ZodOptional<z.ZodNumber>;
    aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
    personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
}, "passthrough", z.ZodTypeAny, z.objectOutputType<{
    apiKey: z.ZodOptional<z.ZodString>;
    numberOfImages: z.ZodOptional<z.ZodNumber>;
    aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
    personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
}, z.ZodTypeAny, "passthrough">, z.objectInputType<{
    apiKey: z.ZodOptional<z.ZodString>;
    numberOfImages: z.ZodOptional<z.ZodNumber>;
    aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
    personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
}, z.ZodTypeAny, "passthrough">>;
type ImagenConfigSchemaType = typeof ImagenConfigSchema;
type ImagenConfig = z.infer<ImagenConfigSchemaType>;
type ConfigSchemaType = ImagenConfigSchemaType;
declare const KNOWN_MODELS: {
    readonly 'imagen-4.0-fast-generate-001': ModelReference<z.ZodObject<{
        apiKey: z.ZodOptional<z.ZodString>;
        numberOfImages: z.ZodOptional<z.ZodNumber>;
        aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
        personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
    }, "passthrough", z.ZodTypeAny, z.objectOutputType<{
        apiKey: z.ZodOptional<z.ZodString>;
        numberOfImages: z.ZodOptional<z.ZodNumber>;
        aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
        personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
    }, z.ZodTypeAny, "passthrough">, z.objectInputType<{
        apiKey: z.ZodOptional<z.ZodString>;
        numberOfImages: z.ZodOptional<z.ZodNumber>;
        aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
        personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
    }, z.ZodTypeAny, "passthrough">>>;
    readonly 'imagen-4.0-generate-001': ModelReference<z.ZodObject<{
        apiKey: z.ZodOptional<z.ZodString>;
        numberOfImages: z.ZodOptional<z.ZodNumber>;
        aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
        personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
    }, "passthrough", z.ZodTypeAny, z.objectOutputType<{
        apiKey: z.ZodOptional<z.ZodString>;
        numberOfImages: z.ZodOptional<z.ZodNumber>;
        aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
        personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
    }, z.ZodTypeAny, "passthrough">, z.objectInputType<{
        apiKey: z.ZodOptional<z.ZodString>;
        numberOfImages: z.ZodOptional<z.ZodNumber>;
        aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
        personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
    }, z.ZodTypeAny, "passthrough">>>;
    readonly 'imagen-4.0-ultra-generate-001': ModelReference<z.ZodObject<{
        apiKey: z.ZodOptional<z.ZodString>;
        numberOfImages: z.ZodOptional<z.ZodNumber>;
        aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
        personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
    }, "passthrough", z.ZodTypeAny, z.objectOutputType<{
        apiKey: z.ZodOptional<z.ZodString>;
        numberOfImages: z.ZodOptional<z.ZodNumber>;
        aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
        personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
    }, z.ZodTypeAny, "passthrough">, z.objectInputType<{
        apiKey: z.ZodOptional<z.ZodString>;
        numberOfImages: z.ZodOptional<z.ZodNumber>;
        aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
        personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
    }, z.ZodTypeAny, "passthrough">>>;
};
type KnownModels = keyof typeof KNOWN_MODELS;
type ImagenModelName = `imagen-${string}`;
declare function isImagenModelName(value?: string): value is ImagenModelName;
declare function model(version: string, config?: ImagenConfig): ModelReference<ConfigSchemaType>;
declare function listActions(models: Model[]): ActionMetadata[];
declare function listKnownModels(options?: GoogleAIPluginOptions): ModelAction<z.ZodTypeAny>[];
declare function defineModel(name: string, pluginOptions?: GoogleAIPluginOptions): ModelAction;
declare function fromImagenPrediction(p: ImagenPrediction): MediaPart;
declare function toImagenParameters(request: GenerateRequest<typeof ImagenConfigSchema>): ImagenParameters;
declare const TEST_ONLY: {
    toImagenParameters: typeof toImagenParameters;
    fromImagenPrediction: typeof fromImagenPrediction;
    GENERIC_MODEL: ModelReference<z.ZodObject<{
        apiKey: z.ZodOptional<z.ZodString>;
        numberOfImages: z.ZodOptional<z.ZodNumber>;
        aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
        personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
    }, "passthrough", z.ZodTypeAny, z.objectOutputType<{
        apiKey: z.ZodOptional<z.ZodString>;
        numberOfImages: z.ZodOptional<z.ZodNumber>;
        aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
        personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
    }, z.ZodTypeAny, "passthrough">, z.objectInputType<{
        apiKey: z.ZodOptional<z.ZodString>;
        numberOfImages: z.ZodOptional<z.ZodNumber>;
        aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
        personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
    }, z.ZodTypeAny, "passthrough">>>;
    KNOWN_MODELS: {
        readonly 'imagen-4.0-fast-generate-001': ModelReference<z.ZodObject<{
            apiKey: z.ZodOptional<z.ZodString>;
            numberOfImages: z.ZodOptional<z.ZodNumber>;
            aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
            personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
        }, "passthrough", z.ZodTypeAny, z.objectOutputType<{
            apiKey: z.ZodOptional<z.ZodString>;
            numberOfImages: z.ZodOptional<z.ZodNumber>;
            aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
            personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
        }, z.ZodTypeAny, "passthrough">, z.objectInputType<{
            apiKey: z.ZodOptional<z.ZodString>;
            numberOfImages: z.ZodOptional<z.ZodNumber>;
            aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
            personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
        }, z.ZodTypeAny, "passthrough">>>;
        readonly 'imagen-4.0-generate-001': ModelReference<z.ZodObject<{
            apiKey: z.ZodOptional<z.ZodString>;
            numberOfImages: z.ZodOptional<z.ZodNumber>;
            aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
            personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
        }, "passthrough", z.ZodTypeAny, z.objectOutputType<{
            apiKey: z.ZodOptional<z.ZodString>;
            numberOfImages: z.ZodOptional<z.ZodNumber>;
            aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
            personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
        }, z.ZodTypeAny, "passthrough">, z.objectInputType<{
            apiKey: z.ZodOptional<z.ZodString>;
            numberOfImages: z.ZodOptional<z.ZodNumber>;
            aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
            personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
        }, z.ZodTypeAny, "passthrough">>>;
        readonly 'imagen-4.0-ultra-generate-001': ModelReference<z.ZodObject<{
            apiKey: z.ZodOptional<z.ZodString>;
            numberOfImages: z.ZodOptional<z.ZodNumber>;
            aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
            personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
        }, "passthrough", z.ZodTypeAny, z.objectOutputType<{
            apiKey: z.ZodOptional<z.ZodString>;
            numberOfImages: z.ZodOptional<z.ZodNumber>;
            aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
            personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
        }, z.ZodTypeAny, "passthrough">, z.objectInputType<{
            apiKey: z.ZodOptional<z.ZodString>;
            numberOfImages: z.ZodOptional<z.ZodNumber>;
            aspectRatio: z.ZodOptional<z.ZodEnum<["1:1", "9:16", "16:9", "3:4", "4:3"]>>;
            personGeneration: z.ZodOptional<z.ZodEnum<["dont_allow", "allow_adult", "allow_all"]>>;
        }, z.ZodTypeAny, "passthrough">>>;
    };
};

export { type ImagenConfig, ImagenConfigSchema, type ImagenConfigSchemaType, type ImagenModelName, type KnownModels, TEST_ONLY, defineModel, isImagenModelName, listActions, listKnownModels, model };