/**
 * 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, StatusName, type ErrorResponseMetadata } from 'genkit';
import { logger } from 'genkit/logging';
import { runInNewSpan } from 'genkit/tracing';
import { GoogleAuth } from 'google-auth-library';
import {
  CreateInteractionRequest,
  GeminiInteraction,
  InteractionStreamResult,
} from '../common/interaction-types.js';
import { isObject } from '../common/types.js';
import {
  INTERACTIONS_API_REVISION,
  buildTraceMetadataInput,
  extractErrMsg,
  getGenkitClientHeader,
  httpStatusToGenkitStatus,
  interactionErrorCodeToGenkitStatus,
  interactionProcessStream,
  isInteractionUrl,
  parseRetryAfterMs,
  processStream,
  withoutHeaders,
} from '../common/utils.js';
import {
  ClientOptions,
  EmbedContentRequest,
  EmbedContentResponse,
  EmbedPredictRequest,
  EmbedPredictResponse,
  GenerateContentRequest,
  GenerateContentResponse,
  GenerateContentStreamResult,
  ImagenPredictRequest,
  ImagenPredictResponse,
  ListModelsResponse,
  LyriaPredictRequest,
  LyriaPredictResponse,
  Model,
  VeoOperation,
  VeoOperationRequest,
  VeoPredictRequest,
} from './types.js';
import { calculateApiKey, checkSupportedResourceMethod } from './utils.js';

export const tracingHooks = { runInNewSpan };

interface TraceOptions<T> {
  request?: unknown;
  model?: string;
  streaming?: boolean;
  processFn?: (resp: Response) => Promise<T>;
  clientOptions?: ClientOptions;
}

async function maybeTraceRequest<T>(
  url: string,
  fetchOptions: RequestInit,
  traceOptions: TraceOptions<T>
): Promise<T> {
  const call = async () => {
    const response = await makeRequest(url, fetchOptions);
    let processedResponse: Promise<T>;
    if (traceOptions.processFn) {
      // This is for streaming etc.
      processedResponse = traceOptions.processFn(response);
    } else {
      // default processing is just get the json response
      processedResponse = await response.json();
    }
    return processedResponse;
  };

  if (traceOptions.clientOptions?.experimental_debugTraces) {
    return tracingHooks.runInNewSpan(
      {
        metadata: {
          name: 'httpRequest',
          // Seed input up front so the realtime "pending" span export carries it.
          input: buildTraceMetadataInput(url, fetchOptions, traceOptions),
        },
      },
      async (metadata) => {
        const processedResponse = await call();

        if (traceOptions.streaming) {
          metadata.output = '[Streaming Response]';
        } else {
          metadata.output = processedResponse;
        }

        return processedResponse;
      }
    );
  }
  return call();
}

export async function listModels(
  clientOptions: ClientOptions
): Promise<Model[]> {
  const url = getVertexAIUrl({
    includeProjectAndLocation: false,
    resourcePath: 'publishers/google/models',
    clientOptions,
  });
  const fetchOptions = await getFetchOptions({
    method: 'GET',
    clientOptions,
  });
  const response = await makeRequest(url, fetchOptions);
  const modelResponse = (await response.json()) as ListModelsResponse;
  return modelResponse.publisherModels;
}

/**
 * Resource path of the Interactions API. Express mode has no project, but
 * still needs the location: `v1beta1/locations/global/interactions`.
 */
function getInteractionsResourcePath(clientOptions: ClientOptions): string {
  return clientOptions.kind == 'express'
    ? 'locations/global/interactions'
    : 'interactions';
}

export async function createInteraction(
  createInteractionRequest: CreateInteractionRequest,
  clientOptions: ClientOptions
): Promise<GeminiInteraction> {
  const url = getVertexAIUrl({
    includeProjectAndLocation: true,
    resourcePath: getInteractionsResourcePath(clientOptions),
    clientOptions,
  });

  const fetchOptions = await getFetchOptions({
    method: 'POST',
    clientOptions,
    body: JSON.stringify(createInteractionRequest),
    isInteraction: true,
  });

  return maybeTraceRequest<GeminiInteraction>(url, fetchOptions, {
    request: createInteractionRequest,
    clientOptions,
  });
}

export async function createInteractionStream(
  createInteractionRequest: CreateInteractionRequest,
  clientOptions: ClientOptions
): Promise<InteractionStreamResult> {
  const url = getVertexAIUrl({
    includeProjectAndLocation: true,
    resourcePath: getInteractionsResourcePath(clientOptions),
    clientOptions,
  });

  const fetchOptions = await getFetchOptions({
    method: 'POST',
    clientOptions,
    body: JSON.stringify(createInteractionRequest),
    isInteraction: true,
  });

  return maybeTraceRequest<InteractionStreamResult>(url, fetchOptions, {
    request: createInteractionRequest,
    streaming: true,
    clientOptions,
    processFn: async (response) => {
      return interactionProcessStream(response);
    },
  });
}

