/**
 * 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 * as assert from 'assert';
import { afterEach, beforeEach, describe, it } from 'node:test';
import * as sinon from 'sinon';
import { TextEncoder } from 'util';
import {
  FinishReason,
  HarmCategory,
  HarmProbability,
} from '../../src/common/types.js';
import { getGenkitClientHeader } from '../../src/common/utils.js';
import {
  TEST_ONLY,
  createInteraction,
  embedContent,
  generateContent,
  generateContentStream,
  getGoogleAIUrl,
  listModels,
  veoCheckOperation,
  veoPredict,
} from '../../src/googleai/client.js';
import {
  ClientOptions,
  CreateInteractionRequest,
  EmbedContentRequest,
  EmbedContentResponse,
  GenerateContentRequest,
  GenerateContentResponse,
  GenerateContentStreamResult,
  Model,
  VeoOperation,
  VeoPredictRequest,
} from '../../src/googleai/types.js';

const { getAbortSignal } = TEST_ONLY;

describe('Google AI Client', () => {
  let fetchSpy: sinon.SinonStub;
  let runInNewSpanSpy: sinon.SinonStub;
  const apiKey = 'test-api-key';
  const defaultBaseUrl = 'https://generativelanguage.googleapis.com';
  const defaultApiVersion = 'v1beta';

  beforeEach(() => {
    // In case a previous test failed and left stubs hanging, ensure a clean state
    sinon.restore();

    fetchSpy = sinon.stub(global, 'fetch');

    // We stub runInNewSpan via tracingHooks instead of importing the tracing
    // module directly. This avoids ESM immutable binding issues where Sinon
    // throws "descriptor for property 'runInNewSpan' is non-configurable".
    runInNewSpanSpy = sinon
      .stub(TEST_ONLY.tracingHooks, 'runInNewSpan')
      .callsFake((async (...args: any[]) => {
        // runInNewSpan takes either 2 or 3 arguments depending on the overload.
        // We pop the last argument, which is always the callback function.
        const fn = args.pop();
        const options = args[0];
        // Pass options.metadata to the callback so the caller can mutate it,
        // allowing us to assert on input/output fields in our tests.
        return await fn(options.metadata);
      }) as any);
  });

  afterEach(() => {
    sinon.restore();
  });

  function mockFetchResponse(
    body: any,
    ok = true,
    status = 200,
    statusText = 'OK',
    contentType = 'application/json'
  ) {
    const bodyString =
      body === null || body === undefined
        ? ''
        : contentType === 'application/json'
          ? JSON.stringify(body)
          : String(body);
    const response = new Response(bodyString, {
      status: status,
      statusText: statusText,
      headers: { 'Content-Type': contentType },
    });
    fetchSpy.resolves(response);
  }

  function createMockStream(chunks: string[]): Response {
    const stream = new ReadableStream({
      start(controller) {
        for (const chunk of chunks) {
          controller.enqueue(new TextEncoder().encode(chunk));
        }
        controller.close();
      },
    });
    return new Response(stream, {
      headers: { 'Content-Type': 'application/json' },
    });
  }

  describe('getGoogleAIUrl', () => {
    it('should build basic URL', () => {
      const url = getGoogleAIUrl({ resourcePath: 'models' });
      assert.strictEqual(url, `${defaultBaseUrl}/${defaultApiVersion}/models`);
    });

    it('should build URL with resourceMethod', () => {
      const url = getGoogleAIUrl({
        resourcePath: 'models/gemini-2.5-pro',
        resourceMethod: 'generateContent',
      });
      assert.strictEqual(
        url,
        `${defaultBaseUrl}/${defaultApiVersion}/models/gemini-2.5-pro:generateContent`
      );
    });

    it('should build URL with queryParams', () => {
      const url = getGoogleAIUrl({
        resourcePath: 'models',
        queryParams: 'pageSize=10',
      });
      assert.strictEqual(
        url,
        `${defaultBaseUrl}/${defaultApiVersion}/models?pageSize=10`
      );
    });

    it('should add alt=sse for streamGenerateContent', () => {
      const url = getGoogleAIUrl({
        resourcePath: 'models/gemini-2.5-flash',
        resourceMethod: 'streamGenerateContent',
      });
      assert.strictEqual(
        url,
        `${defaultBaseUrl}/${defaultApiVersion}/models/gemini-2.5-flash:streamGenerateContent?alt=sse`
      );
    });

    it('should add alt=sse for streamGenerateContent with other queryParams', () => {
      const url = getGoogleAIUrl({
        resourcePath: 'models/gemini-2.5-flash',
        resourceMethod: 'streamGenerateContent',
        queryParams: 'test=123',
      });
      assert.strictEqual(
        url,
        `${defaultBaseUrl}/${defaultApiVersion}/models/gemini-2.5-flash:streamGenerateContent?test=123&alt=sse`
      );
    });

    it('should use custom apiVersion and baseUrl from RequestOptions', () => {
      const clientOptions: ClientOptions = {
        apiVersion: 'v1',
        baseUrl: 'https://custom.googleapis.com',
      };
      const url = getGoogleAIUrl({ resourcePath: 'models', clientOptions });
      assert.strictEqual(url, 'https://custom.googleapis.com/v1/models');
    });
  });

  describe('listModels', () => {
    it('should return a list of models', async () => {
      const mockModels: Model[] = [
        { name: 'models/gemini-2.5-pro' } as Model,
        { name: 'models/gemini-2.5-flash' } as Model,
      ];
      mockFetchResponse({ models: mockModels });

      const result = await listModels(apiKey);
      assert.deepStrictEqual(result, mockModels);

      const expectedUrl =
        'https://generativelanguage.googleapis.com/v1beta/models?pageSize=1000';
      sinon.assert.calledOnceWithExactly(fetchSpy, expectedUrl, {
        method: 'GET',
        headers: {
          'Content-Type': 'application/json',
          'x-goog-api-key': apiKey,
          'x-goog-api-client': getGenkitClientHeader(),
        },
        redirect: 'manual',
      });
    });

    it('should throw an error if fetch fails with JSON error', async () => {
      const errorResponse = { error: { message: 'Internal Error' } };
      mockFetchResponse(errorResponse, false, 500, 'Internal Server Error');

      await assert.rejects(listModels(apiKey), (err: any) => {
        assert.strictEqual(err.name, 'GenkitError');
        assert.strictEqual(err.status, 'INTERNAL');
        assert.match(
          err.message,
          /Error fetching from .* \[500 Internal Server Error\] Internal Error/
        );
        return true;
      });
    });

    it('should throw an error if fetch fails with non-JSON error', async () => {
      mockFetchResponse(
        '<html><body><h1>Server Error</h1></body></html>',
        false,
        500,
        'Internal Server Error',
        'text/html'
      );

      await assert.rejects(listModels(apiKey), (err: any) => {
        assert.strictEqual(err.name, 'GenkitError');
        assert.strictEqual(err.status, 'INTERNAL');
        assert.match(
          err.message,
          /Error fetching from .* \[500 Internal Server Error\] <html><body><h1>Server Error<\/h1><\/body><\/html>/
        );
        return true;
      });
    });

    it('should throw an error if fetch fails with empty response body', async () => {
      mockFetchResponse(null, false, 502, 'Bad Gateway');

      await assert.rejects(listModels(apiKey), (err: any) => {
        assert.strictEqual(err.name, 'GenkitError');
        assert.strictEqual(err.status, 'UNKNOWN');
        assert.match(err.message, /Error fetching from .* \[502 Bad Gateway\]/);
        return true;
      });
    });

    it('should throw a resource exhausted error on 429', async () => {
      const errorResponse = { error: { message: 'Too many requests' } };
      mockFetchResponse(errorResponse, false, 429, 'Too Many Requests');

      await assert.rejects(listModels(apiKey), (err: any) => {
        assert.strictEqual(err.name, 'GenkitError');
        assert.strictEqual(err.status, 'RESOURCE_EXHAUSTED');
        assert.match(
          err.message,
          /Error fetching from .* \[429 Too Many Requests\] Too many requests/
        );
        return true;
      });
    });

    it('should throw an error on network failure', async () => {
      fetchSpy.rejects(new Error('Network connection failed'));

      await assert.rejects(
        listModels(apiKey),
        /Failed to fetch from .* Network connection failed/
      );
    });

    it('does not apply Interactions error-code mapping to non-Interactions URLs', async () => {
      // `safety` would map to FAILED_PRECONDITION on the Interactions path.
      mockFetchResponse(
        { error: { code: 'safety', message: 'blocked' } },
        false,
        400,
        'Bad Request'
      );
      await assert.rejects(listModels(apiKey), (err: any) => {
        assert.strictEqual(err.status, 'INVALID_ARGUMENT');
        return true;
      });
    });

    it('should include custom headers', async () => {
      mockFetchResponse({ models: [] });
      const clientOptions: ClientOptions = {
        customHeaders: { 'X-Custom-Header': 'test' },
      };
      await listModels(apiKey, clientOptions);

      sinon.assert.calledOnce(fetchSpy);
      const headers = fetchSpy.firstCall.args[1].headers;
      assert.strictEqual(headers['X-Custom-Header'], 'test');
      assert.strictEqual(headers['x-goog-api-key'], apiKey);
    });
  });

  describe('createInteraction headers', () => {
    const request: CreateInteractionRequest = {
      model: 'gemini-flash-latest',
      input: 'hi',
    };

    it('sends the Api-Revision header', async () => {
      mockFetchResponse({ status: 'completed', steps: [] });
      await createInteraction(apiKey, request);
      const headers = fetchSpy.lastCall.args[1].headers;
      assert.strictEqual(headers['Api-Revision'], '2026-05-20');
    });

    it('does not let custom headers in any case change protected headers', async () => {
      mockFetchResponse({ status: 'completed', steps: [] });
      await createInteraction(apiKey, request, {
        customHeaders: {
          'api-revision': '2025-01-01',
          'X-Goog-Api-Key': 'other-key',
          'x-custom': 'a',
        },
      });

      // fetch merges case variants of a header, so check the merged value.
      const headers = new Headers(fetchSpy.lastCall.args[1].headers);
      assert.strictEqual(headers.get('api-revision'), '2026-05-20');
      assert.strictEqual(headers.get('x-goog-api-key'), apiKey);
      assert.strictEqual(headers.get('x-custom'), 'a');
    });
  });

  describe('Interactions API HTTP errors', () => {
    const request = { model: 'gemini-flash-latest', input: 'hi' } as any;

    async function expectStatus(
      body: any,
      httpStatus: number,
      statusText: string,
      expected: string
    ) {
      mockFetchResponse(body, false, httpStatus, statusText);
      await assert.rejects(createInteraction(apiKey, request), (err: any) => {
        assert.strictEqual(err.name, 'GenkitError');
        assert.strictEqual(err.status, expected);
        return true;
      });
    }

    it('maps a pre-execution safety block (400 safety) to FAILED_PRECONDITION', async () => {
      await expectStatus(
        {
          error: {
            code: 'safety',
            message:
              'Request blocked due to safety violations (harmful content). Please modify your input and retry.',
          },
        },
        400,
        'Bad Request',
        'FAILED_PRECONDITION'
      );
    });

    it('maps 404 model_not_found to NOT_FOUND', async () => {
      await expectStatus(
        { error: { code: 'model_not_found', message: 'no such model' } },
        404,
        'Not Found',
        'NOT_FOUND'
      );
    });

    it('maps 429 quota_exceeded to RESOURCE_EXHAUSTED', async () => {
      await expectStatus(
        { error: { code: 'quota_exceeded', message: 'quota' } },
        429,
        'Too Many Requests',
        'RESOURCE_EXHAUSTED'
      );
    });

    it('maps 401 authentication to UNAUTHENTICATED', async () => {
      await expectStatus(
        { error: { code: 'authentication', message: 'bad key' } },
        401,
        'Unauthorized',
        'UNAUTHENTICATED'
      );
    });

    it('falls back to the HTTP status for an unrecognized code', async () => {
      await expectStatus(
        { error: { code: 'something_new', message: 'huh' } },
        403,
        'Forbidden',
        'PERMISSION_DENIED'
      );
    });

    it('detects Interactions URLs by path segment', () => {
      const { isInteractionUrl } = TEST_ONLY;
      assert.strictEqual(
        isInteractionUrl(`${defaultBaseUrl}/v1beta/interactions`),
        true
      );
      assert.strictEqual(
        isInteractionUrl(`${defaultBaseUrl}/v1beta/interactions/abc/cancel`),
        true
      );
      assert.strictEqual(
        isInteractionUrl(
          `${defaultBaseUrl}/v1beta/models/gemini-2.5-flash:generateContent`
        ),
        false
      );
      assert.strictEqual(
        isInteractionUrl(`${defaultBaseUrl}/v1beta/models?q=interactions`),
        false
      );
    });
  });

  describe('generateContent', () => {
    const model = 'gemini-flash-latest';
    const request: GenerateContentRequest = {
      contents: [{ role: 'user', parts: [{ text: 'hello' }] }],
    };

    it('should return GenerateContentResponse', async () => {
      const mockResponse: GenerateContentResponse = {
        candidates: [
          { index: 0, content: { role: 'model', parts: [{ text: 'world' }] } },
        ],
      };
      mockFetchResponse(mockResponse);

      const result = await generateContent(apiKey, model, request);
      assert.deepStrictEqual(result, mockResponse);

      const expectedUrl = `https://generativelanguage.googleapis.com/v1beta/models/${model}:generateContent`;
      sinon.assert.calledOnceWithExactly(fetchSpy, expectedUrl, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'x-goog-api-key': apiKey,
          'x-goog-api-client': getGenkitClientHeader(),
        },
        redirect: 'manual',
        body: JSON.stringify(request),
      });
      // Should not trace by default
      sinon.assert.notCalled(runInNewSpanSpy);
    });

    it('should trace generateContent if experimental_debugTraces is true', async () => {
      const mockResponse: GenerateContentResponse = {
        candidates: [
          { index: 0, content: { role: 'model', parts: [{ text: 'world' }] } },
        ],
      };
      mockFetchResponse(mockResponse);

      const clientOptions: ClientOptions = { experimental_debugTraces: true };
      const result = await generateContent(
        apiKey,
        model,
        request,
        clientOptions
      );
      assert.deepStrictEqual(result, mockResponse);

      sinon.assert.calledOnce(runInNewSpanSpy);

      const traceMetadata = runInNewSpanSpy.firstCall.args[0].metadata;
      assert.strictEqual(traceMetadata.name, 'httpRequest');

      assert.deepStrictEqual(traceMetadata.input, {
        apiEndpoint: `https://generativelanguage.googleapis.com/v1beta/models/${model}:generateContent`,
        request: request,
        model: model,
        headers: {
          'Content-Type': 'application/json',
          'x-goog-api-key': '<REDACTED> (12 characters)', // "test-api-key"
          'x-goog-api-client': getGenkitClientHeader(),
        },
      });
      assert.deepStrictEqual(traceMetadata.output, mockResponse);
    });

    it('should throw on API error with JSON body', async () => {
      const errorResponse = { error: { message: 'Invalid Request' } };
      mockFetchResponse(errorResponse, false, 400, 'Bad Request');

      await assert.rejects(
        generateContent(apiKey, model, request),
        (err: any) => {
          assert.strictEqual(err.name, 'GenkitError');
          assert.strictEqual(err.status, 'INVALID_ARGUMENT');
          assert.match(
            err.message,
            /Error fetching from .* \[400 Bad Request\] Invalid Request/
          );
          return true;
        }
      );
    });

    it('should throw on API error with non-JSON body', async () => {
      mockFetchResponse('Bad Request', false, 400, 'Bad Request', 'text/plain');

      await assert.rejects(
        generateContent(apiKey, model, request),
        (err: any) => {
          assert.strictEqual(err.name, 'GenkitError');
          assert.strictEqual(err.status, 'INVALID_ARGUMENT');
          assert.match(
            err.message,
            /Error fetching from .* \[400 Bad Request\] Bad Request/
          );
          return true;
        }
      );
    });

    it('should throw on network failure', async () => {
      fetchSpy.rejects(new TypeError('Failed to fetch'));
      await assert.rejects(
        generateContent(apiKey, model, request),
        /Failed to fetch from .* Failed to fetch/
      );
    });
  });

  describe('embedContent', () => {
    const model = 'text-embedding-005';
    const request: EmbedContentRequest = {
      content: { role: 'user', parts: [{ text: 'test content' }] },
    };

    it('should return EmbedContentResponse', async () => {
      const mockResponse: EmbedContentResponse = {
        embedding: { values: [0.1, 0.2, 0.3] },
      };
      mockFetchResponse(mockResponse);

      const result = await embedContent(apiKey, model, request);
      assert.deepStrictEqual(result, mockResponse);

      const expectedUrl = `https://generativelanguage.googleapis.com/v1beta/models/${model}:embedContent`;
      sinon.assert.calledOnceWithExactly(fetchSpy, expectedUrl, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'x-goog-api-key': apiKey,
          'x-goog-api-client': getGenkitClientHeader(),
        },
        redirect: 'manual',
        body: JSON.stringify(request),
      });
      sinon.assert.notCalled(runInNewSpanSpy);
    });

    it('should trace embedContent if experimental_debugTraces is true', async () => {
      const mockResponse: EmbedContentResponse = {
        embedding: { values: [0.1, 0.2, 0.3] },
      };
      mockFetchResponse(mockResponse);

      const clientOptions: ClientOptions = { experimental_debugTraces: true };
      const result = await embedContent(apiKey, model, request, clientOptions);
      assert.deepStrictEqual(result, mockResponse);

      sinon.assert.calledOnce(runInNewSpanSpy);

      const traceMetadata = runInNewSpanSpy.firstCall.args[0].metadata;
      assert.strictEqual(traceMetadata.name, 'httpRequest');

      assert.deepStrictEqual(traceMetadata.input, {
        apiEndpoint: `https://generativelanguage.googleapis.com/v1beta/models/${model}:embedContent`,
        request: request,
        model: model,
        headers: {
          'Content-Type': 'application/json',
          'x-goog-api-key': '<REDACTED> (12 characters)',
          'x-goog-api-client': getGenkitClientHeader(),
        },
      });
      assert.deepStrictEqual(traceMetadata.output, mockResponse);
    });

    it('should throw on API error with JSON body', async () => {
      const errorResponse = { error: { message: 'Embedding failed' } };
      mockFetchResponse(errorResponse, false, 500, 'Internal Server Error');

      await assert.rejects(embedContent(apiKey, model, request), (err: any) => {
        assert.strictEqual(err.name, 'GenkitError');
        assert.strictEqual(err.status, 'INTERNAL');
        assert.match(
          err.message,
          /Error fetching from .* \[500 Internal Server Error\] Embedding failed/
        );
        return true;
      });
    });

    it('should throw on API error with non-JSON body', async () => {
      mockFetchResponse(
        'Internal Server Error',
        false,
        500,
        'Internal Server Error',
        'text/plain'
      );

      await assert.rejects(embedContent(apiKey, model, request), (err: any) => {
        assert.strictEqual(err.name, 'GenkitError');
        assert.strictEqual(err.status, 'INTERNAL');
        assert.match(
          err.message,
          /Error fetching from .* \[500 Internal Server Error\] Internal Server Error/
        );
        return true;
      });
    });

    it('should throw on network failure', async () => {
      fetchSpy.rejects(new TypeError('Network error'));
      await assert.rejects(
        embedContent(apiKey, model, request),
        /Failed to fetch from .* Network error/
      );
    });
  });

  describe('generateContentStream', () => {
    const model = 'gemini-2.5-flash';
    const defaultRequest: GenerateContentRequest = {
      contents: [{ role: 'user', parts: [{ text: 'stream test' }] }],
    };

    it('should trace generateContentStream if experimental_debugTraces is true', async () => {
      const chunks = [
        'data: {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": "Hello "}]}}]}\n\n',
      ];
      fetchSpy.resolves(createMockStream(chunks));

      const clientOptions: ClientOptions = { experimental_debugTraces: true };
      const result = await generateContentStream(
        apiKey,
        model,
        defaultRequest,
        clientOptions
      );

      const streamResults: GenerateContentResponse[] = [];
      for await (const item of result.stream) {
        streamResults.push(item);
      }

      sinon.assert.calledOnce(runInNewSpanSpy);

      const traceMetadata = runInNewSpanSpy.firstCall.args[0].metadata;
      assert.strictEqual(traceMetadata.name, 'httpRequest');

      assert.deepStrictEqual(traceMetadata.input, {
        apiEndpoint: `https://generativelanguage.googleapis.com/v1beta/models/${model}:streamGenerateContent?alt=sse`,
        request: defaultRequest,
        model: model,
        headers: {
          'Content-Type': 'application/json',
          'x-goog-api-key': '<REDACTED> (12 characters)',
          'x-goog-api-client': getGenkitClientHeader(),
        },
      });
      assert.strictEqual(traceMetadata.output, '[Streaming Response]');
    });

    it('should process stream and return stream and aggregated response', async () => {
      const chunks = [
        'data: {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": "Hello "}]}}]}\n\n',
        'data: {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": "World!"}]}}], "usageMetadata": {"totalTokenCount": 10}}\n\n',
      ];
      fetchSpy.resolves(createMockStream(chunks));

      const result: GenerateContentStreamResult = await generateContentStream(
        apiKey,
        model,
        defaultRequest
      );

      const streamResults: GenerateContentResponse[] = [];
      for await (const item of result.stream) {
        streamResults.push(item);
      }

      assert.deepStrictEqual(streamResults, [
        {
          candidates: [
            {
              index: 0,
              content: { role: 'model', parts: [{ text: 'Hello ' }] },
            },
          ],
        },
        {
          candidates: [
            {
              index: 0,
              content: { role: 'model', parts: [{ text: 'World!' }] },
            },
          ],
          usageMetadata: { totalTokenCount: 10 },
        },
      ]);

      const aggregated = await result.response;
      assert.deepStrictEqual(aggregated, {
        candidates: [
          {
            index: 0,
            content: {
              role: 'model',
              parts: [{ text: 'Hello ' }, { text: 'World!' }],
            },
          },
        ],
        usageMetadata: { totalTokenCount: 10 },
      });

      const expectedUrl = `https://generativelanguage.googleapis.com/v1beta/models/${model}:streamGenerateContent?alt=sse`;
      sinon.assert.calledOnceWithExactly(fetchSpy, expectedUrl, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'x-goog-api-key': apiKey,
          'x-goog-api-client': getGenkitClientHeader(),
        },
        redirect: 'manual',
        body: JSON.stringify(defaultRequest),
      });
    });

    it('should aggregate parts for multiple candidates', async () => {
      const chunks = [
        'data: {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": "C0 A"}]}}, {"index": 1, "content": {"role": "model", "parts": [{"text": "C1 A"}]}}]}\n\n',
        'data: {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": " C0 B"}]}}, {"index": 1, "content": {"role": "model", "parts": [{"text": " C1 B"}]}}]}\n\n',
      ];
      fetchSpy.resolves(createMockStream(chunks));
      const result = await generateContentStream(apiKey, model, defaultRequest);
      const aggregated = await result.response;

      assert.strictEqual(aggregated.candidates?.length, 2);
      const sortedCandidates = aggregated.candidates!.sort(
        (a, b) => a.index - b.index
      );

      assert.deepStrictEqual(sortedCandidates[0], {
        index: 0,
        content: {
          role: 'model',
          parts: [{ text: 'C0 A' }, { text: ' C0 B' }],
        },
      });
      assert.deepStrictEqual(sortedCandidates[1], {
        index: 1,
        content: {
          role: 'model',
          parts: [{ text: 'C1 A' }, { text: ' C1 B' }],
        },
      });
    });

    it('should use latest metadata from chunks for each candidate', async () => {
      const chunks = [
        'data: {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": "A"}]}, "finishReason": "STOP"}, {"index": 1, "content": {"role": "model", "parts": [{"text": "C1 "}]}}]}\n\n',
        'data: {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": "B"}]}, "finishReason": "MAX_TOKENS", "safetyRatings": [{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "probability": "LOW"}]}, {"index": 1, "content": {"role": "model", "parts": [{"text": "D1"}]}, "finishReason": "STOP" }], "usageMetadata": {"totalTokenCount": 20}}\n\n',
      ];
      fetchSpy.resolves(createMockStream(chunks));
      const result = await generateContentStream(apiKey, model, defaultRequest);
      const aggregated = await result.response;

      assert.strictEqual(aggregated.candidates?.length, 2);
      const sortedCandidates = aggregated.candidates!.sort(
        (a, b) => a.index - b.index
      );

      const cand0 = sortedCandidates[0];
      assert.strictEqual(cand0.index, 0);
      assert.strictEqual(cand0.finishReason, FinishReason.MAX_TOKENS);
      assert.deepStrictEqual(cand0.safetyRatings, [
        {
          category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
          probability: HarmProbability.LOW,
        },
      ]);
      assert.deepStrictEqual(cand0.content.parts, [
        { text: 'A' },
        { text: 'B' },
      ]);

      const cand1 = sortedCandidates[1];
      assert.strictEqual(cand1.index, 1);
      assert.strictEqual(cand1.finishReason, FinishReason.STOP);
      assert.strictEqual(cand1.safetyRatings, undefined);
      assert.deepStrictEqual(cand1.content.parts, [
        { text: 'C1 ' },
        { text: 'D1' },
      ]);

      assert.deepStrictEqual(aggregated.usageMetadata, { totalTokenCount: 20 });
    });

    it('should aggregate different part types', async () => {
      const chunks = [
        'data: {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": "Action: "}]}}]}\n\n',
        'data: {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"functionCall": {"name": "tool_call", "args": {}}}]}}]}\n\n',
      ];
      fetchSpy.resolves(createMockStream(chunks));
      const result = await generateContentStream(apiKey, model, defaultRequest);
      const aggregated = await result.response;
      const cand = aggregated.candidates![0];

      assert.deepStrictEqual(cand, {
        index: 0,
        content: {
          role: 'model',
          parts: [
            { text: 'Action: ' },
            { functionCall: { name: 'tool_call', args: {} } },
          ],
        },
      });
    });

    it('should handle stream with malformed JSON', async () => {
      const chunks = [
        'data: {"candidates": [{"index": 0, "content": {"role": "model", "parts": [{"text": "Hello "}]}}]}\n\n',
        'data: {"candi dates": []}}}}\n\n', // Malformed
      ];
      fetchSpy.resolves(createMockStream(chunks));

      const result: GenerateContentStreamResult = await generateContentStream(
        apiKey,
        model,
        defaultRequest
      );

      const streamResults: GenerateContentResponse[] = [];
      try {
        for await (const item of result.stream) {
          streamResults.push(item);
        }
        assert.fail('Stream should have thrown an error');
      } catch (e: any) {
        assert.match(e.message, /Error parsing JSON response/);
      }

      try {
        await result.response;
        assert.fail('Response promise should have thrown an error');
      } catch (e: any) {
        assert.match(e.message, /Error parsing JSON response/);
      }
    });

    it('should throw error if fetch rejects for stream', async () => {
      fetchSpy.rejects(new Error('Stream failed to connect'));

      await assert.rejects(
        generateContentStream(apiKey, model, defaultRequest),
        /Failed to fetch from .* Stream failed to connect/
      );
    });

    it('should respect AbortSignal', async () => {
      const controller = new AbortController(); // The test's controller
      const clientOptions: ClientOptions = { signal: controller.signal };
      const model = 'gemini-2.5-flash';
      const defaultRequest: GenerateContentRequest = {
        contents: [{ role: 'user', parts: [{ text: 'stream test' }] }],
      };
      const apiKey = 'test-api-key';

      let capturedSignal: AbortSignal | undefined;

      fetchSpy.callsFake((url, options) => {
        return new Promise((resolve, reject) => {
          const signal = options?.signal;
          capturedSignal = signal; // Capture the signal passed to fetch

          if (signal) {
            if (signal.aborted) {
              return reject(
                new DOMException('The operation was aborted.', 'AbortError')
              );
            }
            signal.addEventListener(
              'abort',
              () => {
                reject(
                  new DOMException('The operation was aborted.', 'AbortError')
                );
              },
              { once: true }
            );
          }
          // Keep promise pending to simulate a real network request
        });
      });

      const promise = generateContentStream(
        apiKey,
        model,
        defaultRequest,
        clientOptions
      );

      // Short delay to allow fetchSpy to be called
      await new Promise((resolve) => setTimeout(resolve, 0));

      sinon.assert.calledOnce(fetchSpy);
      assert.ok(capturedSignal, 'A signal should be passed to fetch');
      assert.notStrictEqual(
        capturedSignal,
        controller.signal,
        'The signal passed to fetch should be a new signal, not the original one'
      );
      assert.strictEqual(
        capturedSignal!.aborted,
        false,
        'The passed signal should not be aborted initially'
      );

      // Abort the test's controller
      controller.abort();

      // Verify the promise from generateContentStream rejects due to the abort
      await assert.rejects(promise, (err: Error) => {
        assert.match(
          err.message,
          /Failed to fetch from .* The operation was aborted./
        );
        return true;
      });

      // Verify the signal captured within fetchSpy is now aborted
      assert.strictEqual(
        capturedSignal!.aborted,
        true,
        'The passed signal should be aborted after the original controller is aborted'
      );
    });
  });

  describe('veoPredict', () => {
    const model = 'veo-1.5-flash-0804';
    const request: VeoPredictRequest = {
      instances: [{ prompt: 'A video of a sunset' }],
      parameters: {},
    };

    it('should return VeoOperation', async () => {
      const mockResponse: VeoOperation = {
        name: 'operations/12345',
        done: false,
      };
      mockFetchResponse(mockResponse);

      const result = await veoPredict(apiKey, model, request);
      assert.deepStrictEqual(result, mockResponse);

      const expectedUrl = `https://generativelanguage.googleapis.com/v1beta/models/${model}:predictLongRunning`;
      sinon.assert.calledOnceWithExactly(fetchSpy, expectedUrl, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'x-goog-api-key': apiKey,
          'x-goog-api-client': getGenkitClientHeader(),
        },
        redirect: 'manual',
        body: JSON.stringify(request),
      });
    });

    it('should throw on API error', async () => {
      const errorResponse = { error: { message: 'Veo failed to start' } };
      mockFetchResponse(errorResponse, false, 500, 'Internal Server Error');

      await assert.rejects(veoPredict(apiKey, model, request), (err: any) => {
        assert.strictEqual(err.name, 'GenkitError');
        assert.strictEqual(err.status, 'INTERNAL');
        assert.match(
          err.message,
          /Error fetching from .* \[500 Internal Server Error\] Veo failed to start/
        );
        return true;
      });
    });
  });

  describe('veoCheckOperation', () => {
    const operationName = 'operations/12345';

    it('should return VeoOperation state', async () => {
      const mockResponse: VeoOperation = {
        name: operationName,
        done: true,
        response: {
          generateVideoResponse: {
            generatedSamples: [{ video: { uri: 'gs://bucket/video.mp4' } }],
          },
        },
      };
      mockFetchResponse(mockResponse);

      const result = await veoCheckOperation(apiKey, operationName);
      assert.deepStrictEqual(result, mockResponse);

      const expectedUrl = `https://generativelanguage.googleapis.com/v1beta/${operationName}`;
      sinon.assert.calledOnceWithExactly(fetchSpy, expectedUrl, {
        method: 'GET',
        headers: {
          'Content-Type': 'application/json',
          'x-goog-api-key': apiKey,
          'x-goog-api-client': getGenkitClientHeader(),
        },
        redirect: 'manual',
      });
    });

    it('should throw on API error', async () => {
      const errorResponse = { error: { message: 'Operation not found' } };
      mockFetchResponse(errorResponse, false, 404, 'Not Found');

      await assert.rejects(
        veoCheckOperation(apiKey, operationName),
        (err: any) => {
          assert.strictEqual(err.name, 'GenkitError');
          assert.strictEqual(err.status, 'UNKNOWN');
          assert.match(
            err.message,
            /Error fetching from .* \[404 Not Found\] Operation not found/
          );
          return true;
        }
      );
    });
  });

  describe('TEST_ONLY.getAbortSignal', () => {
    let clock: sinon.SinonFakeTimers;

    beforeEach(() => {
      clock = sinon.useFakeTimers();
    });

    afterEach(() => {
      clock.restore();
    });

    it('should return undefined if no signal or timeout', () => {
      const signal = getAbortSignal({});
      assert.strictEqual(signal, undefined);
    });

    it('should return a signal that aborts after timeout', () => {
      const clientOptions: ClientOptions = { timeout: 100 };
      const signal = getAbortSignal(clientOptions);
      assert.ok(signal);
      assert.strictEqual(signal!.aborted, false);
      clock.tick(100);
      assert.strictEqual(signal!.aborted, true);
    });

    it('should return a signal linked to the provided signal', () => {
      const controller = new AbortController();
      const clientOptions: ClientOptions = { signal: controller.signal };
      const signal = getAbortSignal(clientOptions);
      assert.ok(signal);
      assert.strictEqual(signal!.aborted, false);
      controller.abort();
      assert.strictEqual(signal!.aborted, true);
    });

    it('should return a signal that aborts on timeout or provided signal', () => {
      const controller = new AbortController();
      const clientOptions: ClientOptions = {
        signal: controller.signal,
        timeout: 100,
      };
      const signal = getAbortSignal(clientOptions);
      assert.ok(signal);

      clock.tick(50);
      assert.strictEqual(signal!.aborted, false);
      controller.abort();
      assert.strictEqual(signal!.aborted, true);

      const signal2 = getAbortSignal(clientOptions);
      clock.tick(100);
      assert.strictEqual(signal2!.aborted, true);
    });
  });
});