import { from, map, type Observable, of } from 'rxjs';

import {
  CustomVariableSupport,
  type DataQueryRequest,
  type DataQueryResponse,
  type MetricFindValue,
} from '@grafana/data';

import { getProfileTypeLabel, VariableQueryEditor } from './VariableQueryEditor';
import { type PyroscopeDataSource } from './datasource';
import { type ProfileTypeMessage, type VariableQuery } from './types';

export interface DataAPI {
  getProfileTypes(start: number, end: number): Promise<ProfileTypeMessage[]>;
  getLabelNames(query: string, start: number, end: number): Promise<string[]>;
  getLabelValues(query: string, label: string, start: number, end: number): Promise<string[]>;
}

export class VariableSupport extends CustomVariableSupport<PyroscopeDataSource> {
  constructor(private readonly dataAPI: DataAPI) {
    super();
  }

  editor = VariableQueryEditor;

  query(request: DataQueryRequest<VariableQuery>): Observable<DataQueryResponse> {
    if (request.targets[0].type === 'profileType') {
      return from(this.dataAPI.getProfileTypes(request.range.from.valueOf(), request.range.to.valueOf())).pipe(
        map((values) => {
          return { data: values.map<MetricFindValue>((v) => ({ text: v.label, value: v.id })) };
        })
      );
    }

    if (request.targets[0].type === 'label') {
      if (!request.targets[0].profileTypeId) {
        return of({ data: [] });
      }
      return from(
        this.dataAPI.getLabelNames(
          getProfileTypeLabel(request.targets[0].profileTypeId),
          request.range.from.valueOf(),
          request.range.to.valueOf()
        )
      ).pipe(
        map((values) => {
          return { data: values.map((v) => ({ text: v })) };
        })
      );
    }

    if (request.targets[0].type === 'labelValue') {
      if (!request.targets[0].labelName || !request.targets[0].profileTypeId) {
        return of({ data: [] });
      }
      return from(
        this.dataAPI.getLabelValues(
          getProfileTypeLabel(request.targets[0].profileTypeId),
          request.targets[0].labelName,
          request.range.from.valueOf(),
          request.range.to.valueOf()
        )
      ).pipe(
        map((values) => {
          return { data: values.map((v) => ({ text: v })) };
        })
      );
    }

    return of({ data: [] });
  }
}
