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

import {
  afterAll,
  beforeAll,
  beforeEach,
  describe,
  it,
  jest,
} from '@jest/globals';
import {
  DataPoint,
  Histogram,
  HistogramMetricData,
  ScopeMetrics,
  SumMetricData,
} from '@opentelemetry/sdk-metrics';
import { ReadableSpan } from '@opentelemetry/sdk-trace-base';
import * as assert from 'assert';
import { GenerateResponseData, Genkit, genkit, z } from 'genkit';
import { SPAN_TYPE_ATTR, appendSpan } from 'genkit/tracing';
import {
  GcpOpenTelemetry,
  __forceFlushSpansForTesting,
  __getMetricExporterForTesting,
  __getSpanExporterForTesting,
  enableGoogleCloudTelemetry,
} from '../src/index.js';

jest.mock('../src/auth.js', () => {
  const original = jest.requireActual('../src/auth.js');
  return {
    ...(original || {}),
    resolveCurrentPrincipal: jest.fn().mockImplementation(() => {
      return Promise.resolve({
        projectId: 'test',
        serviceAccountEmail: 'test@test.com',
      });
    }),
    credentialsFromEnvironment: jest.fn().mockImplementation(() => {
      return Promise.resolve({
        projectId: 'test',
        credentials: {
          client_email: 'test@genkit.com',
          private_key: '-----BEGIN PRIVATE KEY-----',
        },
      });
    }),
  };
});

