import { LLMService } from '../src/services/LLMService';
import { GoogleGenerativeAI } from '@google/generative-ai';

// Mock the GoogleGenerativeAI module
jest.mock('@google/generative-ai');

describe('LLMService', () => {
    let llmService: LLMService;
    let mockGenerateContent: jest.Mock;

    beforeEach(() => {
        // Reset mocks before each test
        jest.clearAllMocks();

        // Mock the generateContent method of the GenerativeModel
        mockGenerateContent = jest.fn();
        (GoogleGenerativeAI as jest.Mock).mockImplementation(() => ({
            getGenerativeModel: jest.fn(() => ({
                generateContent: mockGenerateContent,
            })),
        }));

        llmService = new LLMService('test-api-key');
    });

    describe('analyzeDesires', () => {
        it('should call the LLM with the correct prompt and handle afraidToAsk ideas', async () => {
            const desireSets = [
                { wants: ['apple'], accepts: [], noGoes: [], afraidToAsk: 'banana' },
                { wants: ['orange'], accepts: ['banana'], noGoes: [], afraidToAsk: 'grape' },
            ];

            mockGenerateContent.mockResolvedValue({
                response: {
                    text: () => JSON.stringify({
                        goTo: 'apple',
                        alsoGood: 'banana',
                        considerable: 'orange',
                        noGoes: '',
                        needsDiscussion: '',
                    }),
                },
            });

            const result = await llmService.analyzeDesires(desireSets);

            expect(mockGenerateContent).toHaveBeenCalledTimes(1);
            const prompt = mockGenerateContent.mock.calls[0][0];
            expect(prompt).toContain(JSON.stringify(desireSets));

            expect(prompt).toContain('If an \'afraidToAsk\' idea matches, it should be treated as a \'want\'');
            expect(result).toEqual({
                goTo: 'apple',
                alsoGood: 'banana',
                considerable: 'orange',
                noGoes: '',
                needsDiscussion: '',
            });
        });

        it('should throw an error if LLM response is not valid JSON', async () => {
            const desireSets = [
                { wants: ['apple'], accepts: [], noGoes: [], afraidToAsk: 'banana' },
            ];

            mockGenerateContent.mockResolvedValue({
                response: {
                    text: () => 'invalid json',
                },
            });

            await expect(llmService.analyzeDesires(desireSets)).rejects.toThrow('Failed to parse LLM response as JSON.');
        });
    });

    describe('checkForInnerContradictions', () => {
        it('should return null if no contradictions are found', async () => {
            const desireSet = { wants: ['apple'], accepts: ['banana'], noGoes: ['grape'], afraidToAsk: 'kiwi' };
            mockGenerateContent.mockResolvedValue({
                response: {
                    text: () => 'null',
                },
            });

            const result = await llmService.checkForInnerContradictions(desireSet);
            expect(result).toBeNull();
        });

        it('should return a contradiction message if contradictions are found', async () => {
            const desireSet = { wants: ['apple', 'no apple'], accepts: [], noGoes: [], afraidToAsk: '' };
            mockGenerateContent.mockResolvedValue({
                response: {
                    text: () => 'Contradiction: apple and no apple in wants.',
                },
            });

            const result = await llmService.checkForInnerContradictions(desireSet);
            expect(result).toBe('Contradiction: apple and no apple in wants.');
        });
    });
});
