import { lastValueFrom, of } from 'rxjs';

import {
  type DataQueryRequest,
  type DataQueryResponse,
  type Field,
  FieldType,
  LogLevel,
  LogRowContextQueryDirection,
  type LogRowModel,
} from '@grafana/data';

import {
  type CloudWatchLogsAnomaliesQuery,
  type CloudWatchLogsQuery,
  LogsMode,
  LogsQueryLanguage,
} from '../dataquery.gen';
import { logGroupNamesVariable, regionVariable } from '../mocks/CloudWatchDataSource';
import { setupMockedLogsQueryRunner } from '../mocks/LogsQueryRunner';
import { LogsRequestMock } from '../mocks/Request';
import { validLogsQuery } from '../mocks/queries';
import { TimeRangeMock } from '../mocks/timeRange';
import { LOG_GROUP_ACCOUNT_MAX, LOG_GROUP_PREFIX_MAX } from '../utils/logGroupsConstants';

import {
  LOGSTREAM_IDENTIFIER_INTERNAL,
  LOG_IDENTIFIER_INTERNAL,
  convertTrendHistogramToSparkline,
} from './CloudWatchLogsQueryRunner';

describe('CloudWatchLogsQueryRunner', () => {
  beforeEach(() => {
    jest.clearAllMocks();
  });

  describe('interpolateLogsQueryVariables', () => {
    it('returns logGroups with arn and name values sourced from the log group template variable', () => {
      const { runner } = setupMockedLogsQueryRunner({ variables: [logGroupNamesVariable] });

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
        logGroups: [{ arn: '$groups', name: '$groups' }],
      };

      const { logGroups } = runner.interpolateLogsQueryVariables(query, {});

      expect(logGroups).toEqual([
        { arn: 'templatedGroup-arn-1', name: 'templatedGroup-1' },
        { arn: 'templatedGroup-arn-2', name: 'templatedGroup-2' },
      ]);
    });

    it('filters out duplicate log group arns when query already includes an expanded value', () => {
      const { runner } = setupMockedLogsQueryRunner({ variables: [logGroupNamesVariable] });

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
        logGroups: [
          { arn: 'templatedGroup-arn-1', name: 'existing-group-name' },
          { arn: '$groups', name: '$groups' },
        ],
      };

      const { logGroups } = runner.interpolateLogsQueryVariables(query, {});

      expect(logGroups).toEqual([
        { arn: 'templatedGroup-arn-1', name: 'existing-group-name' },
        { arn: 'templatedGroup-arn-2', name: 'templatedGroup-2' },
      ]);
    });

    it('keeps log groups with duplicate names as long as arns are unique', () => {
      const { runner } = setupMockedLogsQueryRunner({ variables: [logGroupNamesVariable] });

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
        logGroups: [
          { arn: 'arn-1', name: 'templatedGroup-1' },
          { arn: '$groups', name: '$groups' },
        ],
      };

      const { logGroups } = runner.interpolateLogsQueryVariables(query, {});

      expect(logGroups).toEqual([
        { arn: 'arn-1', name: 'templatedGroup-1' },
        { arn: 'templatedGroup-arn-1', name: 'templatedGroup-1' },
        { arn: 'templatedGroup-arn-2', name: 'templatedGroup-2' },
      ]);
    });

    it('returns expanded logGroupPrefixes from template variable', () => {
      const prefixVariable = {
        ...logGroupNamesVariable,
        id: 'prefixes',
        name: 'prefixes',
        current: {
          value: ['prefix-1', 'prefix-2'],
          text: ['prefix-1', 'prefix-2'],
          selected: true,
        },
        multi: true,
      };
      const { runner } = setupMockedLogsQueryRunner({ variables: [prefixVariable] });

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
        logGroupPrefixes: ['$prefixes'],
      };

      const { logGroupPrefixes } = runner.interpolateLogsQueryVariables(query, {});

      expect(logGroupPrefixes).toEqual(['prefix-1', 'prefix-2']);
    });

    it('deduplicates logGroupPrefixes after template variable expansion', () => {
      const prefixVariable = {
        ...logGroupNamesVariable,
        id: 'prefixes',
        name: 'prefixes',
        current: {
          value: ['prefix-1', 'prefix-2'],
          text: ['prefix-1', 'prefix-2'],
          selected: true,
        },
        multi: true,
      };
      const { runner } = setupMockedLogsQueryRunner({ variables: [prefixVariable] });

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
        logGroupPrefixes: ['prefix-1', '$prefixes'],
      };

      const { logGroupPrefixes } = runner.interpolateLogsQueryVariables(query, {});

      expect(logGroupPrefixes).toEqual(['prefix-1', 'prefix-2']);
    });

    it('throws error when expanded logGroupPrefixes exceeds 5', () => {
      const prefixVariable = {
        ...logGroupNamesVariable,
        id: 'prefixes',
        name: 'prefixes',
        current: {
          value: ['prefix-1', 'prefix-2', 'prefix-3', 'prefix-4', 'prefix-5', 'prefix-6'],
          text: ['prefix-1', 'prefix-2', 'prefix-3', 'prefix-4', 'prefix-5', 'prefix-6'],
          selected: true,
        },
        multi: true,
      };
      const { runner } = setupMockedLogsQueryRunner({ variables: [prefixVariable] });

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
        logGroupPrefixes: ['$prefixes'],
      };

      expect(() => runner.interpolateLogsQueryVariables(query, {})).toThrow(
        `Expanded prefix count (6) exceeds maximum of ${LOG_GROUP_PREFIX_MAX}`
      );
    });

    it('returns expanded selectedAccountIds from template variable', () => {
      const accountVariable = {
        ...logGroupNamesVariable,
        id: 'accounts',
        name: 'accounts',
        current: {
          value: ['123456789012', '234567890123'],
          text: ['123456789012', '234567890123'],
          selected: true,
        },
        multi: true,
      };
      const { runner } = setupMockedLogsQueryRunner({ variables: [accountVariable] });

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
        selectedAccountIds: ['$accounts'],
      };

      const { selectedAccountIds } = runner.interpolateLogsQueryVariables(query, {});

      expect(selectedAccountIds).toEqual(['123456789012', '234567890123']);
    });

    it('deduplicates selectedAccountIds after template variable expansion', () => {
      const accountVariable = {
        ...logGroupNamesVariable,
        id: 'accounts',
        name: 'accounts',
        current: {
          value: ['123456789012', '234567890123'],
          text: ['123456789012', '234567890123'],
          selected: true,
        },
        multi: true,
      };
      const { runner } = setupMockedLogsQueryRunner({ variables: [accountVariable] });

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
        selectedAccountIds: ['123456789012', '$accounts'],
      };

      const { selectedAccountIds } = runner.interpolateLogsQueryVariables(query, {});

      expect(selectedAccountIds).toEqual(['123456789012', '234567890123']);
    });

    it('returns undefined for logGroupPrefixes when not set in query', () => {
      const { runner } = setupMockedLogsQueryRunner();

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
      };

      const { logGroupPrefixes } = runner.interpolateLogsQueryVariables(query, {});

      expect(logGroupPrefixes).toBeUndefined();
    });

    it('returns undefined for selectedAccountIds when not set in query', () => {
      const { runner } = setupMockedLogsQueryRunner();

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
      };

      const { selectedAccountIds } = runner.interpolateLogsQueryVariables(query, {});

      expect(selectedAccountIds).toBeUndefined();
    });

    it('filters out "all" from selectedAccountIds as it is a special UI value', () => {
      const accountVariable = {
        ...logGroupNamesVariable,
        id: 'accounts',
        name: 'accounts',
        current: {
          value: ['all', '123456789012'],
          text: ['All', '123456789012'],
          selected: true,
        },
        multi: true,
      };
      const { runner } = setupMockedLogsQueryRunner({ variables: [accountVariable] });

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
        selectedAccountIds: ['$accounts'],
      };

      const { selectedAccountIds } = runner.interpolateLogsQueryVariables(query, {});

      expect(selectedAccountIds).toEqual(['123456789012']);
    });

    it('returns undefined when selectedAccountIds only contains "all"', () => {
      const { runner } = setupMockedLogsQueryRunner();

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
        selectedAccountIds: ['all'],
      };

      const { selectedAccountIds } = runner.interpolateLogsQueryVariables(query, {});

      expect(selectedAccountIds).toBeUndefined();
    });

    it('throws error when expanded selectedAccountIds exceeds 20', () => {
      const accountVariable = {
        ...logGroupNamesVariable,
        id: 'accounts',
        name: 'accounts',
        current: {
          value: [
            '111111111111',
            '222222222222',
            '333333333333',
            '444444444444',
            '555555555555',
            '666666666666',
            '777777777777',
            '888888888888',
            '999999999999',
            '101010101010',
            '111111111112',
            '222222222223',
            '333333333334',
            '444444444445',
            '555555555556',
            '666666666667',
            '777777777778',
            '888888888889',
            '999999999990',
            '101010101011',
            '121212121212',
          ],
          text: [
            '111111111111',
            '222222222222',
            '333333333333',
            '444444444444',
            '555555555555',
            '666666666666',
            '777777777777',
            '888888888888',
            '999999999999',
            '101010101010',
            '111111111112',
            '222222222223',
            '333333333334',
            '444444444445',
            '555555555556',
            '666666666667',
            '777777777778',
            '888888888889',
            '999999999990',
            '101010101011',
            '121212121212',
          ],
          selected: true,
        },
        multi: true,
      };
      const { runner } = setupMockedLogsQueryRunner({ variables: [accountVariable] });

      const query: CloudWatchLogsQuery = {
        ...validLogsQuery,
        selectedAccountIds: ['$accounts'],
      };

      expect(() => runner.interpolateLogsQueryVariables(query, {})).toThrow(
        `Expanded account count (21) exceeds maximum of ${LOG_GROUP_ACCOUNT_MAX}`
      );
    });
  });

  describe('getLogRowContext', () => {
    it('replaces parameters correctly in the query', async () => {
      const { runner, queryMock } = setupMockedLogsQueryRunner({ variables: [regionVariable] });
      const row: LogRowModel = {
        entryFieldIndex: 0,
        rowIndex: 0,
        dataFrame: {
          refId: 'B',
          length: 1,
          fields: [
            { name: 'ts', type: FieldType.time, values: [1], config: {} },
            { name: LOG_IDENTIFIER_INTERNAL, type: FieldType.string, values: ['foo'], labels: {}, config: {} },
            { name: LOGSTREAM_IDENTIFIER_INTERNAL, type: FieldType.string, values: ['bar'], labels: {}, config: {} },
          ],
        },
        entry: '4',
        labels: {},
        hasAnsi: false,
        hasUnescapedContent: false,
        raw: '4',
        logLevel: LogLevel.info,
        timeEpochMs: 4,
        timeEpochNs: '4000000',
        timeFromNow: '',
        timeLocal: '',
        timeUtc: '',
        uid: '1',
      };
      await runner.getLogRowContext(row, undefined, queryMock);
      expect(queryMock.mock.calls[0][0].targets[0].endTime).toBe(4);
      // sets the default region if region is empty
      expect(queryMock.mock.calls[0][0].targets[0].region).toBe('us-west-1');

      await runner.getLogRowContext(row, { direction: LogRowContextQueryDirection.Forward }, queryMock, {
        ...validLogsQuery,
        region: '$region',
      });
      expect(queryMock.mock.calls[1][0].targets[0].startTime).toBe(4);
      expect(queryMock.mock.calls[1][0].targets[0].region).toBe('templatedRegion');
    });
  });

  describe('filterQuery', () => {
    it('allows query with namePrefix scope and at least one prefix', async () => {
      const { runner } = setupMockedLogsQueryRunner();

      const queryWithNamePrefix: CloudWatchLogsQuery = {
        ...validLogsQuery,
        logGroups: [],
        logsQueryScope: 'namePrefix',
        logGroupPrefixes: ['/aws/lambda/'],
      };

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: [queryWithNamePrefix],
      };

      const queryFn = jest
        .fn()
        .mockReturnValueOnce(of(startQuerySuccessResponseStub))
        .mockReturnValueOnce(of(getQuerySuccessResponseStub));
      await expect(runner.handleLogQueries([queryWithNamePrefix], options, queryFn)).toEmitValuesWith(() => {
        expect(queryFn).toHaveBeenCalled();
      });
    });

    it('filters out query with namePrefix scope but no prefixes', async () => {
      const { runner } = setupMockedLogsQueryRunner();

      const queryWithEmptyPrefixes: CloudWatchLogsQuery = {
        ...validLogsQuery,
        logGroups: [],
        logsQueryScope: 'namePrefix',
        logGroupPrefixes: [],
      };

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: [queryWithEmptyPrefixes],
      };

      const queryFn = jest.fn().mockReturnValue(of({ data: [] }));
      const response = runner.handleLogQueries([queryWithEmptyPrefixes], options, queryFn);
      await expect(response).toEmitValuesWith(() => {
        expect(queryFn).toHaveBeenCalledWith(
          expect.objectContaining({
            targets: [],
          })
        );
      });
    });

    it('allows query with allLogGroups scope without log groups', async () => {
      const { runner } = setupMockedLogsQueryRunner();

      const queryWithAllLogGroups: CloudWatchLogsQuery = {
        ...validLogsQuery,
        logGroups: [],
        logsQueryScope: 'allLogGroups',
      };

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: [queryWithAllLogGroups],
      };

      const queryFn = jest
        .fn()
        .mockReturnValueOnce(of(startQuerySuccessResponseStub))
        .mockReturnValueOnce(of(getQuerySuccessResponseStub));
      await expect(runner.handleLogQueries([queryWithAllLogGroups], options, queryFn)).toEmitValuesWith(() => {
        expect(queryFn).toHaveBeenCalled();
      });
    });

    it('filters out non-CWLI query when namePrefix scope would otherwise bypass missing log groups', async () => {
      const { runner } = setupMockedLogsQueryRunner();

      const queryWithNamePrefix: CloudWatchLogsQuery = {
        ...validLogsQuery,
        queryLanguage: LogsQueryLanguage.PPL,
        logGroups: [],
        logsQueryScope: 'namePrefix',
        logGroupPrefixes: ['/aws/lambda/'],
      };

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: [queryWithNamePrefix],
      };

      const queryFn = jest.fn().mockReturnValue(of({ data: [] }));
      const response = runner.handleLogQueries([queryWithNamePrefix], options, queryFn);
      await expect(response).toEmitValuesWith(() => {
        expect(queryFn).toHaveBeenCalledWith(
          expect.objectContaining({
            targets: [],
          })
        );
      });
    });

    it('filters out non-CWLI query when allLogGroups scope would otherwise bypass missing log groups', async () => {
      const { runner } = setupMockedLogsQueryRunner();

      const queryWithAllLogGroups: CloudWatchLogsQuery = {
        ...validLogsQuery,
        queryLanguage: LogsQueryLanguage.PPL,
        logGroups: [],
        logsQueryScope: 'allLogGroups',
      };

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: [queryWithAllLogGroups],
      };

      const queryFn = jest.fn().mockReturnValue(of({ data: [] }));
      const response = runner.handleLogQueries([queryWithAllLogGroups], options, queryFn);
      await expect(response).toEmitValuesWith(() => {
        expect(queryFn).toHaveBeenCalledWith(
          expect.objectContaining({
            targets: [],
          })
        );
      });
    });
  });

  describe('handleLogQueries', () => {
    it('appends -logs to the requestId', async () => {
      const { runner, queryMock } = setupMockedLogsQueryRunner();

      const request = {
        ...LogsRequestMock,
        requestId: 'mockId',
      };
      await expect(runner.handleLogQueries(LogsRequestMock.targets, request, queryMock)).toEmitValuesWith(() => {
        expect(queryMock.mock.calls[0][0].requestId).toEqual('mockId-logs');
      });
    });

    it('does not append -logs to the requestId if requestId is not provided', async () => {
      const { runner, queryMock } = setupMockedLogsQueryRunner();

      const request = {
        ...LogsRequestMock,
      };
      await expect(runner.handleLogQueries(LogsRequestMock.targets, request, queryMock)).toEmitValuesWith(() => {
        expect(queryMock.mock.calls[0][0].requestId).toEqual('');
      });
    });

    it('should request to start each query and then request to get the query results', async () => {
      const { runner } = setupMockedLogsQueryRunner();

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: rawLogQueriesStub,
      };

      const queryFn = jest
        .fn()
        .mockReturnValueOnce(of(startQuerySuccessResponseStub))
        .mockReturnValueOnce(of(getQuerySuccessResponseStub));

      const response = runner.handleLogQueries(rawLogQueriesStub, options, queryFn);
      const results = await lastValueFrom(response);
      expect(queryFn).toHaveBeenCalledTimes(2);
      expect(queryFn).toHaveBeenNthCalledWith(
        1,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StartQuery' })]),
        })
      );
      expect(queryFn).toHaveBeenNthCalledWith(
        2,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );

      expect(results).toEqual({
        ...getQuerySuccessResponseStub,
        errors: [],
        key: 'test-key',
      });
    });

    it('should call getQueryResults until the query returns with a status of complete', async () => {
      const { runner } = setupMockedLogsQueryRunner();

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: rawLogQueriesStub,
      };

      const queryFn = jest
        .fn()
        .mockReturnValueOnce(of(startQuerySuccessResponseStub))
        .mockReturnValueOnce(of(getQueryLoadingResponseStub))
        .mockReturnValueOnce(of(getQueryLoadingResponseStub))
        .mockReturnValueOnce(of(getQueryLoadingResponseStub))
        .mockReturnValueOnce(of(getQuerySuccessResponseStub));

      const response = runner.handleLogQueries(rawLogQueriesStub, options, queryFn);
      const results = await lastValueFrom(response);
      expect(queryFn).toHaveBeenCalledTimes(5);

      // first call to start query
      expect(queryFn).toHaveBeenNthCalledWith(
        1,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StartQuery' })]),
        })
      );
      // second call we try to get the results
      expect(queryFn).toHaveBeenNthCalledWith(
        2,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );
      // after getting a loading response we wait and try again
      expect(queryFn).toHaveBeenNthCalledWith(
        3,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );
      // after getting a loading response we wait and try again
      expect(queryFn).toHaveBeenNthCalledWith(
        4,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );
      // after getting a loading response we wait and try again
      expect(queryFn).toHaveBeenNthCalledWith(
        5,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );

      expect(results).toEqual({
        ...getQuerySuccessResponseStub,
        errors: [],
        key: 'test-key',
      });
    });

    it('should call getQueryResults until the query returns even if it the startQuery gets a rate limiting error from aws', async () => {
      const { runner } = setupMockedLogsQueryRunner();

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: rawLogQueriesStub,
      };

      const queryFn = jest
        .fn()
        .mockReturnValueOnce(of(startQueryErrorWhenRateLimitedResponseStub))
        .mockReturnValueOnce(of(startQuerySuccessResponseStub))
        .mockReturnValueOnce(of(getQuerySuccessResponseStub));

      const response = runner.handleLogQueries(rawLogQueriesStub, options, queryFn);
      const results = await lastValueFrom(response);
      expect(queryFn).toHaveBeenCalledTimes(3);

      // first call
      expect(queryFn).toHaveBeenNthCalledWith(
        1,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StartQuery' })]),
        })
      );
      // we retry because the first call failed with the rate limiting error
      expect(queryFn).toHaveBeenNthCalledWith(
        2,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StartQuery' })]),
        })
      );
      // we get results because second call was successful
      expect(queryFn).toHaveBeenNthCalledWith(
        3,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );

      expect(results).toEqual({
        ...getQuerySuccessResponseStub,
        errors: [],
        key: 'test-key',
      });
    });

    it('should call getQueryResults until the query returns even if it the startQuery gets a throttling error from aws', async () => {
      const { runner } = setupMockedLogsQueryRunner();

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: rawLogQueriesStub,
      };

      const queryFn = jest
        .fn()
        .mockReturnValueOnce(of(startQueryErrorWhenThrottlingResponseStub))
        .mockReturnValueOnce(of(startQuerySuccessResponseStub))
        .mockReturnValueOnce(of(getQuerySuccessResponseStub));

      const response = runner.handleLogQueries(rawLogQueriesStub, options, queryFn);
      const results = await lastValueFrom(response);
      expect(queryFn).toHaveBeenCalledTimes(3);

      // first call
      expect(queryFn).toHaveBeenNthCalledWith(
        1,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StartQuery' })]),
        })
      );
      // we retry because the first call failed with the rate limiting error
      expect(queryFn).toHaveBeenNthCalledWith(
        2,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StartQuery' })]),
        })
      );
      // we get results because second call was successful
      expect(queryFn).toHaveBeenNthCalledWith(
        3,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );

      expect(results).toEqual({
        ...getQuerySuccessResponseStub,
        errors: [],
        key: 'test-key',
      });
    });

    it('should return an error if it timesout before the start queries can get past a rate limiting error', async () => {
      const { runner } = setupMockedLogsQueryRunner();
      // first time timeout is called it will not be timed out, second time it will be timed out
      const timeoutFunc = jest
        .fn()
        .mockImplementationOnce(() => false)
        .mockImplementationOnce(() => true);
      runner.createTimeoutFn = jest.fn(() => timeoutFunc);

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: rawLogQueriesStub,
      };

      // running query fn will always return the rate limit
      const queryFn = jest.fn().mockReturnValue(of(startQueryErrorWhenRateLimitedResponseStub));

      const response = runner.handleLogQueries(rawLogQueriesStub, options, queryFn);
      const results = await lastValueFrom(response);

      expect(queryFn).toHaveBeenCalledTimes(2);

      // first call starts the query, but it fails with rate limiting error
      expect(queryFn).toHaveBeenNthCalledWith(
        1,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StartQuery' })]),
        })
      );

      // we retry because the first call failed with the rate limiting error and we haven't timed out yet
      expect(queryFn).toHaveBeenNthCalledWith(
        2,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StartQuery' })]),
        })
      );

      expect(results).toEqual({
        ...startQueryErrorWhenRateLimitedResponseStub,
        key: 'test-key',
        state: 'Done',
      });
    });

    it('should return an error if the start query fails with an error that is not a rate limiting error', async () => {
      const { runner } = setupMockedLogsQueryRunner();

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: rawLogQueriesStub,
      };

      const queryFn = jest.fn().mockReturnValueOnce(of(startQueryErrorWhenBadSyntaxResponseStub));

      const response = runner.handleLogQueries(rawLogQueriesStub, options, queryFn);
      const results = await lastValueFrom(response);

      // only one query is made, it gets the error and returns the error
      expect(queryFn).toHaveBeenCalledTimes(1);
      expect(queryFn).toHaveBeenNthCalledWith(
        1,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StartQuery' })]),
        })
      );
      expect(results).toEqual({
        ...startQueryErrorWhenBadSyntaxResponseStub,
        key: 'test-key',
        state: 'Done',
      });
    });

    it('should return an error and stop querying if get query results has finished with errors', async () => {
      const { runner } = setupMockedLogsQueryRunner();

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: rawLogQueriesStub,
      };

      const queryFn = jest
        .fn()
        .mockReturnValueOnce(of(startQuerySuccessResponseStub))
        .mockReturnValueOnce(of(getQueryLoadingResponseStub))
        .mockReturnValueOnce(of(getQueryErrorResponseStub))
        .mockReturnValueOnce(of(stopQueryResponseStub));

      const response = runner.handleLogQueries(rawLogQueriesStub, options, queryFn);
      const results = await lastValueFrom(response);

      expect(queryFn).toHaveBeenCalledTimes(4);
      expect(queryFn).toHaveBeenNthCalledWith(
        1,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StartQuery' })]),
        })
      );
      expect(queryFn).toHaveBeenNthCalledWith(
        2,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );
      expect(queryFn).toHaveBeenNthCalledWith(
        3,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );
      expect(queryFn).toHaveBeenNthCalledWith(
        4,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StopQuery' })]),
        })
      );
      expect(results).toEqual({
        ...getQueryErrorResponseStub,
        key: 'test-key',
        state: 'Done',
      });
    });

    it('should return an error and any partial data if it timesout before getting back all the results', async () => {
      const { runner } = setupMockedLogsQueryRunner();
      // mocking running for a while and then timing out
      const timeoutFunc = jest
        .fn()
        .mockImplementationOnce(() => false)
        .mockImplementationOnce(() => false)
        .mockImplementationOnce(() => false)
        .mockImplementationOnce(() => true);
      runner.createTimeoutFn = jest.fn(() => timeoutFunc);

      const queryFn = jest
        .fn()
        .mockReturnValueOnce(of(startQuerySuccessResponseStub))
        .mockReturnValueOnce(of(getQueryLoadingResponseStub))
        .mockReturnValueOnce(of(getQueryLoadingResponseStub))
        .mockReturnValueOnce(of(getQueryLoadingResponseStub))
        .mockReturnValueOnce(of(getQueryLoadingResponseStub))
        .mockReturnValueOnce(of(stopQueryResponseStub));

      const options: DataQueryRequest<CloudWatchLogsQuery> = {
        ...LogsRequestMock,
        targets: rawLogQueriesStub,
      };
      const response = runner.handleLogQueries(rawLogQueriesStub, options, queryFn);
      const results = await lastValueFrom(response);
      expect(queryFn).toHaveBeenCalledTimes(6);
      expect(queryFn).toHaveBeenNthCalledWith(
        1,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StartQuery' })]),
        })
      );
      expect(queryFn).toHaveBeenNthCalledWith(
        2,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );
      expect(queryFn).toHaveBeenNthCalledWith(
        3,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );
      expect(queryFn).toHaveBeenNthCalledWith(
        4,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );
      expect(queryFn).toHaveBeenNthCalledWith(
        5,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'GetQueryResults' })]),
        })
      );
      expect(queryFn).toHaveBeenNthCalledWith(
        6,
        expect.objectContaining({
          targets: expect.arrayContaining([expect.objectContaining({ subtype: 'StopQuery' })]),
        })
      );
      expect(results).toEqual({
        ...getQueryLoadingResponseStub,
        errors: [
          {
            message:
              'Error: Query hit timeout before completing after 3 attempts, partial results may be shown. To increase the timeout window update your datasource configuration.',
            refId: 'A',
            type: 'timeout',
          },
        ],
        key: 'test-key',
        state: 'Done',
      });
    });
  });

  describe('handleLogAnomaliesQueries', () => {
    it('appends -anomalies to the requestId', async () => {
      const { runner, queryMock } = setupMockedLogsQueryRunner();
      const logsAnomaliesRequestMock: DataQueryRequest<CloudWatchLogsAnomaliesQuery> = {
        requestId: 'mockId',
        range: TimeRangeMock,
        rangeRaw: { from: TimeRangeMock.from, to: TimeRangeMock.to },
        targets: [
          {
            id: '1',
            logsMode: LogsMode.Anomalies,
            queryMode: 'Logs',
            refId: 'A',
            region: 'us-east-1',
          },
        ],
        interval: '',
        intervalMs: 0,
        scopedVars: { __interval: { value: '20s' } },
        timezone: '',
        app: '',
        startTime: 0,
      };
      await expect(
        runner.handleLogAnomaliesQueries(LogsRequestMock.targets, logsAnomaliesRequestMock, queryMock)
      ).toEmitValuesWith(() => {
        expect(queryMock.mock.calls[0][0].requestId).toEqual('mockId-logsAnomalies');
      });
    });

    it('processes log trend histogram data correctly', async () => {
      const response = structuredClone(anomaliesQueryResponse);

      convertTrendHistogramToSparkline(response);

      expect(response.data[0].fields.find((field: Field) => field.name === 'Log trend')).toEqual({
        name: 'Log trend',
        type: 'frame',
        config: {
          custom: {
            drawStyle: 'bars',
            cellOptions: {
              type: 'sparkline',
              hideValue: true,
            },
          },
        },
        values: [
          {
            name: 'Trend_row_0',
            length: 8,
            fields: [
              {
                name: 'time',
                type: 'time',
                values: [
                  1760454000000, 1760544000000, 1760724000000, 1761282000000, 1761300000000, 1761354000000,
                  1761372000000, 1761390000000,
                ],
                config: {},
              },
              {
                name: 'value',
                type: 'number',
                values: [81, 35, 35, 36, 36, 36, 72, 36],
                config: {},
              },
            ],
          },
          {
            name: 'Trend_row_1',
            length: 2,
            fields: [
              {
                name: 'time',
                type: 'time',
                values: [1760687665000, 1760687670000],
                config: {},
              },
              {
                name: 'value',
                type: 'number',
                values: [3, 3],
                config: {},
              },
            ],
          },
        ],
      });
    });

    it('replaces log trend histogram field at the same index in the frame', () => {
      const response = structuredClone(anomaliesQueryResponse);
      convertTrendHistogramToSparkline(response);
      expect(response.data[0].fields[4].name).toEqual('Log trend');
    });

    it('ignore invalid timestamps in log trend histogram', () => {
      const response = structuredClone(anomaliesQueryResponse);

      response.data[0].fields[4].values[1] = {
        invalidTimestamp: 3,
        '1760687670000': 3,
        anotherInvalidTimestamp: 2,
        '1760687670010': 3,
      };

      convertTrendHistogramToSparkline(response);

      expect(response.data[0].fields[4].values[1].fields[0].values.length).toEqual(2);
      expect(response.data[0].fields[4].values[1].fields[1].values.length).toEqual(2);
    });
  });
});