describe('GoogleCloudMetrics', () => {
  let ai: Genkit;

  beforeAll(async () => {
    process.env.GCLOUD_PROJECT = 'test';
    process.env.GENKIT_ENV = 'dev';
    await enableGoogleCloudTelemetry({
      projectId: 'test',
      forceDevExport: false,
      metricExportIntervalMillis: 100,
      metricExportTimeoutMillis: 100,
    });
    ai = genkit({});
  });
  beforeEach(async () => {
    __getMetricExporterForTesting().reset();
    __getSpanExporterForTesting().reset();
  });
  afterAll(async () => {
    await ai.stopServers();
  });

  it('writes action metrics for a successful flow', async () => {
    const testFlow = createFlow(ai, 'testFlow');

    await testFlow();
    await testFlow();

    await getExportedSpans();

    const actionRequestCounter = await getCounterMetric(
      'genkit/action/requests'
    );
    const actionLatencyHistogram = await getHistogramMetric(
      'genkit/action/latency'
    );
    assert.equal(actionRequestCounter.value, 2);
    assert.equal(actionRequestCounter.attributes.name, 'testFlow');
    assert.equal(actionRequestCounter.attributes.source, 'ts');
    assert.equal(actionRequestCounter.attributes.status, 'success');
    assert.ok(actionRequestCounter.attributes.sourceVersion);
    assert.equal(actionLatencyHistogram.value.count, 2);
    assert.equal(actionLatencyHistogram.attributes.name, 'testFlow');
    assert.equal(actionLatencyHistogram.attributes.source, 'ts');
    assert.equal(actionLatencyHistogram.attributes.status, 'success');
    assert.ok(actionLatencyHistogram.attributes.sourceVersion);
  }, 10000); //timeout

  it('writes action metrics for a failing flow', async () => {
    const testFlow = createFlow(ai, 'testFlow', async () => {
      const nothing: { missing?: any } = { missing: 1 };
      delete nothing.missing;
      return nothing.missing.explode;
    });

    assert.rejects(async () => {
      await testFlow();
    });

    await getExportedSpans();

    const requestCounter = await getCounterMetric('genkit/action/requests');
    assert.equal(requestCounter.value, 1);
    assert.equal(requestCounter.attributes.name, 'testFlow');
    assert.equal(requestCounter.attributes.source, 'ts');
    assert.equal(requestCounter.attributes.error, 'TypeError');
    assert.equal(requestCounter.attributes.status, 'failure');
  }, 10000); //timeout

  it('writes feature metrics for a successful flow', async () => {
    const testFlow = createFlow(ai, 'testFlow');

    await testFlow();
    await testFlow();

    await getExportedSpans();

    const requestCounter = await getCounterMetric('genkit/feature/requests');
    const latencyHistogram = await getHistogramMetric('genkit/feature/latency');
    assert.equal(requestCounter.value, 2);
    assert.equal(requestCounter.attributes.name, 'testFlow');
    assert.equal(requestCounter.attributes.source, 'ts');
    assert.equal(requestCounter.attributes.status, 'success');
    assert.ok(requestCounter.attributes.sourceVersion);
    assert.equal(latencyHistogram.value.count, 2);
    assert.equal(latencyHistogram.attributes.name, 'testFlow');
    assert.equal(latencyHistogram.attributes.source, 'ts');
    assert.equal(latencyHistogram.attributes.status, 'success');
    assert.ok(latencyHistogram.attributes.sourceVersion);
  });

  it('writes feature metrics for a failing flow', async () => {
    const testFlow = createFlow(ai, 'testFlow', async () => {
      const nothing: { missing?: any } = { missing: 1 };
      delete nothing.missing;
      return nothing.missing.explode;
    });

    assert.rejects(async () => {
      await testFlow();
    });

    await getExportedSpans();

    const requestCounter = await getCounterMetric('genkit/feature/requests');
    assert.equal(requestCounter.value, 1);
    assert.equal(requestCounter.attributes.name, 'testFlow');
    assert.equal(requestCounter.attributes.source, 'ts');
    assert.equal(requestCounter.attributes.error, 'TypeError');
    assert.equal(requestCounter.attributes.status, 'failure');
  }, 10000); //timeout

  it('writes action metrics', async () => {
    const testAction = createAction(ai, 'testAction');
    const testFlow = createFlow(ai, 'testFlowWithActions', async () => {
      await Promise.all([
        testAction(undefined),
        testAction(undefined),
        testAction(undefined),
      ]);
    });

    await testFlow();
    await testFlow();

    await getExportedSpans();

    const requestCounters = await getCounterDataPoints(
      'genkit/action/requests'
    );
    const latencyHistograms = await getHistogramDataPoints(
      'genkit/action/latency'
    );
    const requestCounter = requestCounters[0]; // the action
    const latencyHistogram = latencyHistograms[0];

    assert.equal(requestCounter.value, 6);
    assert.equal(requestCounter.attributes.name, 'testAction');
    assert.equal(requestCounter.attributes.source, 'ts');
    assert.equal(requestCounter.attributes.status, 'success');
    assert.ok(requestCounter.attributes.sourceVersion);
    assert.equal(requestCounter.attributes.featureName, 'testFlowWithActions');
    assert.equal(latencyHistogram.value.count, 6);
    assert.equal(latencyHistogram.attributes.name, 'testAction');
    assert.equal(latencyHistogram.attributes.source, 'ts');
    assert.equal(latencyHistogram.attributes.status, 'success');
    assert.ok(latencyHistogram.attributes.sourceVersion);
    assert.equal(requestCounter.attributes.featureName, 'testFlowWithActions');
  });

  it('writes feature metrics for an action', async () => {
    const testAction = createAction(ai, 'featureAction');

    await testAction(null);
    await testAction(null);

    await getExportedSpans();

    const requestCounter = await getCounterMetric('genkit/feature/requests');
    const latencyHistogram = await getHistogramMetric('genkit/feature/latency');
    assert.equal(requestCounter.value, 2);
    assert.equal(requestCounter.attributes.name, 'featureAction');
    assert.equal(requestCounter.attributes.source, 'ts');
    assert.equal(requestCounter.attributes.status, 'success');
    assert.ok(requestCounter.attributes.sourceVersion);
    assert.equal(latencyHistogram.value.count, 2);
    assert.equal(latencyHistogram.attributes.name, 'featureAction');
    assert.equal(latencyHistogram.attributes.source, 'ts');
    assert.equal(latencyHistogram.attributes.status, 'success');
    assert.ok(latencyHistogram.attributes.sourceVersion);
  });

  // it('writes feature metrics for prompts')
  // after PR #1029

  it('writes feature metrics for generate', async () => {
    const testModel = createTestModel(ai, 'helloModel');
    await ai.generate({ model: testModel, prompt: 'Hi' });
    await ai.generate({ model: testModel, prompt: 'Yo' });

    await getExportedSpans();

    const requestCounter = await getCounterMetric('genkit/feature/requests');
    const latencyHistogram = await getHistogramMetric('genkit/feature/latency');
    assert.equal(requestCounter.value, 2);
    assert.equal(requestCounter.attributes.name, 'generate');
    assert.equal(requestCounter.attributes.source, 'ts');
    assert.equal(requestCounter.attributes.status, 'success');
    assert.ok(requestCounter.attributes.sourceVersion);
    assert.equal(latencyHistogram.value.count, 2);
    assert.equal(latencyHistogram.attributes.name, 'generate');
    assert.equal(latencyHistogram.attributes.source, 'ts');
    assert.equal(latencyHistogram.attributes.status, 'success');
    assert.ok(latencyHistogram.attributes.sourceVersion);
  });

  it('truncates metric dimensions', async () => {
    const testFlow = createFlow(ai, 'anExtremelyLongFlowNameThatIsTooBig');

    await testFlow();

    await getExportedSpans();

    const requestCounter = await getCounterMetric('genkit/feature/requests');
    const latencyHistogram = await getHistogramMetric('genkit/feature/latency');
    assert.equal(
      requestCounter.attributes.name,
      'anExtremelyLongFlowNameThatIsToo'
    );
    assert.equal(
      latencyHistogram.attributes.name,
      'anExtremelyLongFlowNameThatIsToo'
    );
  });

  it('writes action failure metrics', async () => {
    const testAction = ai.defineTool(
      { name: 'testActionWithFailure', description: 'Just a test' },
      async () => {
        const nothing: { missing?: any } = { missing: 1 };
        delete nothing.missing;
        return nothing.missing.explode;
      }
    );

    assert.rejects(async () => {
      return testAction(null);
    });
    await getExportedSpans();

    const requestCounter = await getCounterMetric('genkit/action/requests');
    assert.equal(requestCounter.value, 1);
    assert.equal(requestCounter.attributes.name, 'testActionWithFailure');
    assert.equal(requestCounter.attributes.source, 'ts');
    assert.equal(requestCounter.attributes.status, 'failure');
    assert.equal(requestCounter.attributes.error, 'TypeError');
  }, 10000); //timeout

  it('writes generate metrics for a successful model action', async () => {
    const testModel = createTestModel(ai, 'testModel');

    await ai.generate({
      model: testModel,
      prompt: 'test prompt',
      config: {
        temperature: 1.0,
        topK: 3,
        topP: 5,
        maxOutputTokens: 7,
      },
    });

    await getExportedSpans();

    const requestCounter = await getCounterMetric(
      'genkit/ai/generate/requests'
    );
    const inputTokenCounter = await getCounterMetric(
      'genkit/ai/generate/input/tokens'
    );
    const outputTokenCounter = await getCounterMetric(
      'genkit/ai/generate/output/tokens'
    );
    const inputCharacterCounter = await getCounterMetric(
      'genkit/ai/generate/input/characters'
    );
    const outputCharacterCounter = await getCounterMetric(
      'genkit/ai/generate/output/characters'
    );
    const inputImageCounter = await getCounterMetric(
      'genkit/ai/generate/input/images'
    );
    const outputImageCounter = await getCounterMetric(
      'genkit/ai/generate/output/images'
    );
    const latencyHistogram = await getHistogramMetric(
      'genkit/ai/generate/latency'
    );
    assert.equal(requestCounter.value, 1);
    assert.equal(inputTokenCounter.value, 10);
    assert.equal(outputTokenCounter.value, 14);
    assert.equal(inputCharacterCounter.value, 8);
    assert.equal(outputCharacterCounter.value, 16);
    assert.equal(inputImageCounter.value, 1);
    assert.equal(outputImageCounter.value, 3);
    assert.equal(latencyHistogram.value.count, 1);
    for (const metric of [
      requestCounter,
      inputTokenCounter,
      outputTokenCounter,
      inputCharacterCounter,
      outputCharacterCounter,
      inputImageCounter,
      outputImageCounter,
      latencyHistogram,
    ]) {
      assert.equal(metric.attributes.modelName, 'testModel');
      assert.equal(metric.attributes.source, 'ts');
      assert.equal(metric.attributes.status, 'success');
      assert.equal(metric.attributes.featureName, 'generate');
      assert.ok(metric.attributes.sourceVersion);
    }
  });

  it('writes generate metrics for a failing model action', async () => {
    const testModel = createModel(ai, 'failingTestModel', async () => {
      const nothing: { missing?: any } = { missing: 1 };
      delete nothing.missing;
      return nothing.missing.explode;
    });

    assert.rejects(async () => {
      return ai.generate({
        model: testModel,
        prompt: 'test prompt',
        config: {
          temperature: 1.0,
          topK: 3,
          topP: 5,
          maxOutputTokens: 7,
        },
      });
    });

    await getExportedSpans();

    const requestCounter = await getCounterMetric(
      'genkit/ai/generate/requests'
    );
    assert.equal(requestCounter.value, 1);
    assert.equal(requestCounter.attributes.modelName, 'failingTestModel');
    assert.equal(requestCounter.attributes.source, 'ts');
    assert.equal(requestCounter.attributes.status, 'failure');
    assert.equal(requestCounter.attributes.error, 'TypeError');
    assert.ok(requestCounter.attributes.sourceVersion);
  }, 10000); //timeout

  it('writes flow label to action metrics when running inside flow', async () => {
    const testAction = createAction(ai, 'testAction');
    const flow = createFlow(ai, 'flowNameLabelTestFlow', async () => {
      return await testAction(undefined);
    });

    await flow();

    await getExportedSpans();

    const requestCounter = await getCounterMetric('genkit/action/requests');
    const latencyHistogram = await getHistogramMetric('genkit/action/latency');
    assert.equal(
      requestCounter.attributes.featureName,
      'flowNameLabelTestFlow'
    );
    assert.equal(
      latencyHistogram.attributes.featureName,
      'flowNameLabelTestFlow'
    );
  });

  it('writes feature label to generate and action metrics when running inside an action', async () => {
    const testModel = createTestModel(ai, 'testModel');

    const testAction = createAction(ai, 'testGenerateAction', async () => {
      await ai.generate({
        model: testModel,
        prompt: 'helllloooooooo',
      });
    });

    testAction(null);

    await getExportedSpans();

    const requestCounter = await getCounterMetric('genkit/action/requests');
    const latencyHistogram = await getHistogramMetric('genkit/action/latency');
    const generateRequestCounter = await getCounterMetric(
      'genkit/ai/generate/requests'
    );
    assert.equal(requestCounter.attributes.featureName, 'testGenerateAction');
    assert.equal(latencyHistogram.attributes.featureName, 'testGenerateAction');
    assert.equal(
      generateRequestCounter.attributes.featureName,
      'testGenerateAction'
    );
  }, 10000); //timeout

  it('writes feature label to generate metrics when running inside a flow', async () => {
    const testModel = createModel(ai, 'testModel', async () => {
      return {
        message: {
          role: 'user',
          content: [
            {
              text: 'response',
            },
          ],
        },
        finishReason: 'stop',
        usage: {
          inputTokens: 10,
          outputTokens: 14,
          inputCharacters: 8,
          outputCharacters: 16,
          inputImages: 1,
          outputImages: 3,
        },
      };
    });
    const flow = createFlow(ai, 'testFlow', async () => {
      return await ai.generate({
        model: testModel,
        prompt: 'test prompt',
      });
    });

    await flow();

    await getExportedSpans();

    const metrics = [
      await getCounterMetric('genkit/ai/generate/requests'),
      await getCounterMetric('genkit/ai/generate/input/tokens'),
      await getCounterMetric('genkit/ai/generate/output/tokens'),
      await getCounterMetric('genkit/ai/generate/input/characters'),
      await getCounterMetric('genkit/ai/generate/output/characters'),
      await getCounterMetric('genkit/ai/generate/input/images'),
      await getCounterMetric('genkit/ai/generate/output/images'),
      await getHistogramMetric('genkit/ai/generate/latency'),
    ];
    for (const metric of metrics) {
      assert.equal(metric.attributes.featureName, 'testFlow');
    }
  });

  it('writes path metrics for a successful flow', async () => {
    const flow = createFlow(ai, 'pathTestFlow', async () => {
      await ai.run('step1', async () => {
        return await ai.run('substep_a', async () => {
          return await ai.run('substep_b', async () => 'res1');
        });
      });
      await ai.run('step2', async () => 'res2');
      return;
    });

    await flow();

    await getExportedSpans();

    const expectedPaths = new Set([
      '/{pathTestFlow,t:flow}/{step2,t:flowStep}',
      '/{pathTestFlow,t:flow}/{step1,t:flowStep}/{substep_a,t:flowStep}/{substep_b,t:flowStep}',
    ]);
    const pathCounterPoints = await getCounterDataPoints(
      'genkit/feature/path/requests'
    );
    const pathLatencyPoints = await getHistogramDataPoints(
      'genkit/feature/path/latency'
    );
    const paths = new Set(
      pathCounterPoints.map((point) => point.attributes.path)
    );
    assert.deepEqual(paths, expectedPaths);
    pathCounterPoints.forEach((point) => {
      assert.equal(point.value, 1);
      assert.equal(point.attributes.featureName, 'pathTestFlow');
      assert.equal(point.attributes.source, 'ts');
      assert.equal(point.attributes.status, 'success');
      assert.ok(point.attributes.sourceVersion);
    });
    pathLatencyPoints.forEach((point) => {
      assert.equal(point.value.count, 1);
      assert.equal(point.attributes.featureName, 'pathTestFlow');
      assert.equal(point.attributes.source, 'ts');
      assert.equal(point.attributes.status, 'success');
      assert.ok(point.attributes.sourceVersion);
    });
  });

  it('writes path metrics for a failing flow with exception in root', async () => {
    const flow = createFlow(ai, 'testFlow', async () => {
      await ai.run('sub-action', async () => {
        return 'done';
      });
      return Promise.reject(new Error('failed'));
    });

    assert.rejects(async () => {
      await flow();
    });

    await getExportedSpans();

    const reqPoints = await getCounterDataPoints(
      'genkit/feature/path/requests'
    );
    const reqStatuses = reqPoints.map((p) => [
      p.attributes.path,
      p.attributes.status,
    ]);
    assert.deepEqual(reqStatuses, [
      ['/{testFlow,t:flow}/{sub-action,t:flowStep}', 'success'],
      ['/{testFlow,t:flow}', 'failure'],
    ]);
    const latencyPoints = await getHistogramDataPoints(
      'genkit/feature/path/latency'
    );
    const latencyStatuses = latencyPoints.map((p) => [
      p.attributes.path,
      p.attributes.status,
    ]);
    assert.deepEqual(latencyStatuses, [
      ['/{testFlow,t:flow}/{sub-action,t:flowStep}', 'success'],
      ['/{testFlow,t:flow}', 'failure'],
    ]);
  }, 10000); //timeout

  it('writes path metrics for a failing flow with exception in subaction', async () => {
    const flow = createFlow(ai, 'testFlow', async () => {
      await ai.run('sub-action-1', async () => {
        await ai.run('sub-action-2', async () => {
          return Promise.reject(new Error('failed'));
        });
        return 'done';
      });
      return 'done';
    });

    assert.rejects(async () => {
      await flow();
    });

    await getExportedSpans();

    const reqPoints = await getCounterDataPoints(
      'genkit/feature/path/requests'
    );
    const reqStatuses = reqPoints.map((p) => [
      p.attributes.path,
      p.attributes.status,
    ]);
    assert.deepEqual(reqStatuses, [
      [
        '/{testFlow,t:flow}/{sub-action-1,t:flowStep}/{sub-action-2,t:flowStep}',
        'failure',
      ],
    ]);
    const latencyPoints = await getHistogramDataPoints(
      'genkit/feature/path/latency'
    );
    const latencyStatuses = latencyPoints.map((p) => [
      p.attributes.path,
      p.attributes.status,
    ]);
    assert.deepEqual(latencyStatuses, [
      [
        '/{testFlow,t:flow}/{sub-action-1,t:flowStep}/{sub-action-2,t:flowStep}',
        'failure',
      ],
    ]);
  }, 10000); //timeout

  it('writes path metrics for a flow with exception in action', async () => {
    const flow = createFlow(ai, 'testFlow', async () => {
      await ai.run('sub-action-1', async () => {
        await ai.run('sub-action-2', async () => {
          return 'done';
        });
        return Promise.reject(new Error('failed'));
      });
      return 'done';
    });

    assert.rejects(async () => {
      await flow();
    });

    await getExportedSpans();

    const reqPoints = await getCounterDataPoints(
      'genkit/feature/path/requests'
    );
    const reqStatuses = reqPoints.map((p) => [
      p.attributes.path,
      p.attributes.status,
    ]);
    assert.deepEqual(reqStatuses, [
      [
        '/{testFlow,t:flow}/{sub-action-1,t:flowStep}/{sub-action-2,t:flowStep}',
        'success',
      ],
      ['/{testFlow,t:flow}/{sub-action-1,t:flowStep}', 'failure'],
    ]);
    const latencyPoints = await getHistogramDataPoints(
      'genkit/feature/path/latency'
    );
    const latencyStatuses = latencyPoints.map((p) => [
      p.attributes.path,
      p.attributes.status,
    ]);
    assert.deepEqual(latencyStatuses, [
      [
        '/{testFlow,t:flow}/{sub-action-1,t:flowStep}/{sub-action-2,t:flowStep}',
        'success',
      ],
      ['/{testFlow,t:flow}/{sub-action-1,t:flowStep}', 'failure'],
    ]);
  }, 10000); //timeout

  it('writes path metrics for a flow with an exception in a serial action', async () => {
    const flow = createFlow(ai, 'testFlow', async () => {
      await ai.run('sub-action-1', async () => {
        return 'done';
      });
      await ai.run('sub-action-2', async () => {
        return Promise.reject(new Error('failed'));
      });
      return 'done';
    });

    assert.rejects(async () => {
      await flow();
    });

    await getExportedSpans();

    const reqPoints = await getCounterDataPoints(
      'genkit/feature/path/requests'
    );
    const reqStatuses = reqPoints.map((p) => [
      p.attributes.path,
      p.attributes.status,
    ]);
    assert.deepEqual(reqStatuses, [
      ['/{testFlow,t:flow}/{sub-action-1,t:flowStep}', 'success'],
      ['/{testFlow,t:flow}/{sub-action-2,t:flowStep}', 'failure'],
    ]);
    const latencyPoints = await getHistogramDataPoints(
      'genkit/feature/path/latency'
    );
    const latencyStatuses = latencyPoints.map((p) => [
      p.attributes.path,
      p.attributes.status,
    ]);
    assert.deepEqual(latencyStatuses, [
      ['/{testFlow,t:flow}/{sub-action-1,t:flowStep}', 'success'],
      ['/{testFlow,t:flow}/{sub-action-2,t:flowStep}', 'failure'],
    ]);
  }, 10000); //timeout

  it('writes user feedback metrics', async () => {
    await appendSpan(
      'trace1',
      'parent1',
      {
        name: 'user-feedback',
        path: '/{flowName}',
        metadata: {
          subtype: 'userFeedback',
          feedbackValue: 'negative',
          textFeedback: 'terrible',
        },
      },
      { [SPAN_TYPE_ATTR]: 'userEngagement' }
    );

    await getExportedSpans();
    const dataPoints = await getCounterDataPoints('genkit/engagement/feedback');

    const points = dataPoints.map((p) => [
      p.attributes.name,
      p.attributes.value,
      p.attributes.hasText,
      p.attributes.source,
    ]);
    assert.deepEqual(points, [['flowName', 'negative', true, 'ts']]);
    assert.ok(dataPoints[0].attributes.sourceVersion);
  });

  it('writes user acceptance metrics', async () => {
    await appendSpan(
      'trace1',
      'parent1',
      {
        name: 'user-acceptance',
        path: '/{flowName}',
        metadata: { subtype: 'userAcceptance', acceptanceValue: 'rejected' },
      },
      { [SPAN_TYPE_ATTR]: 'userEngagement' }
    );

    await getExportedSpans();
    const dataPoints = await getCounterDataPoints(
      'genkit/engagement/acceptance'
    );

    const points = dataPoints.map((p) => [
      p.attributes.name,
      p.attributes.value,
      p.attributes.source,
    ]);
    assert.deepEqual(points, [['flowName', 'rejected', 'ts']]);
    assert.ok(dataPoints[0].attributes.sourceVersion);
  });

  describe('Configuration', () => {
    it('should export only traces', async () => {
      const telemetry = new GcpOpenTelemetry({
        export: true,
        disableMetrics: true,
        disableTraces: false,
      } as any);
      assert.equal(telemetry['shouldExportTraces'](), true);
      assert.equal(telemetry['shouldExportMetrics'](), false);
    });

    it('should export only metrics', async () => {
      const telemetry = new GcpOpenTelemetry({
        export: true,
        disableTraces: true,
        disableMetrics: false,
      } as any);
      assert.equal(telemetry['shouldExportTraces'](), false);
      assert.equal(telemetry['shouldExportMetrics'](), true);
    });
  });

  /** Polls the in memory metric exporter until the genkit scope is found. */
  async function getGenkitMetrics(
    name: string = 'genkit',
    maxAttempts: number = 100
  ): Promise<ScopeMetrics | undefined> {
    var attempts = 0;
    while (attempts++ < maxAttempts) {
      await new Promise((resolve) => setTimeout(resolve, 50));
      const found = __getMetricExporterForTesting()
        .getMetrics()
        .find((e) => e.scopeMetrics.map((sm) => sm.scope.name).includes(name));
      if (found) {
        return found.scopeMetrics.find((e) => e.scope.name === name);
      }
    }
    assert.fail(`Waiting for metric ${name} but it has not been written.`);
  }

  /** Polls the in memory metric exporter until the genkit scope is found. */
  async function getExportedSpans(
    maxAttempts: number = 200
  ): Promise<ReadableSpan[]> {
    __forceFlushSpansForTesting();
    var attempts = 0;
    while (attempts++ < maxAttempts) {
      await new Promise((resolve) => setTimeout(resolve, 50));
      const found = __getSpanExporterForTesting().getFinishedSpans();
      if (found.length > 0) {
        return found;
      }
    }
    assert.fail(`Timed out while waiting for spans to be exported.`);
  }

  /** Finds all datapoints for a counter metric with the given name in the in memory exporter */
  async function getCounterDataPoints(
    metricName: string
  ): Promise<Array<DataPoint<number>>> {
    const genkitMetrics = await getGenkitMetrics();
    if (genkitMetrics) {
      const counterMetric: SumMetricData = genkitMetrics.metrics.find(
        (e) =>
          e.descriptor.name === metricName && e.descriptor.type === 'COUNTER'
      ) as SumMetricData;
      if (counterMetric) {
        return counterMetric.dataPoints;
      }
      assert.fail(
        `No counter metric named ${metricName} was found. Only found: ${genkitMetrics.metrics.map((e) => e.descriptor.name)}`
      );
    } else {
      assert.fail(`No genkit metrics found.`);
    }
  }

  /** Finds a counter metric with the given name in the in memory exporter */
  async function getCounterMetric(
    metricName: string
  ): Promise<DataPoint<number>> {
    const counter = await getCounterDataPoints(metricName).then((points) =>
      points.at(-1)
    );
    if (counter === undefined) {
      assert.fail(`Counter not found`);
    } else {
      return counter;
    }
  }

  /**
   * Finds all datapoints for a histogram metric with the given name in the in
   * memory exporter.
   */
  async function getHistogramDataPoints(
    metricName: string
  ): Promise<Array<DataPoint<Histogram>>> {
    const genkitMetrics = await getGenkitMetrics();
    if (genkitMetrics === undefined) {
      assert.fail('Genkit metrics not found');
    } else {
      const histogramMetric = genkitMetrics.metrics.find(
        (e) =>
          e.descriptor.name === metricName && e.descriptor.type === 'HISTOGRAM'
      ) as HistogramMetricData;
      if (histogramMetric === undefined) {
        assert.fail(
          `No histogram metric named ${metricName} was found. Only found: ${genkitMetrics.metrics.map((e) => e.descriptor.name)}`
        );
      } else {
        return histogramMetric.dataPoints;
      }
    }
  }

  /** Finds a histogram metric with the given name in the in memory exporter */
  async function getHistogramMetric(
    metricName: string
  ): Promise<DataPoint<Histogram>> {
    const metric = await getHistogramDataPoints(metricName).then((points) =>
      points.at(-1)
    );
    if (metric === undefined) {
      assert.fail('Metric not found');
    } else {
      return metric;
    }
  }

  /** Helper to create a flow with no inputs or outputs */
  function createFlow(
    ai: Genkit,
    name: string,
    fn: () => Promise<any> = async () => {}
  ) {
    return ai.defineFlow(
      {
        name,
        inputSchema: z.void(),
        outputSchema: z.any(),
      },
      fn
    );
  }

  /** Helper to create an action with no inputs or outputs */
  function createAction(
    ai: Genkit,
    name: string,
    fn: () => Promise<void> = async () => {}
  ) {
    return ai.defineTool(
      {
        name,
        description: "I don't do much...",
      },
      fn
    );
  }

  /** Helper to create a model that returns the value produced by the given
   * response function. */
  function createModel(
    ai: Genkit,
    name: string,
    respFn: () => Promise<GenerateResponseData>
  ) {
    return ai.defineModel({ name }, (req) => respFn());
  }

  function createTestModel(ai: Genkit, name: string) {
    return createModel(ai, name, async () => {
      return {
        message: {
          role: 'model',
          content: [
            {
              text: 'Oh hello',
            },
          ],
        },
        finishReason: 'stop',
        usage: {
          inputTokens: 10,
          outputTokens: 14,
          inputCharacters: 8,
          outputCharacters: 16,
          inputImages: 1,
          outputImages: 3,
        },
      };
    });
  }
});