import { GoogleGenerativeAI } from '@google/generative-ai';
import { Tone } from '../../types';

jest.mock('@google/generative-ai');
jest.mock('../../utils/logger');

const mockGetGeminiConfig = jest.fn();
const mockOnConfigChange = jest.fn();

jest.mock('../../config', () => ({
  getGeminiConfig: mockGetGeminiConfig,
  onConfigChange: mockOnConfigChange,
}));

mockGetGeminiConfig.mockReturnValue({ apiKey: 'truncation-test-key' });

import { geminiClient } from '../../services/geminiClient';

const strings = [
  { id: 'a', content: 'Alpha' },
  { id: 'b', content: 'Bravo' },
  { id: 'c', content: 'Charlie' },
  { id: 'd', content: 'Delta' },
];

type ResponseOptions = {
  finishReason?: string;
  usageMetadata?: Record<string, number>;
};

const responseFor = (body: string, options: ResponseOptions = {}) => ({
  response: {
    text: () => body,
    usageMetadata: options.usageMetadata ?? { promptTokenCount: 2, candidatesTokenCount: 3 },
    ...(options.finishReason ? { candidates: [{ finishReason: options.finishReason }] } : {}),
  },
});

const translated = (id: string, text: string) => ({ id, text });

describe('Gemini bulk truncation retries', () => {
  let mockGenerateContent: jest.Mock;
  let mockGetGenerativeModel: jest.Mock;
  let testCounter = 0;

  beforeEach(() => {
    jest.clearAllMocks();
    testCounter += 1;
    mockGenerateContent = jest.fn();
    mockGetGenerativeModel = jest.fn().mockReturnValue({ generateContent: mockGenerateContent });
    (GoogleGenerativeAI as jest.Mock).mockImplementation(() => ({
      getGenerativeModel: mockGetGenerativeModel,
    }));
    mockGetGeminiConfig.mockReturnValue({ apiKey: `truncation-test-key-${testCounter}` });
    geminiClient.refreshConfig();
  });

  it('passes a complete response through without retrying', async () => {
    mockGenerateContent.mockResolvedValue(responseFor(JSON.stringify([
      translated('a', 'Alfa'),
      translated('b', 'Bravo translated'),
    ])));

    const result = await geminiClient.translateBulk(strings.slice(0, 2), 'en', 'es', Tone.NEUTRAL);

    expect(mockGenerateContent).toHaveBeenCalledTimes(1);
    expect(result.results).toEqual([
      { id: 'a', translation: 'Alfa', success: true },
      { id: 'b', translation: 'Bravo translated', success: true },
    ]);
  });

  it('retries malformed JSON as incomplete when Gemini omits a finish reason', async () => {
    mockGenerateContent
      .mockResolvedValueOnce(responseFor('[{"id":"a","text":"Alfa"'))
      .mockResolvedValueOnce(responseFor(JSON.stringify([translated('a', 'Alfa')])));

    const result = await geminiClient.translateBulk(strings.slice(0, 1), 'en', 'es');

    expect(mockGenerateContent).toHaveBeenCalledTimes(2);
    expect(result.results).toEqual([{ id: 'a', translation: 'Alfa', success: true }]);
  });

  it('retries a missing id once and merges results in input order', async () => {
    mockGenerateContent
      .mockResolvedValueOnce(responseFor(JSON.stringify([
        translated('a', 'Alfa'),
        translated('c', 'Charlie translated'),
      ])))
      .mockResolvedValueOnce(responseFor(JSON.stringify([translated('b', 'Bravo traducido')])));

    const result = await geminiClient.translateBulk(strings.slice(0, 3), 'en', 'es');

    expect(mockGenerateContent).toHaveBeenCalledTimes(2);
    expect(mockGenerateContent.mock.calls[1][0]).toContain('"id":"b"');
    expect(result.results).toEqual([
      { id: 'a', translation: 'Alfa', success: true },
      { id: 'b', translation: 'Bravo traducido', success: true },
      { id: 'c', translation: 'Charlie translated', success: true },
    ]);
  });

  it('treats an empty translation as incomplete and retries it', async () => {
    mockGenerateContent
      .mockResolvedValueOnce(responseFor(JSON.stringify([translated('a', '')])))
      .mockResolvedValueOnce(responseFor(JSON.stringify([translated('a', 'Alfa')])));

    const result = await geminiClient.translateBulk(strings.slice(0, 1), 'en', 'es');

    expect(mockGenerateContent).toHaveBeenCalledTimes(2);
    expect(result.results).toEqual([{ id: 'a', translation: 'Alfa', success: true }]);
  });

  it('keeps earlier successes and returns explicit failures after persistent truncation', async () => {
    mockGenerateContent.mockImplementation(() => Promise.resolve(responseFor(
      JSON.stringify([translated('a', 'Alfa')]),
      { finishReason: 'MAX_TOKENS' }
    )));

    const result = await geminiClient.translateBulk(strings.slice(0, 2), 'en', 'es');

    expect(mockGenerateContent).toHaveBeenCalledTimes(4);
    expect(result.results).toEqual([
      { id: 'a', translation: 'Alfa', success: true },
      { id: 'b', translation: '', success: false, error: 'gemini output truncated' },
    ]);
  });

  it('splits the remaining strings after a no-progress retry round', async () => {
    mockGenerateContent
      .mockResolvedValueOnce(responseFor('[]'))
      .mockResolvedValueOnce(responseFor('[]'))
      .mockResolvedValueOnce(responseFor(JSON.stringify([
        translated('a', 'Alfa'),
        translated('b', 'Bravo'),
      ])))
      .mockResolvedValueOnce(responseFor(JSON.stringify([
        translated('c', 'Charlie'),
        translated('d', 'Delta'),
      ])));

    const result = await geminiClient.translateBulk(strings, 'en', 'es');

    expect(mockGenerateContent).toHaveBeenCalledTimes(4);
    expect(mockGenerateContent.mock.calls[2][0]).toContain('"id":"a"');
    expect(mockGenerateContent.mock.calls[2][0]).toContain('"id":"b"');
    expect(mockGenerateContent.mock.calls[2][0]).not.toContain('"id":"c"');
    expect(mockGenerateContent.mock.calls[3][0]).toContain('"id":"c"');
    expect(mockGenerateContent.mock.calls[3][0]).toContain('"id":"d"');
    expect(result.results.every(item => item.success)).toBe(true);
  });
});