const rawLogQueriesStub: CloudWatchLogsQuery[] = [
  {
    refId: 'A',
    id: '',
    region: 'us-east-2',
    logGroups: [
      {
        accountId: 'accountId',
        arn: 'somearn',
        name: 'nameOfLogGroup',
      },
    ],
    queryMode: 'Logs',
    expression: 'fields @timestamp, @message |\n sort @timestamp desc |\n limit 20',
    datasource: {
      type: 'cloudwatch',
      uid: 'ff87aa43-7618-42ee-ae9c-4a405378728b',
    },
  },
];

const startQuerySuccessResponseStub = {
  data: [
    {
      name: 'A',
      refId: 'A',
      meta: {
        typeVersion: [0, 0],
        custom: { Region: 'us-east-2' },
      },
      fields: [
        {
          name: 'queryId',
          type: 'string',
          typeInfo: { frame: 'string' },
          config: {},
          values: ['123'],
          entities: {},
        },
      ],
      length: 1,
      state: 'Done',
    },
  ],
};

const startQueryErrorWhenRateLimitedResponseStub = {
  data: [],
  errors: [
    {
      refId: 'A',
      message:
        'failed to execute log action with subtype: StartQuery: LimitExceededException: LimitExceededException: Account maximum query concurrency limit of [30] reached.',
      status: 500,
    },
  ],
};

