import type { AggregateOptions, Document } from 'mongodb';
import { cancellableWait, raceWithAbort } from '@mongodb-js/compass-utils';
import {
  FULL_SCAN_STAGES,
  REQUIRED_AS_FIRST_STAGE,
} from '@mongodb-js/mongodb-constants';
import isEqual from 'lodash/isEqual';
import type { PreferencesAccess } from '@mongodb-js/compass-preferences-model';
import type { DataService } from '../data-service';
import {
  getStageOperator,
  getLastStageOperator,
  isLastStageOutputStage,
} from '../../utils/stage';
import { aggregatePipeline } from '../../utils/cancellable-aggregation';

// Used in the `comment` on the aggregate command to help identify
// the operation in server logs and currentOp.
export const PREVIEW_AGGREGATION_COMMENT = 'Aggregation preview';

export const DEFAULT_SAMPLE_SIZE = 100000;

export const DEFAULT_PREVIEW_LIMIT = 10;

export const DEFAULT_PREVIEW_DEBOUNCE_MS = 700;

/**
 * Ops that must scan the entire results before moving to the
 * next stage.
 */
const FULL_SCAN_OPS = FULL_SCAN_STAGES.map((stage) => stage.value) as string[];

/**
 * Stage operators that are required to be the first stage.
 */
const REQUIRED_AS_FIRST_STAGE_OPS = REQUIRED_AS_FIRST_STAGE.map(
  (stage) => stage.value
) as string[];

export interface PreviewOptions extends AggregateOptions {
  debounceMs?: number;
  totalDocumentCount?: number;
  sampleSize?: number;
  previewSize?: number;
}

export function createPreviewAggregation(
  pipeline: Document[],
  options: Pick<
    PreviewOptions,
    'sampleSize' | 'previewSize' | 'totalDocumentCount'
  > = {}
) {
  if (isLastStageOutputStage(pipeline)) {
    throw new Error('Cannot preview pipeline with last stage as output stage');
  }

  const stages = [];
  for (const stage of pipeline) {
    if (
      (!options.totalDocumentCount ||
        options.totalDocumentCount >
          (options.sampleSize ?? DEFAULT_SAMPLE_SIZE)) &&
      // If stage can cause a full scan on the collection, prepend it with a
      // $limit
      FULL_SCAN_OPS.includes(getStageOperator(stage) ?? '')
    ) {
      stages.push({ $limit: options.sampleSize ?? DEFAULT_SAMPLE_SIZE });
    }
    stages.push(stage);
  }
  if (
    // TODO: super unsure what this is doing, half of these are not even
    // selectable stage operators in UI
    !REQUIRED_AS_FIRST_STAGE_OPS.includes(getLastStageOperator(stages))
  ) {
    stages.push({ $limit: options.previewSize ?? DEFAULT_PREVIEW_LIMIT });
  }
  return stages;
}

export class PipelinePreviewManager {
  private queue = new Map<number, AbortController>();
  private lastPipeline = new Map<number, Document[]>();
  private dataService: DataService;
  private preferences: PreferencesAccess;
  constructor(dataService: DataService, preferences: PreferencesAccess) {
    this.dataService = dataService;
    this.preferences = preferences;
  }

  /**
   * Request aggregation results with a default debounce
   */
  async getPreviewForStage(
    idx: number,
    namespace: string,
    getPipeline: () => Promise<Document[]>,
    getPreviewOptions: () => Promise<PreviewOptions>,
    force = false
  ): Promise<Document[]> {
    this.queue.get(idx)?.abort();
    const controller = new AbortController();
    this.queue.set(idx, controller);
    // Registered before the pipeline and options (collation) resolve so that
    // cancelling while they are still parsing aborts this request.
    const [
      pipeline,
      { sampleSize, previewSize, totalDocumentCount, ...options },
    ] = await raceWithAbort(
      Promise.all([getPipeline(), getPreviewOptions()]),
      controller.signal
    );
    if (!force) {
      await cancellableWait(
        options.debounceMs ?? DEFAULT_PREVIEW_DEBOUNCE_MS,
        controller.signal
      );
    }
    this.lastPipeline.set(idx, pipeline);
    const result = await aggregatePipeline({
      dataService: this.dataService,
      preferences: this.preferences,
      signal: controller.signal,
      namespace,
      pipeline: createPreviewAggregation(pipeline, {
        sampleSize,
        previewSize,
        totalDocumentCount,
      }),
      options: {
        comment: PREVIEW_AGGREGATION_COMMENT,
        ...options,
      },
    });
    this.queue.delete(idx);
    return result;
  }

  isLastPipelineEqual(idx: number, pipeline: Document[]): boolean {
    return isEqual(pipeline, this.lastPipeline.get(idx));
  }

  /**
   * Cancel aggregation request by id
   */
  cancelPreviewForStage(idx: number): boolean {
    this.queue.get(idx)?.abort();
    this.lastPipeline.delete(idx);
    return this.queue.delete(idx);
  }

  /**
   * Clear request queue starting from an index
   */
  clearQueue(from = 0) {
    for (const [idx, controller] of Array.from(this.queue.entries())) {
      if (idx >= from) {
        controller.abort();
        this.lastPipeline.delete(idx);
        this.queue.delete(idx);
      }
    }
  }
}