/**
 * 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 { beforeEach, describe, expect, it } from '@jest/globals';
import { genkit } from 'genkit';
import { modelArmor } from '../src/model-armor.js';

// Mock ModelArmorClient
class MockModelArmorClient {
  sanitizeUserPrompt = async () => [{}];
  sanitizeModelResponse = async () => [{}];
}

function createEmptyResult() {
  return { sanitizationResult: {} };
}

function createSdpResult(replacementText: string) {
  return {
    sanitizationResult: {
      filterResults: {
        sdp: {
          sdpFilterResult: {
            deidentifyResult: {
              data: { text: replacementText },
            },
          },
        },
      },
    },
  };
}

function createRaiBlockResult() {
  return {
    sanitizationResult: {
      filterMatchState: 'MATCH_FOUND',
      filterResults: {
        rai: { raiFilterResult: { matchState: 'MATCH_FOUND' } },
      },
    },
  };
}

function createSdpBlockResult(replacementText: string) {
  const res = createSdpResult(replacementText);
  (res.sanitizationResult as any).filterMatchState = 'MATCH_FOUND';
  return res;
}

describe('modelArmor', () => {
  let ai: any;
  let mockClient: any;

  beforeEach(() => {
    ai = genkit({});
    ai.defineModel({ name: 'echoModel' }, async (req: any) => {
      return {
        message: {
          role: 'model',
          content: [{ text: `Echo: ${req.messages[0].content[0].text}` }],
        },
      };
    });
    mockClient = new MockModelArmorClient();
  });

  it('passes through when no sanitization triggers', async () => {
    mockClient.sanitizeUserPrompt = async () => [createEmptyResult()];
    mockClient.sanitizeModelResponse = async () => [createEmptyResult()];

    const response = await ai.generate({
      model: 'echoModel',
      prompt: 'hello',
      use: [modelArmor({ templateName: 'test', client: mockClient as any })],
    });

    expect(response.text).toMatch(/Echo: hello/);
  });

  it('replaces user prompt on SDP match', async () => {
    mockClient.sanitizeUserPrompt = async () => [
      createSdpResult('sanitized_hello'),
    ];
    mockClient.sanitizeModelResponse = async () => [createEmptyResult()];

    const response = await ai.generate({
      model: 'echoModel',
      prompt: 'hello',
      use: [
        modelArmor({
          templateName: 'test',
          client: mockClient as any,
          applyDeidentificationResults: true,
        }),
      ],
    });

    // The echo model should receive the SANITIZED prompt
    expect(response.text).toMatch(/Echo: sanitized_hello/);
  });

  it('blocks user prompt on RAI match', async () => {
    mockClient.sanitizeUserPrompt = async () => [createRaiBlockResult()];

    await expect(
      ai.generate({
        model: 'echoModel',
        prompt: 'bad stuff',
        use: [modelArmor({ templateName: 'test', client: mockClient as any })],
      })
    ).rejects.toThrow(
      expect.objectContaining({
        status: 'PERMISSION_DENIED',
        message: expect.stringContaining('Model Armor blocked user prompt.'),
      })
    );
  });

  it('replaces model response on SDP match', async () => {
    mockClient.sanitizeUserPrompt = async () => [createEmptyResult()];
    mockClient.sanitizeModelResponse = async () => [
      createSdpResult('sanitized_response'),
    ];

    const response = await ai.generate({
      model: 'echoModel',
      prompt: 'hello',
      use: [
        modelArmor({
          templateName: 'test',
          client: mockClient as any,
          applyDeidentificationResults: true,
        }),
      ],
    });

    expect(response.text).toBe('sanitized_response');
  });

  it('blocks model response on RAI match', async () => {
    mockClient.sanitizeUserPrompt = async () => [createEmptyResult()];
    mockClient.sanitizeModelResponse = async () => [createRaiBlockResult()];

    await expect(
      ai.generate({
        model: 'echoModel',
        prompt: 'hello',
        use: [modelArmor({ templateName: 'test', client: mockClient as any })],
      })
    ).rejects.toThrow(
      expect.objectContaining({
        status: 'PERMISSION_DENIED',
        message: expect.stringContaining('Model Armor blocked model response.'),
      })
    );
  });

  it('respects protectionTarget=userPrompt', async () => {
    // Should sanitize prompt but NOT response
    mockClient.sanitizeUserPrompt = async () => [createSdpResult('sanitized')];
    // This one would block if called
    mockClient.sanitizeModelResponse = async () => [createRaiBlockResult()];

    const response = await ai.generate({
      model: 'echoModel',
      prompt: 'hello',
      use: [
        modelArmor({
          templateName: 'test',
          client: mockClient as any,
          protectionTarget: 'userPrompt',
          applyDeidentificationResults: true,
        }),
      ],
    });

    expect(response.text).toMatch(/Echo: sanitized/);
  });

  it('strictSdpEnforcement blocks even if remediated', async () => {
    mockClient.sanitizeUserPrompt = async () => [
      createSdpBlockResult('sanitized_hello'),
    ];

    await expect(
      ai.generate({
        model: 'echoModel',
        prompt: 'sensitive',
        use: [
          modelArmor({
            templateName: 'test',
            client: mockClient as any,
            strictSdpEnforcement: true,
            applyDeidentificationResults: true,
          }),
        ],
      })
    ).rejects.toThrow(
      expect.objectContaining({
        status: 'PERMISSION_DENIED',
        message: expect.stringContaining('Model Armor blocked user prompt.'),
      })
    );
  });

  it('respects filters option', async () => {
    // RAI match found, but we only filter 'csam'
    mockClient.sanitizeUserPrompt = async () => [
      {
        sanitizationResult: {
          filterMatchState: 'MATCH_FOUND',
          filterResults: {
            rai: { raiFilterResult: { matchState: 'MATCH_FOUND' } },
          },
        },
      },
    ];

    const response = await ai.generate({
      model: 'echoModel',
      prompt: 'bad stuff',
      use: [
        modelArmor({
          templateName: 'test',
          client: mockClient as any,
          filters: ['csam'],
        }),
      ],
    });

    expect(response.text).toMatch(/Echo: bad stuff/);
  });

  it('preserves non-text parts when SDP replaces text', async () => {
    mockClient.sanitizeUserPrompt = async () => [
      createSdpResult('sanitized_text'),
    ];

    ai.defineModel({ name: 'inspectionModel' }, async (req: any) => {
      return {
        message: {
          role: 'model',
          content: [{ text: JSON.stringify(req.messages) }],
        },
      };
    });

    const response = await ai.generate({
      model: 'inspectionModel',
      messages: [
        {
          role: 'user',
          content: [{ text: 'old stuff' }],
        },
        {
          role: 'model',
          content: [{ text: 'response' }],
        },
        {
          role: 'user',
          content: [
            { text: 'sensitive info' },
            { media: { url: 'http://example.com/image.png' } },
          ],
        },
      ],
      use: [
        modelArmor({
          templateName: 'test',
          client: mockClient as any,
          applyDeidentificationResults: true,
        }),
      ],
    });

    const content = JSON.parse(response.text);
    // content should have preserved media and replaced text
    expect(content).toEqual([
      {
        role: 'user',
        content: [{ text: 'old stuff' }],
      },
      {
        role: 'model',
        content: [{ text: 'response' }],
      },
      {
        role: 'user',
        content: [
          { media: { url: 'http://example.com/image.png' } },
          { text: 'sanitized_text' },
        ],
      },
    ]);
  });

  it('supports custom function for applying SDP', async () => {
    mockClient.sanitizeUserPrompt = async () => [
      createSdpResult('sanitized_text'),
    ];

    const applyFn = ({ messages, sdpResult }: any) => {
      // Custom logic: replace with "CUSTOM APPLIED" instead of sdpResult data
      const newContent = [{ text: 'CUSTOM APPLIED' }];
      return [{ ...messages[0], content: newContent }];
    };

    const response = await ai.generate({
      model: 'echoModel',
      prompt: 'hello',
      use: [
        modelArmor({
          templateName: 'test',
          client: mockClient as any,
          applyDeidentificationResults: applyFn,
        }),
      ],
    });

    expect(response.text).toMatch(/Echo: CUSTOM APPLIED/);
  });
});