/**
 * 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 type { MessageCreateParamsBase } from '@anthropic-ai/sdk/resources/messages.mjs';
import * as assert from 'assert';
import type { GenerateRequest, GenerateResponseData } from 'genkit';
import { describe, it } from 'node:test';
import {
  fromAnthropicResponse,
  toAnthropicRequest,
  type AnthropicConfigSchema,
} from '../../../src/modelgarden/v2/anthropic';

const MODEL_ID = 'modelid';

describe('toAnthropicRequest', () => {
  const testCases: {
    should: string;
    input: GenerateRequest<typeof AnthropicConfigSchema>;
    expectedOutput: MessageCreateParamsBase;
  }[] = [
    {
      should: 'should transform genkit message (text content) correctly',
      input: {
        messages: [
          {
            role: 'user',
            content: [{ text: 'Tell a joke about dogs.' }],
          },
        ],
      },
      expectedOutput: {
        max_tokens: 4096,
        model: MODEL_ID,
        messages: [
          {
            role: 'user',
            content: [{ type: 'text', text: 'Tell a joke about dogs.' }],
          },
        ],
      },
    },
    {
      should:
        'should transform genkit reasoning part back to anthropic thinking block',
      input: {
        messages: [
          {
            role: 'model',
            content: [
              {
                reasoning: 'I need to think about this.',
                metadata: { thoughtSignature: 'sig_123' },
              },
              { text: 'I thought about it.' },
            ],
          },
        ],
      },
      expectedOutput: {
        max_tokens: 4096,
        model: MODEL_ID,
        messages: [
          {
            role: 'assistant',
            content: [
              {
                type: 'thinking',
                thinking: 'I need to think about this.',
                signature: 'sig_123',
              } as any,
              { type: 'text', text: 'I thought about it.' },
            ],
          },
        ],
      },
    },
    {
      should: 'should transform system message',
      input: {
        messages: [
          {
            role: 'system',
            content: [{ text: 'Talk like a pirate.' }],
          },
          {
            role: 'user',
            content: [{ text: 'Tell a joke about dogs.' }],
          },
        ],
      },
      expectedOutput: {
        max_tokens: 4096,
        model: MODEL_ID,
        system: 'Talk like a pirate.',
        messages: [
          {
            role: 'user',
            content: [{ type: 'text', text: 'Tell a joke about dogs.' }],
          },
        ],
      },
    },
    {
      should: 'should transform thinking and output_config parameters',
      input: {
        messages: [
          {
            role: 'user',
            content: [{ text: 'Refactor this.' }],
          },
        ],
        config: {
          thinking: { adaptive: true, display: 'summarized' },
          output_config: { effort: 'xhigh' },
        },
      } as any,
      expectedOutput: {
        max_tokens: 4096,
        model: MODEL_ID,
        messages: [
          {
            role: 'user',
            content: [{ type: 'text', text: 'Refactor this.' }],
          },
        ],
        thinking: { type: 'adaptive', display: 'summarized' },
        output_config: { effort: 'xhigh' },
      } as any,
    },
    {
      should:
        'should transform genkit message (inline base64 image content) correctly',
      input: {
        messages: [
          {
            role: 'user',
            content: [
              { text: 'describe the following image:' },
              {
                media: {
                  contentType: 'image/jpeg',
                  url: 'data:image/jpeg;base64, ',
                },
              },
            ],
          },
        ],
      },
      expectedOutput: {
        max_tokens: 4096,
        model: MODEL_ID,
        messages: [
          {
            role: 'user',
            content: [
              { type: 'text', text: 'describe the following image:' },
              {
                type: 'image',
                source: {
                  type: 'base64',
                  media_type: 'image/jpeg',
                  data: ' ',
                },
              },
            ],
          },
        ],
      },
    },
  ];
  for (const test of testCases) {
    it(test.should, () => {
      assert.deepEqual(
        toAnthropicRequest(MODEL_ID, test.input),
        test.expectedOutput
      );
    });
  }
});

describe('fromAnthropicResponse', () => {
  const testCases: {
    should: string;
    input: GenerateRequest<typeof AnthropicConfigSchema>;
    response: any;
    expectedOutput: GenerateResponseData;
  }[] = [
    {
      should: 'should transform genkit message (text content) correctly',
      input: {
        messages: [
          {
            role: 'user',
            content: [{ text: 'Tell a joke about dogs.' }],
          },
        ],
      },
      response: {
        id: 'abcd1234',
        model: MODEL_ID,
        role: 'assistant',
        stop_reason: 'end_turn',
        usage: {
          input_tokens: 123,
          output_tokens: 234,
        },
        stop_sequence: null,
        type: 'message',
        content: [
          {
            type: 'text',
            text: 'part 1',
          },
          {
            type: 'text',
            text: 'part 2',
          },
        ],
      },
      expectedOutput: {
        custom: {
          id: 'abcd1234',
          model: MODEL_ID,
          type: 'message',
        },
        finishReason: 'stop',
        message: {
          role: 'model',
          content: [
            {
              text: 'part 1',
            },
            {
              text: 'part 2',
            },
          ],
        },
        usage: {
          inputAudioFiles: 0,
          inputCharacters: 23,
          inputImages: 0,
          inputTokens: 123,
          inputVideos: 0,
          outputAudioFiles: 0,
          outputCharacters: 12,
          outputImages: 0,
          outputTokens: 234,
          outputVideos: 0,
          custom: {
            cache_creation_input_tokens: 0,
            cache_read_input_tokens: 0,
            ephemeral_5m_input_tokens: 0,
            ephemeral_1h_input_tokens: 0,
          },
        },
      },
    },
    {
      should: 'should transform genkit message with thinking block correctly',
      input: {
        messages: [
          {
            role: 'user',
            content: [{ text: 'Solve this.' }],
          },
        ],
      },
      response: {
        id: 'abcd1234',
        model: MODEL_ID,
        role: 'assistant',
        stop_reason: 'end_turn',
        usage: {
          input_tokens: 10,
          output_tokens: 50,
        },
        stop_sequence: null,
        type: 'message',
        content: [
          {
            type: 'thinking',
            thinking: 'Step 1: Ponder deeply.',
            signature: 'sig_123',
          },
          {
            type: 'text',
            text: 'The answer is 42.',
          },
        ],
      },
      expectedOutput: {
        custom: {
          id: 'abcd1234',
          model: MODEL_ID,
          type: 'message',
        },
        finishReason: 'stop',
        message: {
          role: 'model',
          content: [
            {
              reasoning: 'Step 1: Ponder deeply.',
              metadata: { thoughtSignature: 'sig_123' },
            },
            {
              text: 'The answer is 42.',
            },
          ],
        },
        usage: {
          inputAudioFiles: 0,
          inputCharacters: 11,
          inputImages: 0,
          inputTokens: 10,
          inputVideos: 0,
          outputAudioFiles: 0,
          outputCharacters: 17,
          outputImages: 0,
          outputTokens: 50,
          outputVideos: 0,
          custom: {
            cache_creation_input_tokens: 0,
            cache_read_input_tokens: 0,
            ephemeral_5m_input_tokens: 0,
            ephemeral_1h_input_tokens: 0,
          },
        },
      },
    },
    {
      should:
        'should transform response with thinking, text, and multiple tool calls correctly',
      input: {
        messages: [
          {
            role: 'user',
            content: [{ text: 'Do some tasks.' }],
          },
        ],
        tools: [
          {
            name: 'someListFilesTool',
            description: 'Lists files',
            inputSchema: { type: 'object', properties: {} },
          },
        ],
      },
      response: {
        model: 'claude-opus-4-7',
        id: 'msg_vrtx_123',
        type: 'message',
        role: 'assistant',
        content: [
          {
            type: 'thinking',
            thinking: 'I will do multiple tool calls.',
            signature: 'sig_abc',
          },
          {
            type: 'text',
            text: 'Here we go.',
          },
          {
            type: 'tool_use',
            id: 'toolu_vrtx_1',
            name: 'someListFilesTool',
            input: {},
          },
          {
            type: 'tool_use',
            id: 'toolu_vrtx_2',
            name: 'someListFilesTool',
            input: {},
          },
        ],
        stop_reason: 'tool_use',
        stop_sequence: null,
        usage: {
          input_tokens: 800,
          output_tokens: 444,
          cache_creation_input_tokens: 0,
          cache_read_input_tokens: 0,
          cache_creation: {
            ephemeral_5m_input_tokens: 0,
            ephemeral_1h_input_tokens: 0,
          },
        },
      },
      expectedOutput: {
        custom: {
          id: 'msg_vrtx_123',
          model: 'claude-opus-4-7',
          type: 'message',
        },
        finishReason: 'stop',
        message: {
          role: 'model',
          content: [
            {
              reasoning: 'I will do multiple tool calls.',
              metadata: { thoughtSignature: 'sig_abc' },
            },
            {
              text: 'Here we go.',
            },
            {
              toolRequest: {
                name: 'someListFilesTool',
                ref: 'toolu_vrtx_1',
                input: {},
              },
            },
            {
              toolRequest: {
                name: 'someListFilesTool',
                ref: 'toolu_vrtx_2',
                input: {},
              },
            },
          ],
        },
        usage: {
          inputAudioFiles: 0,
          inputCharacters: 14,
          inputImages: 0,
          inputTokens: 800,
          inputVideos: 0,
          outputAudioFiles: 0,
          outputCharacters: 11,
          outputImages: 0,
          outputTokens: 444,
          outputVideos: 0,
          custom: {
            cache_creation_input_tokens: 0,
            cache_read_input_tokens: 0,
            ephemeral_5m_input_tokens: 0,
            ephemeral_1h_input_tokens: 0,
          },
        },
      },
    },
    {
      should: 'should transform genkit tool call correctly',
      input: {
        messages: [
          {
            role: 'user',
            content: [{ text: "What's the weather like today?" }],
          },
        ],
        tools: [
          {
            name: 'get_weather',
            description: 'Get the weather for a location.',
            inputSchema: {
              type: 'object',
              properties: {
                location: {
                  type: 'string',
                  description: 'The city and state, e.g. San Francisco, CA',
                },
              },
              required: ['location'],
            },
          },
        ],
      },
      response: {
        id: 'abcd1234',
        model: MODEL_ID,
        role: 'assistant',
        type: 'message',
        stop_reason: 'tool_use',
        stop_sequence: null,
        usage: {
          input_tokens: 123,
          output_tokens: 234,
        },
        content: [
          {
            id: 'toolu_get_weather',
            name: 'get_weather',
            type: 'tool_use',
            input: {
              type: 'object',
              properties: {
                location: {
                  type: 'string',
                  description: 'The city and state, e.g. San Francisco, CA',
                },
              },
              required: ['location'],
            },
          },
        ],
      },
      expectedOutput: {
        custom: {
          id: 'abcd1234',
          model: MODEL_ID,
          type: 'message',
        },
        finishReason: 'stop',
        message: {
          role: 'model',
          content: [
            {
              toolRequest: {
                name: 'get_weather',
                ref: 'toolu_get_weather',
                input: {
                  type: 'object',
                  properties: {
                    location: {
                      type: 'string',
                      description: 'The city and state, e.g. San Francisco, CA',
                    },
                  },
                  required: ['location'],
                },
              },
            },
          ],
        },
        usage: {
          inputAudioFiles: 0,
          inputCharacters: 30,
          inputImages: 0,
          inputTokens: 123,
          inputVideos: 0,
          outputAudioFiles: 0,
          outputCharacters: 0,
          outputImages: 0,
          outputTokens: 234,
          outputVideos: 0,
          custom: {
            cache_creation_input_tokens: 0,
            cache_read_input_tokens: 0,
            ephemeral_5m_input_tokens: 0,
            ephemeral_1h_input_tokens: 0,
          },
        },
      },
    },
  ];
  for (const test of testCases) {
    it(test.should, () => {
      assert.deepEqual(fromAnthropicResponse(test.input, test.response), {
        ...test.expectedOutput,
        raw: test.response,
      });
    });
  }
});