export async function generateContent(
  model: string,
  generateContentRequest: GenerateContentRequest,
  clientOptions: ClientOptions
): Promise<GenerateContentResponse> {
  let url: string;
  if (model.includes('endpoints/')) {
    // Tuned model
    url = getVertexAIUrl({
      includeProjectAndLocation: !model.startsWith('projects/'),
      resourcePath: model,
      resourceMethod: 'generateContent',
      clientOptions,
    });
  } else {
    url = getVertexAIUrl({
      includeProjectAndLocation: true,
      resourcePath: `publishers/google/models/${model}`,
      resourceMethod: 'generateContent',
      clientOptions,
    });
  }
  const fetchOptions = await getFetchOptions({
    method: 'POST',
    clientOptions,
    body: JSON.stringify(generateContentRequest),
  });

  return maybeTraceRequest<GenerateContentResponse>(url, fetchOptions, {
    model: model,
    request: generateContentRequest,
    clientOptions,
  });
}

export async function generateContentStream(
  model: string,
  generateContentRequest: GenerateContentRequest,
  clientOptions: ClientOptions
): Promise<GenerateContentStreamResult> {
  let url: string;
  if (model.includes('endpoints/')) {
    // Tuned model
    url = getVertexAIUrl({
      includeProjectAndLocation: !model.startsWith('projects/'),
      resourcePath: model,
      resourceMethod: 'streamGenerateContent',
      clientOptions,
    });
  } else {
    url = getVertexAIUrl({
      includeProjectAndLocation: true,
      resourcePath: `publishers/google/models/${model}`,
      resourceMethod: 'streamGenerateContent',
      clientOptions,
    });
  }
  const fetchOptions = await getFetchOptions({
    method: 'POST',
    clientOptions,
    body: JSON.stringify(generateContentRequest),
  });

  return maybeTraceRequest<GenerateContentStreamResult>(url, fetchOptions, {
    model: model,
    request: generateContentRequest,
    streaming: true,
    clientOptions,
    processFn: async (response) => {
      return processStream(response);
    },
  });
}

async function internalPredict<T>(
  model: string,
  requestPayload: unknown,
  clientOptions: ClientOptions
): Promise<T> {
  const url = getVertexAIUrl({
    includeProjectAndLocation: true,
    resourcePath: `publishers/google/models/${model}`,
    resourceMethod: 'predict',
    clientOptions,
  });

  const fetchOptions = await getFetchOptions({
    method: 'POST',
    clientOptions,
    body: JSON.stringify(requestPayload),
  });

  return maybeTraceRequest<T>(url, fetchOptions, {
    model: model,
    request: requestPayload,
    clientOptions,
  });
}

/**
 * Calls the `:predict` embeddings method (gemini-embedding-001,
 * text-embedding-*, multimodalembedding@001).
 */
export async function embedPredict(
  model: string,
  embedPredictRequest: EmbedPredictRequest,
  clientOptions: ClientOptions
): Promise<EmbedPredictResponse> {
  return internalPredict<EmbedPredictResponse>(
    model,
    embedPredictRequest,
    clientOptions
  );
}

/**
 * Calls the `:embedContent` method (gemini-embedding-2).
 */
export async function embedContent(
  model: string,
  embedContentRequest: EmbedContentRequest,
  clientOptions: ClientOptions
): Promise<EmbedContentResponse> {
  const url = getVertexAIUrl({
    includeProjectAndLocation: true,
    resourcePath: `publishers/google/models/${model}`,
    resourceMethod: 'embedContent',
    clientOptions,
  });

  const fetchOptions = await getFetchOptions({
    method: 'POST',
    clientOptions,
    body: JSON.stringify(embedContentRequest),
  });

  return maybeTraceRequest<EmbedContentResponse>(url, fetchOptions, {
    model,
    request: embedContentRequest,
    clientOptions,
  });
}

export async function imagenPredict(
  model: string,
  imagenPredictRequest: ImagenPredictRequest,
  clientOptions: ClientOptions
): Promise<ImagenPredictResponse> {
  return internalPredict<ImagenPredictResponse>(
    model,
    imagenPredictRequest,
    clientOptions
  );
}

