import { GroqClient } from '../../src';
import { AvailableModel, ChatMessageTextContent, ChatMessageImageContent } from '@browserbasehq/stagehand';

const mockLogger = jest.fn();

describe('GroqClient System Prompt Handling', () => {
  let client: GroqClient;
  let mockCreate: jest.Mock;

  beforeEach(() => {
    mockLogger.mockClear();
    client = new GroqClient({
      modelName: 'mixtral-8x7b-32768' as AvailableModel,
      clientOptions: { apiKey: 'test-key' },
      userProvidedInstructions: undefined,
    });
    mockCreate = jest.fn().mockResolvedValue({
      id: 'test-id',
      object: 'chat.completion',
      created: Date.now(),
      model: 'mixtral-8x7b-32768',
      choices: [
        {
          message: { role: 'assistant', content: 'Test response' },
          finish_reason: 'stop',
        },
      ],
      usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
    });
    jest.spyOn(client["client"].chat.completions, 'create').mockImplementation(mockCreate);
  });

  describe('Basic System Prompt Functionality', () => {
    it('should not add system message when no user instructions or system message exist', async () => {
      const messages = [{ role: 'user' as const, content: 'Hello' }];

      await client.createChatCompletion({
        options: {
          messages,
          requestId: 'test-request',
        },
        logger: mockLogger,
      });

      expect(mockCreate).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [{ role: 'user', content: 'Hello' }],
        }),
      );

      // Verify no system message related logs were created
      expect(mockLogger).not.toHaveBeenCalledWith(
        expect.objectContaining({
          message: expect.stringMatching(/system message/i),
        }),
      );
    });

    it('should preserve existing system message when no user instructions provided', async () => {
      const systemMessage = 'Original system message';
      const messages = [
        { role: 'system' as const, content: systemMessage },
        { role: 'user' as const, content: 'Hello' },
      ];

      await client.createChatCompletion({
        options: {
          messages,
          requestId: 'test-request',
        },
        logger: mockLogger,
      });

      expect(mockCreate).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [
            { role: 'system', content: systemMessage },
            { role: 'user', content: 'Hello' },
          ],
        }),
      );

      // Verify no combining logs were created
      expect(mockLogger).not.toHaveBeenCalledWith(
        expect.objectContaining({
          message: 'Combining system message with user instructions',
        }),
      );
    });

    it('should handle empty system message content', async () => {
      const messages = [
        { role: 'system' as const, content: '' },
        { role: 'user' as const, content: 'Hello' },
      ];

      await client.createChatCompletion({
        options: {
          messages,
          requestId: 'test-request',
        },
        logger: mockLogger,
      });

      expect(mockCreate).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [
            { role: 'system', content: '' },
            { role: 'user', content: 'Hello' },
          ],
        }),
      );
    });

    it('should handle non-string system message content', async () => {
      const contentObject = { key: 'value', nested: { data: true } };
      const messages = [
        { role: 'system' as const, content: contentObject as any },
        { role: 'user' as const, content: 'Hello' },
      ];

      await client.createChatCompletion({
        options: {
          messages,
          requestId: 'test-request',
        },
        logger: mockLogger,
      });

      expect(mockCreate).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [
            { role: 'system', content: JSON.stringify(contentObject) },
            { role: 'user', content: 'Hello' },
          ],
        }),
      );
    });

    it('should return original messages when no system message and no user instructions', async () => {
      // Create a client without user instructions
      const clientWithoutInstructions = new GroqClient({
        modelName: 'mixtral-8x7b-32768' as AvailableModel,
        clientOptions: { apiKey: 'test-key' },
        userProvidedInstructions: undefined,
      });
      const mockCreateWithoutInstructions = jest.fn().mockResolvedValue({
        id: 'test-id',
        object: 'chat.completion',
        created: Date.now(),
        model: 'mixtral-8x7b-32768',
        choices: [
          {
            message: { role: 'assistant', content: 'Test response' },
            finish_reason: 'stop',
          },
        ],
        usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
      });
      jest.spyOn(clientWithoutInstructions["client"].chat.completions, 'create').mockImplementation(mockCreateWithoutInstructions);

      const messages = [{ role: 'user' as const, content: 'Hello' }];
      await clientWithoutInstructions.createChatCompletion({
        options: {
          messages,
          requestId: 'test-request',
        },
        logger: mockLogger,
      });

      expect(mockCreateWithoutInstructions).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [{ role: 'user', content: 'Hello' }],
        }),
      );

      // Verify no system message related logs were created
      expect(mockLogger).not.toHaveBeenCalledWith(
        expect.objectContaining({
          message: expect.stringMatching(/system message/i),
        }),
      );
    });

    it('should return original system prompt when userProvidedInstructions is empty', async () => {
      const systemPrompt = 'Original system prompt';
      const messages = [
        { role: 'system' as const, content: systemPrompt },
        { role: 'user' as const, content: 'Hello' },
      ];

      // Create a client with empty user instructions
      const clientWithEmptyInstructions = new GroqClient({
        modelName: 'mixtral-8x7b-32768' as AvailableModel,
        clientOptions: { apiKey: 'test-key' },
        userProvidedInstructions: '',
      });
      const mockCreateEmptyInstructions = jest.fn().mockResolvedValue({
        id: 'test-id',
        object: 'chat.completion',
        created: Date.now(),
        model: 'mixtral-8x7b-32768',
        choices: [
          {
            message: { role: 'assistant', content: 'Test response' },
            finish_reason: 'stop',
          },
        ],
        usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
      });
      jest.spyOn(clientWithEmptyInstructions["client"].chat.completions, 'create').mockImplementation(mockCreateEmptyInstructions);

      await clientWithEmptyInstructions.createChatCompletion({
        options: {
          messages,
          requestId: 'test-request',
        },
        logger: mockLogger,
      });

      expect(mockCreateEmptyInstructions).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [
            { role: 'system', content: systemPrompt },
            { role: 'user', content: 'Hello' },
          ],
        }),
      );
    });

    it('should return original system prompt when userProvidedInstructions contains only whitespace', async () => {
      const systemPrompt = 'Original system prompt';
      const messages = [
        { role: 'system' as const, content: systemPrompt },
        { role: 'user' as const, content: 'Hello' },
      ];

      // Create a client with whitespace-only user instructions
      const clientWithWhitespaceInstructions = new GroqClient({
        modelName: 'mixtral-8x7b-32768' as AvailableModel,
        clientOptions: { apiKey: 'test-key' },
        userProvidedInstructions: '   \n\t   ',  // Only whitespace
      });
      const mockCreateWhitespaceInstructions = jest.fn().mockResolvedValue({
        id: 'test-id',
        object: 'chat.completion',
        created: Date.now(),
        model: 'mixtral-8x7b-32768',
        choices: [
          {
            message: { role: 'assistant', content: 'Test response' },
            finish_reason: 'stop',
          },
        ],
        usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
      });
      jest.spyOn(clientWithWhitespaceInstructions["client"].chat.completions, 'create').mockImplementation(mockCreateWhitespaceInstructions);

      await clientWithWhitespaceInstructions.createChatCompletion({
        options: {
          messages,
          requestId: 'test-request',
        },
        logger: mockLogger,
      });

      expect(mockCreateWhitespaceInstructions).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [
            { role: 'system', content: systemPrompt },
            { role: 'user', content: 'Hello' },
          ],
        }),
      );
    });
  });

  describe('User Instructions Integration', () => {
    it('should add system message when only user instructions are provided', async () => {
      const userInstructions = 'Always respond in JSON format';
      client = new GroqClient({
        modelName: 'mixtral-8x7b-32768' as AvailableModel,
        clientOptions: { apiKey: 'test-key' },
        userProvidedInstructions: userInstructions,
      });
      jest.spyOn(client["client"].chat.completions, 'create').mockImplementation(mockCreate);

      const messages = [{ role: 'user' as const, content: 'Hello' }];

      await client.createChatCompletion({
        options: {
          messages,
          requestId: 'test-request',
        },
        logger: mockLogger,
      });

      expect(mockCreate).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [
            {
              role: 'system',
              content: expect.stringContaining(userInstructions),
            },
            { role: 'user', content: 'Hello' },
          ],
        }),
      );

      expect(mockLogger).toHaveBeenCalledWith(
        expect.objectContaining({
          category: 'groq',
          message: 'Adding system message with user instructions',
          level: 1,
          auxiliary: expect.objectContaining({
            userInstructions: {
              value: userInstructions,
              type: 'string',
            },
          }),
        }),
      );
    });

    it('should combine existing system message with user instructions', async () => {
      const userInstructions = 'Always respond in JSON format';
      const systemMessage = 'Original system message';
      client = new GroqClient({
        modelName: 'mixtral-8x7b-32768' as AvailableModel,
        clientOptions: { apiKey: 'test-key' },
        userProvidedInstructions: userInstructions,
      });
      jest.spyOn(client["client"].chat.completions, 'create').mockImplementation(mockCreate);

      const messages = [
        { role: 'system' as const, content: systemMessage },
        { role: 'user' as const, content: 'Hello' },
      ];

      await client.createChatCompletion({
        options: {
          messages,
          requestId: 'test-request',
        },
        logger: mockLogger,
      });

      expect(mockCreate).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [
            {
              role: 'system',
              content: expect.stringContaining(systemMessage),
            },
            { role: 'user', content: 'Hello' },
          ],
        }),
      );

      const combinedMessage = mockCreate.mock.calls[0][0].messages[0].content;
      expect(combinedMessage).toContain(systemMessage);
      expect(combinedMessage).toContain(userInstructions);
      expect(combinedMessage).toContain('Custom Instructions Provided by the User');

      expect(mockLogger).toHaveBeenCalledWith(
        expect.objectContaining({
          category: 'groq',
          message: 'Combining system message with user instructions',
          level: 1,
          auxiliary: expect.objectContaining({
            originalSystemMessage: {
              value: systemMessage,
              type: 'string',
            },
            userInstructions: {
              value: userInstructions,
              type: 'string',
            },
          }),
        }),
      );
    });

    it('should handle special characters in user instructions', async () => {
      const userInstructions = 'Line 1\nLine 2\tTabbed\n<html>tags</html>\n"quotes" and \\backslashes\\';
      client = new GroqClient({
        modelName: 'mixtral-8x7b-32768' as AvailableModel,
        clientOptions: { apiKey: 'test-key' },
        userProvidedInstructions: userInstructions,
      });
      jest.spyOn(client["client"].chat.completions, 'create').mockImplementation(mockCreate);

      const messages = [{ role: 'user' as const, content: 'Hello' }];

      await client.createChatCompletion({
        options: {
          messages,
          requestId: 'test-request',
        },
        logger: mockLogger,
      });

      const systemMessage = mockCreate.mock.calls[0][0].messages[0].content;
      expect(systemMessage).toContain(userInstructions);
      expect(systemMessage.split('\n').length).toBeGreaterThan(1);
    });
  });

  describe('Message Content Handling', () => {
    it('should handle array message content with mixed types', async () => {
      const messages = [
        {
          role: 'user' as const,
          content: [
            { type: 'text', text: 'plain text' } as ChatMessageTextContent,
            { type: 'image_url', image_url: { url: 'https://example.com/image.jpg' } } as ChatMessageImageContent,
          ],
        },
      ];

      await client.createChatCompletion({
        options: {
          messages,
          requestId: 'test-request',
        },
        logger: mockLogger,
      });

      expect(mockCreate).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [
            {
              role: 'user',
              content: 'plain text\nhttps://example.com/image.jpg',
            },
          ],
        }),
      );
    });

    it('should handle empty or missing content values', async () => {
      const messages = [
        {
          role: 'user' as const,
          content: [
            { type: 'text', text: '' } as ChatMessageTextContent,
            { type: 'image_url', image_url: { url: '' } } as ChatMessageImageContent,
            { type: 'text', text: '' } as ChatMessageTextContent,
            { type: 'image_url', image_url: { url: '' } } as ChatMessageImageContent,
          ],
        },
      ];

      await client.createChatCompletion({
        options: {
          messages,
          requestId: 'test-request',
        },
        logger: mockLogger,
      });

      expect(mockCreate).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [
            {
              role: 'user',
              content: '\n\n\n',
            },
          ],
        }),
      );
    });
  });
});