...

import { createAssistantContextItem } from '@grafana/assistant';
import {
  DataFrame,
  DataQueryResponse,
  DataQueryRequest,
  FieldType,
  MutableDataFrame,
  dateTime,
  PreferredVisualisationType,
} from '@grafana/data';

import { GrafanaPyroscopeDataQuery } from './dataquery.gen';
import { enrichDataFrameWithAssistantContentMapper } from './utils';

// Mock the createContext function
jest.mock('@grafana/assistant', () => ({
  createAssistantContextItem: jest.fn(),
}));

const mockCreateAssistantContextItem = createAssistantContextItem as jest.MockedFunction<
  typeof createAssistantContextItem
>;

describe('enrichDataFrameWithAssistantContentMapper', () => {
  beforeEach(() => {
    jest.clearAllMocks();
    mockCreateAssistantContextItem.mockImplementation((type, data) => ({
      type,
      data,
      node: { id: 'test-id', type: 'test', name: 'test-node', navigable: true },
      occurrences: [],
    }));
  });

afterEach(() => {
    jest.restoreAllMocks();
  });

const createMockRequest = (
    targets: Array<Partial<GrafanaPyroscopeDataQuery>> = []
  ): DataQueryRequest<GrafanaPyroscopeDataQuery> => ({
    requestId: 'test-request',
    interval: '1s',
    intervalMs: 1000,
    maxDataPoints: 1000,
    range: {
      from: dateTime('2023-01-01T00:00:00Z'),
      to: dateTime('2023-01-01T01:00:00Z'),
      raw: {
        from: 'now-1h',
        to: 'now',
      },
    },
    scopedVars: {},
    timezone: 'UTC',
    app: 'explore',
    startTime: Date.now(),
    targets: targets.map((target, index) => ({
      refId: `A${index}`,
      profileTypeId: 'cpu',
      labelSelector: '{service="test"}',
      groupBy: [],
      queryType: 'profile' as const,
      datasource: {
        uid: 'test-uid',
        type: 'grafana-pyroscope-datasource',
      },
      ...target,
    })) as GrafanaPyroscopeDataQuery[],
  });

const createMockDataFrame = (refId: string, preferredVisualisationType?: PreferredVisualisationType): DataFrame => {
    const frame = new MutableDataFrame({
      refId,
      fields: [
        { name: 'time', type: FieldType.time, values: [1672531200000] },
        { name: 'value', type: FieldType.number, values: [100] },
      ],
    });

if (preferredVisualisationType) {
      frame.meta = {
        preferredVisualisationType,
      };
    }

return frame;
  };

const createMockResponse = (dataFrames: DataFrame[]): DataQueryResponse => ({
    data: dataFrames,
  });

describe('when processing flamegraph data frames', () => {
    it('should enrich flamegraph data frame with assistant context', () => {
      const request = createMockRequest([
        {
          refId: 'A0',
          profileTypeId: 'cpu',
          labelSelector: '{service="test"}',
        },
      ]);
      const dataFrame = createMockDataFrame('A0', 'flamegraph');
      const response = createMockResponse([dataFrame]);

const mapper = enrichDataFrameWithAssistantContentMapper(request, 'PyroscopeDatasource');
      const result = mapper(response);

expect(result.data).toHaveLength(1);
      expect(result.data[0].meta?.custom?.assistantContext).toBeDefined();
      expect(mockCreateAssistantContextItem).toHaveBeenCalledTimes(2);

// Verify datasource context
      expect(mockCreateAssistantContextItem).toHaveBeenCalledWith('datasource', {
        datasourceUid: 'test-uid',
      });

// Verify structured context
      expect(mockCreateAssistantContextItem).toHaveBeenCalledWith('structured', {
        title: 'Analyze Flame Graph',
        data: {
          start: request.range.from.valueOf(),
          end: request.range.to.valueOf(),
          profile_type_id: 'cpu',
          label_selector: '{service="test"}',
          operation: 'execute',
        },
      });
    });

it('should preserve existing meta.custom properties', () => {
      const request = createMockRequest([
        {
          refId: 'A0',
          profileTypeId: 'memory',
          labelSelector: '{app="backend"}',
        },
      ]);
      const dataFrame = createMockDataFrame('A0', 'flamegraph');
      dataFrame.meta = {
        preferredVisualisationType: 'flamegraph',
        custom: {
          existingProperty: 'value',
        },
      };

const response = createMockResponse([dataFrame]);
      const mapper = enrichDataFrameWithAssistantContentMapper(request, 'TestDatasource');
      const result = mapper(response);

expect(result.data[0].meta?.custom?.existingProperty).toBe('value');
      expect(result.data[0].meta?.custom?.assistantContext).toBeDefined();
    });

it('should handle multiple flamegraph data frames', () => {
      const request = createMockRequest([
        {
          refId: 'A0',
          profileTypeId: 'cpu',
          labelSelector: '{service="test1"}',
        },
        {
          refId: 'A1',
          profileTypeId: 'memory',
          labelSelector: '{service="test2"}',
        },
      ]);

const dataFrames = [createMockDataFrame('A0', 'flamegraph'), createMockDataFrame('A1', 'flamegraph')];
      const response = createMockResponse(dataFrames);