/**
 * 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 { GenkitError, z, type ActionContext } from 'genkit';
import { GoogleAuth } from 'google-auth-library';
import {
  isMultiRegionalLocation,
  type ClientOptions,
  type ExpressClientOptions,
  type GlobalClientOptions,
  type MultiRegionalClientOptions,
  type RegionalClientOptions,
  type VertexPluginOptions,
} from './types.js';

export {
  checkModelName,
  checkSupportedMimeType,
  cleanSchema,
  extractMedia,
  extractMimeType,
  extractText,
  extractVersion,
  modelName,
} from '../common/utils.js';

const CLOUD_PLATFORM_OAUTH_SCOPE =
  'https://www.googleapis.com/auth/cloud-platform';

function parseFirebaseProjectId(): string | undefined {
  if (!process.env.FIREBASE_CONFIG) return undefined;
  try {
    return JSON.parse(process.env.FIREBASE_CONFIG).projectId as string;
  } catch {
    return undefined;
  }
}

export function shouldUseLegacyEndpoint(model: string) {
  // This only covers lyria models so far
  return model === 'lyria-002';
}

let __mockDerivedOptions: ClientOptions | undefined = undefined;
function setMockDerivedOptions(options: ClientOptions | undefined): void {
  __mockDerivedOptions = options;
}
export const TEST_ONLY = { setMockDerivedOptions };

export async function getDerivedOptions(
  options?: VertexPluginOptions,
  AuthClass: typeof GoogleAuth = GoogleAuth // Injectable testing
): Promise<ClientOptions> {
  if (__mockDerivedOptions) {
    return Promise.resolve(__mockDerivedOptions);
  }

  // Lazy instantiate the authClient
  let authClientPromise: Promise<GoogleAuth> | undefined;
  const getAuthClientInternal = (): Promise<GoogleAuth> => {
    if (!authClientPromise) {
      authClientPromise = getAuthClient(AuthClass, options);
    }
    return authClientPromise;
  };

  // Lazy instantiate the projectId
  let projectIdPromise: Promise<string> | undefined;
  const getProjectIdInternal = (): Promise<string> => {
    if (!projectIdPromise) {
      projectIdPromise = (async () => {
        const authClient = await getAuthClientInternal();
        return getProjectId(authClient, options);
      })();
    }
    return projectIdPromise;
  };

  // Figure out the type of preferred options if possible
  // The order of the if statements is important.
  if (options?.location == 'global') {
    return await getGlobalDerivedOptions(
      await getAuthClientInternal(),
      await getProjectIdInternal(),
      options
    );
  } else if (isMultiRegionalLocation(options?.location)) {
    return await getMultiRegionalDerivedOptions(
      await getAuthClientInternal(),
      await getProjectIdInternal(),
      options
    );
  } else if (options?.location) {
    return await getRegionalDerivedOptions(
      await getAuthClientInternal(),
      await getProjectIdInternal(),
      options
    );
  } else if (options?.apiKey !== undefined) {
    // apiKey = false still indicates apiKey expectation
    return getExpressDerivedOptions(options);
  }

  // If we got here then we're relying on environment variables.
  // Try multi-regional first, it's very narrowly defined location
  try {
    const multiRegionalOptions = await getMultiRegionalDerivedOptions(
      await getAuthClientInternal(),
      await getProjectIdInternal(),
      options
    );
    return multiRegionalOptions;
  } catch (e: unknown) {
    /* no-op - try regional next */
  }
  // Try regional second, it's the most common usage.
  try {
    const regionalOptions = await getRegionalDerivedOptions(
      await getAuthClientInternal(),
      await getProjectIdInternal(),
      options
    );
    return regionalOptions;
  } catch (e: unknown) {
    /* no-op - try global next */
  }
  try {
    const globalOptions = await getGlobalDerivedOptions(
      await getAuthClientInternal(),
      await getProjectIdInternal(),
      options
    );
    return globalOptions;
  } catch (e: unknown) {
    /* no-op - try express last */
  }
  try {
    const expressOptions = getExpressDerivedOptions(options);
    return expressOptions;
  } catch (e: unknown) {
    /* no-op */
  }

  // We did not have enough information in the options or in environment variables
  // to properly determine client options.
  throw new GenkitError({
    status: 'INVALID_ARGUMENT',
    message:
      'Unable to determine client options. Please set either apiKey or projectId and location',
  });
}

async function getGlobalDerivedOptions(
  authClient: GoogleAuth,
  projectId: string,
  options?: VertexPluginOptions
): Promise<GlobalClientOptions> {
  const clientOpt: GlobalClientOptions = {
    kind: 'global',
    location: 'global',
    projectId,
    authClient,
    experimental_debugTraces: options?.experimental_debugTraces,
  };
  if (options?.apiKey) {
    clientOpt.apiKey = options.apiKey;
  }
  if (options?.apiVersion) {
    clientOpt.apiVersion = options.apiVersion;
  }

  return clientOpt;
}