export async function lyriaPredict(
  model: string,
  lyriaPredictRequest: LyriaPredictRequest,
  clientOptions: ClientOptions
): Promise<LyriaPredictResponse> {
  return internalPredict<LyriaPredictResponse>(
    model,
    lyriaPredictRequest,
    clientOptions
  );
}

export async function veoPredict(
  model: string,
  veoPredictRequest: VeoPredictRequest,
  clientOptions: ClientOptions
): Promise<VeoOperation> {
  const url = getVertexAIUrl({
    includeProjectAndLocation: true,
    resourcePath: `publishers/google/models/${model}`,
    resourceMethod: 'predictLongRunning',
    clientOptions,
  });

  const fetchOptions = await getFetchOptions({
    method: 'POST',
    clientOptions,
    body: JSON.stringify(veoPredictRequest),
  });

  return maybeTraceRequest<VeoOperation>(url, fetchOptions, {
    model: model,
    request: veoPredictRequest,
    clientOptions,
  });
}

export async function veoCheckOperation(
  model: string,
  veoOperationRequest: VeoOperationRequest,
  clientOptions: ClientOptions
): Promise<VeoOperation> {
  const url = getVertexAIUrl({
    includeProjectAndLocation: true,
    resourcePath: `publishers/google/models/${model}`,
    resourceMethod: 'fetchPredictOperation',
    clientOptions,
  });
  const fetchOptions = await getFetchOptions({
    method: 'POST',
    clientOptions,
    body: JSON.stringify(veoOperationRequest),
  });

  return maybeTraceRequest<VeoOperation>(url, fetchOptions, {
    model: model,
    request: veoOperationRequest,
    clientOptions,
  });
}

export function getVertexAIUrl(params: {
  includeProjectAndLocation: boolean; // False for listModels, true for most others
  resourcePath: string;
  resourceMethod?: string;
  queryParams?: string;
  clientOptions: ClientOptions;
}): string {
  checkSupportedResourceMethod(params);

  const DEFAULT_API_VERSION = 'v1beta1';
  const API_BASE_PATH = 'aiplatform.googleapis.com';

  let basePath: string;

  if (params.clientOptions.kind == 'regional') {
    basePath = `${params.clientOptions.location}-${API_BASE_PATH}`;
  } else if (params.clientOptions.kind == 'multi-regional') {
    basePath = `aiplatform.${params.clientOptions.location}.rep.googleapis.com`;
  } else {
    basePath = API_BASE_PATH;
  }

  let resourcePath = params.resourcePath;
  if (
    params.clientOptions.kind != 'express' &&
    params.includeProjectAndLocation
  ) {
    const parent = `projects/${params.clientOptions.projectId}/locations/${params.clientOptions.location}`;
    resourcePath = `${parent}/${params.resourcePath}`;
  }

  const version = params.clientOptions.apiVersion ?? DEFAULT_API_VERSION;

  let url = `https://${basePath}/${version}/${resourcePath}`;
  if (params.resourceMethod) {
    url += `:${params.resourceMethod}`;
  }

  let joiner = '?';
  if (params.queryParams) {
    url += `${joiner}${params.queryParams}`;
    joiner = '&';
  }
  if (params.resourceMethod === 'streamGenerateContent') {
    url += `${joiner}alt=sse`;
    joiner = '&';
  }
  return url;
}

async function getFetchOptions(params: {
  method: 'POST' | 'GET';
  body?: string;
  clientOptions: ClientOptions;
  isInteraction?: boolean;
}) {
  const fetchOptions: RequestInit = {
    method: params.method,
    headers: await getHeaders(params.clientOptions, params.isInteraction),
    redirect: 'manual',
  };
  if (params.body) {
    fetchOptions.body = params.body;
  }
  const signal = getAbortSignal(params.clientOptions);
  if (signal) {
    fetchOptions.signal = signal;
  }
  return fetchOptions;
}

function getAbortSignal(clientOptions: ClientOptions): AbortSignal | undefined {
  const hasTimeout = (clientOptions.timeout ?? -1) >= 0;
  if (clientOptions.signal !== undefined || hasTimeout) {
    const controller = new AbortController();
    if (hasTimeout) {
      setTimeout(() => controller.abort(), clientOptions.timeout);
    }
    if (clientOptions?.signal) {
      clientOptions.signal.addEventListener('abort', () => {
        controller.abort();
      });
    }
    return controller.signal;
  }
  return undefined;
}