const startQueryErrorWhenThrottlingResponseStub = {
  data: [],
  errors: [
    {
      refId: 'A',
      message:
        'failed to execute log action with subtype: StartQuery: ThrottlingException: ThrottlingException: Rate exceeded',
      status: 500,
    },
  ],
};

const startQueryErrorWhenBadSyntaxResponseStub = {
  data: [],
  state: 'Error',
  errors: [
    {
      refId: 'A',
      message:
        'failed to execute log action with subtype: StartQuery: MalformedQueryException: unexpected symbol found bad at line 1 and position 843',
      status: 500,
    },
  ],
};

const getQuerySuccessResponseStub = {
  data: [
    {
      name: 'A',
      refId: 'A',
      meta: {
        custom: { Status: 'Complete' },
        typeVersion: [0, 0],
        stats: [
          { displayName: 'Bytes scanned', value: 1000 },
          { displayName: 'Records scanned', value: 1000 },
          { displayName: 'Records matched', value: 1000 },
        ],
      },
      fields: [
        {
          name: '@message',
          type: 'string',
          typeInfo: { frame: 'string' },
          config: {},
          values: ['some log'],
        },
      ],
      length: 1,
      state: 'Done',
    },
  ],
  state: 'Done',
};

const getQueryLoadingResponseStub = {
  data: [
    {
      name: 'A',
      refId: 'A',
      meta: {
        custom: { Status: 'Running' },
        typeVersion: [0, 0],
        stats: [
          { displayName: 'Bytes scanned', value: 1 },
          { displayName: 'Records scanned', value: 1 },
          { displayName: 'Records matched', value: 1 },
        ],
      },
      fields: [
        {
          name: '@message',
          type: 'string',
          typeInfo: { frame: 'string' },
          config: {},
          values: ['some log'],
        },
      ],
      length: 1,
      state: 'Done',
    },
  ],
  state: 'Done',
};