function getExpressDerivedOptions(
  options?: VertexPluginOptions
): ExpressClientOptions {
  const apiKey = checkApiKey(options?.apiKey);
  const clientOpt: ExpressClientOptions = {
    kind: 'express',
    apiKey,
    experimental_debugTraces: options?.experimental_debugTraces,
  };
  if (options?.apiVersion) {
    clientOpt.apiVersion = options.apiVersion;
  }
  return clientOpt;
}

async function getAuthClient(
  AuthClass: typeof GoogleAuth,
  options?: VertexPluginOptions
): Promise<GoogleAuth> {
  let authOptions = options?.googleAuth;
  let authClient: GoogleAuth;
  const providedProjectId =
    options?.projectId ||
    process.env.GCLOUD_PROJECT ||
    parseFirebaseProjectId();
  if (process.env.GCLOUD_SERVICE_ACCOUNT_CREDS) {
    const serviceAccountCreds = JSON.parse(
      process.env.GCLOUD_SERVICE_ACCOUNT_CREDS
    );
    authOptions = {
      credentials: serviceAccountCreds,
      scopes: [CLOUD_PLATFORM_OAUTH_SCOPE],
      projectId: providedProjectId,
    };
    authClient = new AuthClass(authOptions);
  } else {
    authClient = new AuthClass(
      authOptions ?? {
        scopes: [CLOUD_PLATFORM_OAUTH_SCOPE],
        projectId: providedProjectId,
      }
    );
  }
  return authClient;
}

async function getProjectId(
  authClient: GoogleAuth,
  options?: VertexPluginOptions
): Promise<string> {
  const projectId =
    options?.projectId ||
    process.env.GCLOUD_PROJECT ||
    (await authClient.getProjectId());
  if (!projectId) {
    throw new Error(
      `VertexAI Plugin is missing the 'project' configuration. Please set the 'GCLOUD_PROJECT' environment variable or explicitly pass 'project' into genkit config.`
    );
  }
  return projectId;
}

async function getMultiRegionalDerivedOptions(
  authClient: GoogleAuth,
  projectId: string,
  options?: VertexPluginOptions
): Promise<MultiRegionalClientOptions> {
  const location = options?.location || process.env.GCLOUD_LOCATION;

  if (!location) {
    throw new Error(
      `VertexAI Plugin is missing the 'location' configuration. Please set the 'GCLOUD_LOCATION' environment variable or explicitly pass 'location' into genkit config.`
    );
  }
  if (!isMultiRegionalLocation(location)) {
    throw new Error(
      `Tried to set multi regional location to an invalid region: ${location}`
    );
  }

  const clientOpt: MultiRegionalClientOptions = {
    kind: 'multi-regional',
    location,
    projectId,
    authClient,
    experimental_debugTraces: options?.experimental_debugTraces,
  };
  if (options?.apiKey) {
    clientOpt.apiKey = options.apiKey;
  }
  if (options?.apiVersion) {
    clientOpt.apiVersion = options.apiVersion;
  }
  return clientOpt;
}

async function getRegionalDerivedOptions(
  authClient: GoogleAuth,
  projectId: string,
  options?: VertexPluginOptions
): Promise<RegionalClientOptions> {
  // TODO(V2): Default to 'global' for every model. For now only models on
  // the Interactions API (only on global and the us/eu multi-regions) use
  // global when the location is defaulted (see `locationDefaulted`).
  const chosenLocation = options?.location || process.env.GCLOUD_LOCATION;
  const location = chosenLocation || 'us-central1';
  if (!location) {
    throw new Error(
      `VertexAI Plugin is missing the 'location' configuration. Please set the 'GCLOUD_LOCATION' environment variable or explicitly pass 'location' into genkit config.`
    );
  }

  const clientOpt: RegionalClientOptions = {
    kind: 'regional',
    location,
    projectId,
    authClient,
    experimental_debugTraces: options?.experimental_debugTraces,
  };
  if (!chosenLocation) {
    clientOpt.locationDefaulted = true;
  }
  if (options?.apiKey) {
    clientOpt.apiKey = options.apiKey;
  }
  if (options?.apiVersion) {
    clientOpt.apiVersion = options.apiVersion;
  }
  return clientOpt;
}

/**
 * Client options for an Interactions API request. The Interactions API is
 * only on the global location (and the us/eu multi-regions), so when the
 * location is the built-in default (`us-central1`, nothing chosen) the
 * request goes to `global`. A chosen location is used as is.
 */