async function getHeaders(
  clientOptions: ClientOptions,
  isInteraction?: boolean
): Promise<HeadersInit> {
  // The Interactions converters only understand INTERACTIONS_API_REVISION's
  // schema, so custom headers can't change it (in any letter case).
  const customHeaders = withoutHeaders(
    clientOptions.customHeaders,
    isInteraction ? ['Api-Revision'] : []
  );
  if (clientOptions.kind == 'express') {
    const headers: HeadersInit = {
      'x-goog-api-key': calculateApiKey(clientOptions.apiKey, undefined),
      'Content-Type': 'application/json',
      'X-Goog-Api-Client': getGenkitClientHeader(),
      'User-Agent': getGenkitClientHeader(),
      ...customHeaders,
    };
    if (isInteraction) {
      headers['Api-Revision'] = INTERACTIONS_API_REVISION;
    }
    return headers;
  } else {
    const token = await getToken(clientOptions.authClient);
    const headers: HeadersInit = {
      Authorization: `Bearer ${token}`,
      'x-goog-user-project': clientOptions.projectId,
      'Content-Type': 'application/json',
      'X-Goog-Api-Client': getGenkitClientHeader(),
      'User-Agent': getGenkitClientHeader(),
      ...customHeaders,
    };
    if (clientOptions.apiKey) {
      headers['x-goog-api-key'] = clientOptions.apiKey;
    }
    if (isInteraction) {
      headers['Api-Revision'] = INTERACTIONS_API_REVISION;
    }
    return headers;
  }
}

async function getToken(authClient: GoogleAuth): Promise<string> {
  const CREDENTIAL_ERROR_MESSAGE =
    '\nUnable to authenticate your request\
        \nDepending on your run time environment, you can get authentication by\
        \n- if in local instance or cloud shell: `!gcloud auth login`\
        \n- if in Colab:\
        \n    -`from google.colab import auth`\
        \n    -`auth.authenticate_user()`\
        \n- if in service account or other: please follow guidance in https://cloud.google.com/docs/authentication';
  const token = await authClient.getAccessToken().catch((e) => {
    throw new Error(CREDENTIAL_ERROR_MESSAGE, e);
  });
  if (!token) {
    throw new Error(CREDENTIAL_ERROR_MESSAGE);
  }
  return token;
}

async function makeRequest(
  url: string,
  fetchOptions: RequestInit
): Promise<Response> {
  try {
    const response = await fetch(url, fetchOptions);
    if (!response.ok) {
      let errorText = await response.text();
      let errorMessage = errorText;
      let errorDetail: unknown;
      try {
        const json = JSON.parse(errorText);
        errorDetail = json;
        if (json.error && json.error.message) {
          errorMessage = json.error.message;
          if (Array.isArray(json.error.details)) {
            const detailsText = json.error.details
              .map((d: any) => {
                if (d.detail && typeof d.detail === 'string') {
                  const match = d.detail.match(
                    /\[ORIGINAL ERROR\]\s*(?:generic::[^:]+:\s*)?(.*?)(?:\s+\[|\s+\d+\s+\{|$)/
                  );
                  return match ? match[1].trim() : d.detail;
                }
                return JSON.stringify(d);
              })
              .filter(Boolean)
              .join('\n');
            if (detailsText) {
              errorMessage += `\nDetails: ${detailsText}`;
            }
          }
        }
      } catch (e) {
        // Not JSON or expected format, use the raw text
      }
      let status: StatusName = httpStatusToGenkitStatus(response.status);
      // Interactions API errors carry a specific `error.code` (e.g. `safety`,
      // `model_not_found`, `quota_exceeded`) that is more precise than the
      // HTTP status, so prefer it when recognized.
      if (isInteractionUrl(url)) {
        const apiError = isObject(errorDetail) ? errorDetail.error : undefined;
        const codeStatus = interactionErrorCodeToGenkitStatus(
          isObject(apiError) ? apiError.code : undefined
        );
        if (codeStatus) {
          status = codeStatus;
        }
      }
      // Capture Retry-After header for retry middleware to use
      const retryAfterHeader = response.headers.get('retry-after');
      const retryAfterMs = retryAfterHeader
        ? parseRetryAfterMs(retryAfterHeader)
        : undefined;
      const responseMetadata: ErrorResponseMetadata | undefined =
        retryAfterMs !== undefined ? { retryAfterMs } : undefined;

      throw new GenkitError({
        status,
        message: `Error fetching from ${url}: [${response.status} ${response.statusText}] ${errorMessage}`,
        detail: errorDetail,
        responseMetadata,
      });
    }
    return response;
  } catch (e: unknown) {
    logger.error(e);
    if (e instanceof GenkitError) {
      throw e;
    }
    throw new Error(`Failed to fetch from ${url}: ${extractErrMsg(e)}`);
  }
}

export const TEST_ONLY = {
  getFetchOptions,
  getAbortSignal,
  getHeaders,
  makeRequest,
  tracingHooks,
};