const getQueryErrorResponseStub = {
  data: [],
  errors: [
    {
      refId: 'A',
      message: 'failed to execute log action with subtype: GetQueryResults: AWS is down',
      status: 500,
    },
  ],
  state: 'Error',
};

const stopQueryResponseStub = {
  state: 'Done',
};

const anomaliesQueryResponse: DataQueryResponse = {
  data: [
    {
      name: 'Log anomalies',
      refId: 'A',
      meta: {
        preferredVisualisationType: 'table',
      },
      fields: [
        {
          name: 'state',
          type: 'string',
          typeInfo: {
            frame: 'string',
          },
          config: {
            displayName: 'State',
          },
          values: ['Active', 'Active'],
          entities: {},
        },
        {
          name: 'description',
          type: 'string',
          typeInfo: {
            frame: 'string',
          },
          config: {
            displayName: 'Anomaly',
          },
          values: [
            '50.0% increase in count of value "405" for "code"-3',
            '151.3% increase in count of value 1 for "dotnet_collection_count_total"-3',
          ],
          entities: {},
        },
        {
          name: 'priority',
          type: 'string',
          typeInfo: {
            frame: 'string',
          },
          config: {
            displayName: 'Priority',
          },
          values: ['MEDIUM', 'MEDIUM'],
          entities: {},
        },
        {
          name: 'patternString',
          type: 'string',
          typeInfo: {
            frame: 'string',
          },
          config: {
            displayName: 'Log Pattern',
          },
          values: [
            '{"ClusterName":"PetSite","Namespace":"default","Service":"service-petsite","Timestamp":<*>,"Version":<*>,"code":<*>,"container_name":"petsite","http_requests_received_total":<*>,"instance":<*>:<*>,"job":"kubernetes-service-endpoints","kubernetes_node":<*>,"method":<*>,"pod_name":<*>,"prom_metric_type":"counter"}',
            '{"ClusterName":"PetSite","Namespace":"default","Service":"service-petsite","Timestamp":<*>,"Version":<*>,"container_name":"petsite","dotnet_collection_count_total":<*>,"generation":<*>,"instance":<*>:<*>,"job":"kubernetes-service-endpoints","kubernetes_node":<*>,"pod_name":<*>,"prom_metric_type":"counter"}',
          ],
          entities: {},
        },
        {
          name: 'logTrend',
          type: 'other',
          typeInfo: {
            frame: 'json.RawMessage',
            nullable: true,
          },
          config: {
            displayName: 'Log Trend',
          },
          values: [
            {
              '1760454000000': 81,
              '1760544000000': 35,
              '1760724000000': 35,
              '1761282000000': 36,
              '1761300000000': 36,
              '1761354000000': 36,
              '1761372000000': 72,
              '1761390000000': 36,
            },
            {
              '1760687665000': 3,
              '1760687670000': 3,
            },
          ],
          entities: {},
        },
        {
          name: 'firstSeen',
          type: 'time',
          typeInfo: {
            frame: 'time.Time',
          },
          config: {
            displayName: 'First seen',
          },
          values: [1760462460000, 1760687640000],
          entities: {},
        },
        {
          name: 'lastSeen',
          type: 'time',
          typeInfo: {
            frame: 'time.Time',
          },
          config: {
            displayName: 'Last seen',
          },
          values: [1761393660000, 1760687940000],
          entities: {},
        },
        {
          name: 'suppressed',
          type: 'boolean',
          typeInfo: {
            frame: 'bool',
          },
          config: {
            displayName: 'Suppressed?',
          },
          values: [false, false],
          entities: {},
        },
        {
          name: 'logGroupArnList',
          type: 'string',
          typeInfo: {
            frame: 'string',
          },
          config: {
            displayName: 'Log Groups',
          },
          values: [
            'arn:aws:logs:us-east-2:569069006612:log-group:/aws/containerinsights/PetSite/prometheus',
            'arn:aws:logs:us-east-2:569069006612:log-group:/aws/containerinsights/PetSite/prometheus',
          ],
          entities: {},
        },
        {
          name: 'anomalyArn',
          type: 'string',
          typeInfo: {
            frame: 'string',
          },
          config: {
            displayName: 'Anomaly Arn',
          },
          values: [
            'arn:aws:logs:us-east-2:569069006612:anomaly-detector:dca8b129-d09d-4167-86e9-7bf62ede2f95',
            'arn:aws:logs:us-east-2:569069006612:anomaly-detector:dca8b129-d09d-4167-86e9-7bf62ede2f95',
          ],
          entities: {},
        },
      ],
      length: 2,
    },
  ],
};