export function toInteractionsClientOptions(
  clientOptions: ClientOptions
): ClientOptions {
  if (clientOptions.kind === 'regional' && clientOptions.locationDefaulted) {
    const {
      locationDefaulted: _,
      kind: _kind,
      location: _location,
      ...rest
    } = clientOptions;
    const globalOptions: GlobalClientOptions = {
      ...rest,
      kind: 'global',
      location: 'global',
    };
    return globalOptions;
  }
  return clientOptions;
}

export type RequestClientOptions = ClientOptions & {
  signal: AbortSignal;
};

/**
 * If location or apiKey are present in reqConfig, they will
 * override the values in the clientOptions. The newOptions will
 * contain the clientOptions with those overrides.
 * @param clientOptions The client options
 * @param reqConfig The request config
 */
export function calculateRequestOptions<T extends z.ZodObject<any, any, any>>(
  clientOptions: RequestClientOptions,
  reqConfig?: z.infer<T>
): RequestClientOptions;
export function calculateRequestOptions<T extends z.ZodObject<any, any, any>>(
  clientOptions: ClientOptions,
  reqConfig?: z.infer<T>
): ClientOptions;
export function calculateRequestOptions<T extends z.ZodObject<any, any, any>>(
  clientOptions: RequestClientOptions | ClientOptions,
  reqConfig?: z.infer<T>
): RequestClientOptions | ClientOptions {
  let newOptions = { ...clientOptions };
  if (reqConfig?.apiVersion) {
    newOptions.apiVersion = reqConfig.apiVersion;
  }
  // A location set in the request is chosen, even if it equals the default.
  if (
    reqConfig?.location &&
    typeof reqConfig.location == 'string' &&
    newOptions.kind == 'regional'
  ) {
    delete newOptions.locationDefaulted;
  }
  if (
    reqConfig?.location &&
    typeof reqConfig.location == 'string' &&
    newOptions.kind != 'express' &&
    newOptions.location != reqConfig.location
  ) {
    // Override the location if it's specified in the request
    if (reqConfig.location == 'global') {
      newOptions.location = 'global';
      newOptions.kind = 'global';
    } else if (isMultiRegionalLocation(reqConfig.location)) {
      newOptions.kind = 'multi-regional';
      newOptions.location = reqConfig.location;
    } else {
      newOptions.kind = 'regional';
      newOptions.location = reqConfig.location;
    }
  }
  if (
    clientOptions.kind == 'express' &&
    reqConfig?.apiKey &&
    typeof reqConfig.apiKey == 'string'
  ) {
    newOptions.apiKey = calculateApiKey(clientOptions.apiKey, reqConfig.apiKey);
  } else if (reqConfig?.apiKey && typeof reqConfig.apiKey == 'string') {
    // Regional or Global can still use APIKey for billing (not auth)
    newOptions.apiKey = reqConfig.apiKey;
  }
  return newOptions;
}

/**
 * Applies connection-related overrides carried in the action run context
 * (e.g. `options.context` on `start`/`check`/`cancel` for background models)
 * on top of the given client options. This is the mechanism by which callers
 * can supply per-call overrides for calls (like
 * `checkOperation`/`cancelOperation`) that don't have access to the original
 * request's `config`.
 *
 * The API key is read from `context.secrets.apiKey` (scrubbed from traces).
 * All other overrides (e.g. `location`, `apiVersion`) are read from
 * `context.config`, which is where `checkOperation`/`cancelOperation` fold
 * their top-level `config` option. This lets callers write
 * `ai.checkOperation(op, { config: { location } })`.
 *
 * All context-derived values take precedence over whatever is already in
 * `clientOptions`.
 *
 * @param clientOptions The base client options to apply overrides on top of.
 * @param context The action run context (`options.context`), if any.
 */
export function applyContextOverrides(
  clientOptions: ClientOptions,
  context?: ActionContext
): ClientOptions {
  if (!context) {
    return clientOptions;
  }
  let newOptions = { ...clientOptions };
  const contextApiKey = context.secrets?.apiKey;
  if (typeof contextApiKey === 'string') {
    if (newOptions.kind == 'express') {
      newOptions.apiKey = calculateApiKey(newOptions.apiKey, contextApiKey);
    } else {
      // Regional or Global can still use APIKey for billing (not auth)
      newOptions.apiKey = contextApiKey;
    }
  }
  // All non-secret overrides (e.g. location, apiVersion) are carried in
  // `context.config` (from checkOperation/cancelOperation's top-level
  // `config`).
  return calculateRequestOptions(newOptions, context.config);
}

/**
 * Retrieves an API key from environment variables.
 *
 * @returns The API key as a string, or `undefined` if none of the specified
 *          environment variables are set.
 */
export function getApiKeyFromEnvVar(): string | undefined {
  return (
    process.env.VERTEX_API_KEY ||
    process.env.GOOGLE_API_KEY ||
    process.env.GOOGLE_GENAI_API_KEY
  );
}

