/**
 * 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 type { GenerateContentCandidate } from '@google-cloud/vertexai';
import * as assert from 'assert';
import { z, type MessageData } from 'genkit';
import { toJsonSchema } from 'genkit/schema';
import { describe, it } from 'node:test';
import {
  cleanSchema,
  fromGeminiCandidate,
  toGeminiMessage,
  toGeminiSystemInstruction,
  toGeminiTool,
} from '../src/gemini.js';

describe('toGeminiMessages', () => {
  const testCases = [
    {
      should: 'should transform genkit message (text content) correctly',
      inputMessage: {
        role: 'user',
        content: [{ text: 'Tell a joke about dogs.' }],
      },
      expectedOutput: {
        role: 'user',
        parts: [{ text: 'Tell a joke about dogs.' }],
      },
    },
    {
      should:
        'should transform genkit message (tool response content with ref) correctly',
      inputMessage: {
        role: 'tool',
        content: [
          {
            toolResponse: {
              name: 'tellAFunnyJoke',
              output: 'Why did the dogs cross the road?',
              ref: '1',
            },
          },
          {
            toolResponse: {
              name: 'tellAnotherFunnyJoke',
              output: 'To get to the other side.',
              ref: '0',
            },
          },
        ],
      },
      expectedOutput: {
        role: 'function',
        parts: [
          {
            functionResponse: {
              name: 'tellAnotherFunnyJoke',
              response: {
                name: 'tellAnotherFunnyJoke',
                content: 'To get to the other side.',
              },
            },
          },
          {
            functionResponse: {
              name: 'tellAFunnyJoke',
              response: {
                name: 'tellAFunnyJoke',
                content: 'Why did the dogs cross the road?',
              },
            },
          },
        ],
      },
    },
    {
      should:
        'should transform genkit message (tool request content) correctly',
      inputMessage: {
        role: 'model',
        content: [
          { toolRequest: { name: 'tellAFunnyJoke', input: { topic: 'dogs' } } },
        ],
      },
      expectedOutput: {
        role: 'model',
        parts: [
          { functionCall: { name: 'tellAFunnyJoke', args: { topic: 'dogs' } } },
        ],
      },
    },
    {
      should:
        'should transform genkit message (tool response content) correctly',
      inputMessage: {
        role: 'tool',
        content: [
          {
            toolResponse: {
              name: 'tellAFunnyJoke',
              output: 'Why did the dogs cross the road?',
            },
          },
        ],
      },
      expectedOutput: {
        role: 'function',
        parts: [
          {
            functionResponse: {
              name: 'tellAFunnyJoke',
              response: {
                name: 'tellAFunnyJoke',
                content: 'Why did the dogs cross the road?',
              },
            },
          },
        ],
      },
    },
    {
      should:
        'should transform genkit message (inline base64 image content) correctly',
      inputMessage: {
        role: 'user',
        content: [
          { text: 'describe the following image:' },
          {
            media: {
              contentType: 'image/jpeg',
              url: 'data:image/jpeg;base64, ',
            },
          },
        ],
      },
      expectedOutput: {
        role: 'user',
        parts: [
          { text: 'describe the following image:' },
          {
            inlineData: {
              mimeType: 'image/jpeg',
              data: ' ',
            },
          },
        ],
      },
    },
    {
      should: 'should re-populate thoughtSignature from reasoning metadata',
      inputMessage: {
        role: 'model',
        content: [{ reasoning: '', metadata: { thoughtSignature: 'abc123' } }],
      },
      expectedOutput: {
        role: 'model',
        parts: [{ thought: true, thoughtSignature: 'abc123' }],
      },
    },
  ];
  for (const test of testCases) {
    it(test.should, () => {
      assert.deepStrictEqual(
        toGeminiMessage(test.inputMessage as MessageData),
        test.expectedOutput
      );
    });
  }
});

describe('toGeminiSystemInstruction', () => {
  const testCases = [
    {
      should: 'should transform from system to user',
      inputMessage: {
        role: 'system',
        content: [{ text: 'You are an expert in all things cats.' }],
      },
      expectedOutput: {
        role: 'user',
        parts: [{ text: 'You are an expert in all things cats.' }],
      },
    },
    {
      should: 'should transform from system to user with multiple parts',
      inputMessage: {
        role: 'system',
        content: [
          { text: 'You are an expert in all things animals.' },
          { text: 'You love cats.' },
        ],
      },
      expectedOutput: {
        role: 'user',
        parts: [
          { text: 'You are an expert in all things animals.' },
          { text: 'You love cats.' },
        ],
      },
    },
  ];
  for (const test of testCases) {
    it(test.should, () => {
      assert.deepStrictEqual(
        toGeminiSystemInstruction(test.inputMessage as MessageData),
        test.expectedOutput
      );
    });
  }
});

describe('fromGeminiCandidate', () => {
  const testCases = [
    {
      should:
        'should transform gemini candidate to genkit candidate (text parts) correctly',
      // had to delete the probabilityScore, severity, severityScore for the HARM_CATEGORY_SEXUALLY_EXPLICIT safety rating category
      geminiCandidate: {
        content: {
          role: 'model',
          parts: [
            {
              text: 'Why did the dog go to the bank?\n\nTo get his bones cashed!',
            },
          ],
        },
        finishReason: 'STOP',
        safetyRatings: [
          {
            category: 'HARM_CATEGORY_HATE_SPEECH',
            probability: 'NEGLIGIBLE',
            probabilityScore: 0.12074952,
            severity: 'HARM_SEVERITY_NEGLIGIBLE',
            severityScore: 0.18388656,
          },
          {
            category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
            probability: 'NEGLIGIBLE',
            probabilityScore: 0.37874627,
            severity: 'HARM_SEVERITY_LOW',
            severityScore: 0.37227696,
          },
          {
            category: 'HARM_CATEGORY_HARASSMENT',
            probability: 'NEGLIGIBLE',
            probabilityScore: 0.3983479,
            severity: 'HARM_SEVERITY_LOW',
            severityScore: 0.22270013,
          },
          {
            category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
            probability: 'NEGLIGIBLE',
          },
        ],
      },
      expectedOutput: {
        index: 0,
        message: {
          role: 'model',
          content: [
            {
              text: 'Why did the dog go to the bank?\n\nTo get his bones cashed!',
            },
          ],
        },
        finishReason: 'stop',
        finishMessage: undefined,
        custom: {
          citationMetadata: undefined,
          safetyRatings: [
            {
              category: 'HARM_CATEGORY_HATE_SPEECH',
              probability: 'NEGLIGIBLE',
              probabilityScore: 0.12074952,
              severity: 'HARM_SEVERITY_NEGLIGIBLE',
              severityScore: 0.18388656,
            },
            {
              category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
              probability: 'NEGLIGIBLE',
              probabilityScore: 0.37874627,
              severity: 'HARM_SEVERITY_LOW',
              severityScore: 0.37227696,
            },
            {
              category: 'HARM_CATEGORY_HARASSMENT',
              probability: 'NEGLIGIBLE',
              probabilityScore: 0.3983479,
              severity: 'HARM_SEVERITY_LOW',
              severityScore: 0.22270013,
            },
            {
              category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
              probability: 'NEGLIGIBLE',
            },
          ],
        },
      },
    },
    {
      should:
        'should transform gemini candidate to genkit candidate (function call parts) correctly',
      geminiCandidate: {
        content: {
          role: 'model',
          parts: [
            {
              functionCall: { name: 'tellAFunnyJoke', args: { topic: 'dog' } },
            },
          ],
        },
        finishReason: 'STOP',
        safetyRatings: [
          {
            category: 'HARM_CATEGORY_HATE_SPEECH',
            probability: 'NEGLIGIBLE',
            probabilityScore: 0.11858909,
            severity: 'HARM_SEVERITY_NEGLIGIBLE',
            severityScore: 0.11456649,
          },
          {
            category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
            probability: 'NEGLIGIBLE',
            probabilityScore: 0.13857833,
            severity: 'HARM_SEVERITY_NEGLIGIBLE',
            severityScore: 0.11417085,
          },
          {
            category: 'HARM_CATEGORY_HARASSMENT',
            probability: 'NEGLIGIBLE',
            probabilityScore: 0.28012377,
            severity: 'HARM_SEVERITY_NEGLIGIBLE',
            severityScore: 0.112405084,
          },
          {
            category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
            probability: 'NEGLIGIBLE',
          },
        ],
      },
      expectedOutput: {
        index: 0,
        message: {
          role: 'model',
          content: [
            {
              toolRequest: {
                name: 'tellAFunnyJoke',
                input: { topic: 'dog' },
                ref: '0',
              },
            },
          ],
        },
        finishReason: 'stop',
        finishMessage: undefined,
        custom: {
          citationMetadata: undefined,
          safetyRatings: [
            {
              category: 'HARM_CATEGORY_HATE_SPEECH',
              probability: 'NEGLIGIBLE',
              probabilityScore: 0.11858909,
              severity: 'HARM_SEVERITY_NEGLIGIBLE',
              severityScore: 0.11456649,
            },
            {
              category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
              probability: 'NEGLIGIBLE',
              probabilityScore: 0.13857833,
              severity: 'HARM_SEVERITY_NEGLIGIBLE',
              severityScore: 0.11417085,
            },
            {
              category: 'HARM_CATEGORY_HARASSMENT',
              probability: 'NEGLIGIBLE',
              probabilityScore: 0.28012377,
              severity: 'HARM_SEVERITY_NEGLIGIBLE',
              severityScore: 0.112405084,
            },
            {
              category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
              probability: 'NEGLIGIBLE',
            },
          ],
        },
      },
    },
    {
      should:
        'should transform gemini candidate to genkit candidate (thought parts) correctly',
      geminiCandidate: {
        content: {
          role: 'model',
          parts: [
            {
              thought: true,
              thoughtSignature: 'abc123',
            },
            {
              thought: true,
              text: 'thought with text',
              thoughtSignature: 'def456',
            },
          ],
        },
        finishReason: 'STOP',
        safetyRatings: [
          {
            category: 'HARM_CATEGORY_HATE_SPEECH',
            probability: 'NEGLIGIBLE',
            probabilityScore: 0.11858909,
            severity: 'HARM_SEVERITY_NEGLIGIBLE',
            severityScore: 0.11456649,
          },
          {
            category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
            probability: 'NEGLIGIBLE',
            probabilityScore: 0.13857833,
            severity: 'HARM_SEVERITY_NEGLIGIBLE',
            severityScore: 0.11417085,
          },
          {
            category: 'HARM_CATEGORY_HARASSMENT',
            probability: 'NEGLIGIBLE',
            probabilityScore: 0.28012377,
            severity: 'HARM_SEVERITY_NEGLIGIBLE',
            severityScore: 0.112405084,
          },
          {
            category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
            probability: 'NEGLIGIBLE',
          },
        ],
      },
      expectedOutput: {
        index: 0,
        message: {
          role: 'model',
          content: [
            {
              reasoning: '',
              metadata: { thoughtSignature: 'abc123' },
            },
            {
              reasoning: 'thought with text',
              metadata: { thoughtSignature: 'def456' },
            },
          ],
        },
        finishReason: 'stop',
        finishMessage: undefined,
        custom: {
          citationMetadata: undefined,
          safetyRatings: [
            {
              category: 'HARM_CATEGORY_HATE_SPEECH',
              probability: 'NEGLIGIBLE',
              probabilityScore: 0.11858909,
              severity: 'HARM_SEVERITY_NEGLIGIBLE',
              severityScore: 0.11456649,
            },
            {
              category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
              probability: 'NEGLIGIBLE',
              probabilityScore: 0.13857833,
              severity: 'HARM_SEVERITY_NEGLIGIBLE',
              severityScore: 0.11417085,
            },
            {
              category: 'HARM_CATEGORY_HARASSMENT',
              probability: 'NEGLIGIBLE',
              probabilityScore: 0.28012377,
              severity: 'HARM_SEVERITY_NEGLIGIBLE',
              severityScore: 0.112405084,
            },
            {
              category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
              probability: 'NEGLIGIBLE',
            },
          ],
        },
      },
    },
  ];
  for (const test of testCases) {
    it(test.should, () => {
      assert.deepStrictEqual(
        fromGeminiCandidate(
          test.geminiCandidate as GenerateContentCandidate,
          false
        ),
        test.expectedOutput
      );
    });
  }
});

describe('cleanSchema', () => {
  it('strips nulls from type', () => {
    const cleaned = cleanSchema({
      type: 'object',
      properties: {
        title: {
          type: 'string',
        },
        subtitle: {
          type: ['string', 'null'],
        },
      },
      required: ['title'],
      additionalProperties: true,
      $schema: 'http://json-schema.org/draft-07/schema#',
    });

    assert.deepStrictEqual(cleaned, {
      type: 'object',
      properties: {
        title: {
          type: 'string',
        },
        subtitle: {
          type: 'string',
        },
      },
      required: ['title'],
    });
  });
});

describe('toGeminiTool', () => {
  it('converts to gemini tool', async () => {
    const got = toGeminiTool({
      name: 'foo',
      description: 'tool foo',
      inputSchema: toJsonSchema({
        schema: z.object({
          simpleString: z.string().describe('a string').nullable(),
          simpleNumber: z.number().describe('a number'),
          simpleBoolean: z.boolean().describe('a boolean').optional(),
          simpleArray: z.array(z.string()).describe('an array').optional(),
          simpleEnum: z
            .enum(['choice_a', 'choice_b'])
            .describe('an enum')
            .optional(),
        }),
      }),
    });

    const want = {
      description: 'tool foo',
      name: 'foo',
      parameters: {
        properties: {
          simpleArray: {
            description: 'an array',
            items: {
              type: 'STRING',
            },
            type: 'ARRAY',
          },
          simpleBoolean: {
            description: 'a boolean',
            type: 'BOOLEAN',
          },
          simpleEnum: {
            description: 'an enum',
            enum: ['choice_a', 'choice_b'],
            type: 'STRING',
          },
          simpleNumber: {
            description: 'a number',
            type: 'NUMBER',
          },
          simpleString: {
            description: 'a string',
            nullable: true,
            type: 'STRING',
          },
        },
        required: ['simpleString', 'simpleNumber'],
        type: 'OBJECT',
      },
    };
    assert.deepStrictEqual(got, want);
  });

  it('converts empty object schema', async () => {
    const got = toGeminiTool({
      name: 'foo',
      description: 'tool foo',
      inputSchema: {
        type: 'object',
      },
    });

    assert.deepStrictEqual(got, {
      description: 'tool foo',
      name: 'foo',
      parameters: {
        properties: {},
        required: undefined,
        type: 'OBJECT',
      },
    });
  });
});