export const MISSING_API_KEY_ERROR = new GenkitError({
  status: 'FAILED_PRECONDITION',
  message:
    'Please pass in the API key or set the VERTEX_API_KEY or GOOGLE_API_KEY environment variable.\n' +
    'For more details see https://genkit.dev/docs/integrations/google-genai',
});

export const API_KEY_FALSE_ERROR = new GenkitError({
  status: 'INVALID_ARGUMENT',
  message:
    'VertexAI plugin was initialized with {apiKey: false} but no apiKey configuration was passed at call time.',
});

export const NOT_SUPPORTED_IN_EXPRESS_ERROR = new GenkitError({
  status: 'PERMISSION_DENIED',
  message:
    'This method or model is not supported in Vertex AI Express Mode.\n' +
    'For more details see https://cloud.google.com/vertex-ai/generative-ai/docs/start/express-mode/vertex-ai-express-mode-api-reference',
});

/**
 * Checks and retrieves an API key based on the provided argument and environment variables.
 *
 * - If `pluginApiKey` is a non-empty string, it's used as the API key.
 * - If `pluginApiKey` is `undefined` or an empty string, it attempts to fetch the API key from environment
 * - If `pluginApiKey` is `false`, key retrieval from the environment is skipped, and the function
 *   will return `undefined`. This mode indicates that the API key is expected to be provided
 *   at a later stage or in a different context.
 *
 * @param pluginApiKey - An optional API key string, `undefined` to check the environment, or `false` to bypass all checks in this function.
 * @returns The resolved API key as a string, or `undefined` if `pluginApiKey` is `false`.
 * @throws {Error} MISSING_API_KEY_ERROR - Thrown if `pluginApiKey` is not `false` and no API key
 *   can be found either in the `pluginApiKey` argument or from the environment.
 */
export function checkApiKey(
  pluginApiKey: string | false | undefined
): string | undefined {
  let apiKey: string | undefined;

  // Don't get the key from the environment if pluginApiKey is false
  if (pluginApiKey !== false) {
    apiKey = pluginApiKey || getApiKeyFromEnvVar();
  }

  // If pluginApiKey is false, then we don't throw because we are waiting for
  // the apiKey passed into the individual call
  if (pluginApiKey !== false && !apiKey) {
    throw MISSING_API_KEY_ERROR;
  }
  return apiKey;
}

/**
 * Calculates and returns the effective API key based on multiple potential sources.
 * The order of precedence for determining the API key is:
 * 1. `requestApiKey` (if provided)
 * 2. `pluginApiKey` (if provided and not `false`)
 * 3. Environment variable (if `pluginApiKey` is not `false` and `pluginApiKey` is not provided)
 *
 * @param pluginApiKey - The apiKey value provided during plugin initialization.
 * @param requestApiKey - The apiKey provided to an individual generate call.
 * @returns The resolved API key as a string.
 * @throws {Error} API_KEY_FALSE_ERROR - Thrown if `pluginApiKey` is `false` and `requestApiKey` is not provided
 * @throws {Error} MISSING_API_KEY_ERROR - Thrown if no API key can be resolved from any source
 */
export function calculateApiKey(
  pluginApiKey: string | false | undefined,
  requestApiKey: string | undefined
): string {
  let apiKey: string | undefined;

  // Don't get the key from the environment if pluginApiKey is false
  if (pluginApiKey !== false) {
    apiKey = pluginApiKey || getApiKeyFromEnvVar();
  }

  apiKey = requestApiKey || apiKey;

  if (pluginApiKey === false && !requestApiKey) {
    throw API_KEY_FALSE_ERROR;
  }

  if (!apiKey) {
    throw MISSING_API_KEY_ERROR;
  }
  return apiKey;
}

/** Vertex Express Mode lets you try a *subset* of Vertex AI features */
export function checkSupportedResourceMethod(params: {
  clientOptions: ClientOptions;
  resourcePath?: string;
  resourceMethod?: string;
}) {
  if (params.resourcePath == '') {
    // This is how we get a base url for metadata
    return;
  }

  const supportedExpressMethods = [
    'countTokens',
    'generateContent',
    'streamGenerateContent',
  ];

  // The Interactions API has no resource method (`.../interactions`) and is
  // available in Express mode.
  const isInteractions =
    !params.resourceMethod &&
    !!params.resourcePath?.split('/').includes('interactions');

  if (
    params.clientOptions.kind === 'express' &&
    ((!isInteractions &&
      !supportedExpressMethods.includes(params.resourceMethod ?? '')) ||
      params.resourcePath?.includes('endpoints/'))
  ) {
    throw NOT_SUPPORTED_IN_EXPRESS_ERROR;
  }
}