/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

import { ssms as __ssms } from './core/index.ts';
import createDebug from 'debug';

const { SSMS } = __ssms;

import { EventEmitter } from 'eventemitter3';
import { loadPlugins, loadPluginsConfig } from './load-plugins.ts';
import type { PluginLoadFailure } from './load-plugins.ts';

const debug = createDebug('core');

import { promisify as p } from 'node:util';
import _ from 'lodash';
import type { PeriodMetrics } from './core/ssms.ts';
import PlatformLambda from './platform/aws-lambda/index.ts';
import PlatformAzureACI from './platform/az/aci.ts';
import PlatformLocal from './platform/local/index.ts';

// The contract every load-generation platform implements. Worker
// operations beyond these are platform-internal.
export interface LoadPlatform {
  events: {
    on(event: string, listener: (...args: any[]) => void): unknown;
  };
  startJob(): Promise<unknown>;
  shutdown(): Promise<unknown>;
  getDesiredWorkerCount(): number;
}

export interface LauncherOpts {
  platform?: string;
  mode?: string;
  count?: number;
  platformConfig?: Record<string, string>;
  cliArgs?: Record<string, any>;
  testRunId?: string;
}

// Buffered worker message, grouped and printed by flushWorkerMessages.
interface BufferedWorkerMessage {
  id?: unknown;
  ts: number;
  error?: Error;
  msg?: string;
  level?: string;
  [key: string]: any;
}

async function createLauncher(
  script: { config: Record<string, any>; [key: string]: any },
  payload: unknown,
  opts: Record<string, any>,
  launcherOpts?: LauncherOpts
): Promise<Launcher | null> {
  launcherOpts = launcherOpts || {
    platform: 'local',
    mode: 'distribute'
  };
  let l: Launcher;
  try {
    l = new Launcher(script, payload, opts, launcherOpts);
  } catch (err) {
    console.log(err);
    return null;
  }

  return l;
}
class Launcher {
  declare script: { config: Record<string, any>; [key: string]: any };
  declare payload: unknown;
  declare opts: Record<string, any>;
  declare exitedWorkersCount: number;
  declare workerMessageBuffer: BufferedWorkerMessage[];
  declare metricsByPeriod: Record<string, PeriodMetrics[]>;
  declare finalReportsByWorker: Record<string, PeriodMetrics>;
  declare events: EventEmitter;
  declare pluginEvents: EventEmitter;
  declare pluginEventsLegacy: EventEmitter;
  declare launcherOpts: LauncherOpts;
  declare periodsReportedFor: string[];
  declare platform: LoadPlatform;
  declare phaseStartedEventsSeen: Record<string, number>;
  declare phaseCompletedEventsSeen: Record<string, number>;
  declare eventsByWorker: Record<string, unknown>;
  declare i1: ReturnType<typeof setInterval>;
  declare i2: ReturnType<typeof setInterval>;
  declare workerExitWatcher: ReturnType<typeof setInterval>;

  constructor(
    script: { config: Record<string, any>; [key: string]: any },
    payload: unknown,
    opts: Record<string, any>,
    launcherOpts: LauncherOpts
  ) {
    this.script = script;
    this.payload = payload;
    this.opts = opts;

    this.exitedWorkersCount = 0;
    this.workerMessageBuffer = [];

    this.metricsByPeriod = {}; // individual intermediates by worker
    this.finalReportsByWorker = {};

    this.events = new EventEmitter();

    this.pluginEvents = new EventEmitter();
    this.pluginEventsLegacy = new EventEmitter();

    this.launcherOpts = launcherOpts;

    this.periodsReportedFor = [];

    if (launcherOpts.platform === 'local') {
      this.platform = new PlatformLocal(script, payload, opts, launcherOpts);
    } else if (launcherOpts.platform === 'aws:lambda') {
      this.platform = new PlatformLambda(script, payload, opts, launcherOpts);
    } else if (launcherOpts.platform === 'az:aci') {
      this.platform = new PlatformAzureACI(script, payload, opts, launcherOpts);
    } else {
      throw new Error(`Unknown platform: ${launcherOpts.platform}`);
    }

    this.phaseStartedEventsSeen = {};
    this.phaseCompletedEventsSeen = {};

    this.eventsByWorker = {};
  }

  async initWorkerEvents(workerEvents: LoadPlatform['events']) {
    workerEvents.on('workerError', (_workerId, message) => {
      const { id, error, level, aggregatable, logs } = message;

      if (level !== 'warn') {
        this.exitedWorkersCount++;
      }

      if (aggregatable) {
        this.workerMessageBuffer.push(message);
      } else {
        global.artillery.log(`[${id}]: ${error.message}`);
        if (logs) {
          global.artillery.log(logs);
        }
      }

      this.events.emit('workerError', message);
    });

    workerEvents.on('phaseStarted', (_workerId, message) => {
      // Note - we send only the first event for a phase, not all of them
      if (
        typeof this.phaseStartedEventsSeen[message.phase.index] === 'undefined'
      ) {
        this.phaseStartedEventsSeen[message.phase.index] = Date.now();
        const fullPhase = {
          //get back original phase without any splitting for workers
          ...this.script.config.phases[message.phase.index],
          index: message.phase.index,
          id: message.phase.id,
          startTime: this.phaseStartedEventsSeen[message.phase.index]
        };

        this.events.emit('phaseStarted', fullPhase);
        this.pluginEvents.emit('phaseStarted', fullPhase);
        this.pluginEventsLegacy.emit('phaseStarted', fullPhase);

        global.artillery.globalEvents.emit('phaseStarted', fullPhase);
      }
    });

    workerEvents.on('phaseCompleted', (_workerId, message) => {
      if (
        typeof this.phaseCompletedEventsSeen[message.phase.index] ===
        'undefined'
      ) {
        this.phaseCompletedEventsSeen[message.phase.index] = Date.now();
        const fullPhase = {
          //get back original phase without any splitting for workers
          ...this.script.config.phases[message.phase.index],
          id: message.phase.id,
          index: message.phase.index,
          startTime: this.phaseStartedEventsSeen[message.phase.index],
          endTime: message.phase.endTime
        };

        this.events.emit('phaseCompleted', fullPhase);
        this.pluginEvents.emit('phaseCompleted', fullPhase);
        this.pluginEventsLegacy.emit('phaseCompleted', fullPhase);
        global.artillery.globalEvents.emit('phaseCompleted', fullPhase);
      }
    });

    // We are not going to receive stats events from workers
    // which have zero arrivals for a phase. (This can only happen
    // in "distribute" mode.)
    workerEvents.on('stats', (_workerId, message) => {
      const workerStats = SSMS.deserializeMetrics(message.stats);
      const period = workerStats.period;
      if (typeof this.metricsByPeriod[period] === 'undefined') {
        this.metricsByPeriod[period] = [];
      }
      // TODO: might want the full message here, with worker ID etc
      this.metricsByPeriod[period].push(workerStats);
    });

    workerEvents.on('done', async (workerId, message) => {
      this.exitedWorkersCount++;
      this.finalReportsByWorker[workerId] = SSMS.deserializeMetrics(
        message.report
      );
    });

    workerEvents.on('log', async (_workerId, message) => {
      artillery.globalEvents.emit('log', ...message.args);
    });

    workerEvents.on('setSuggestedExitCode', (_workerId, message) => {
      artillery.suggestedExitCode = message.code;
    });
  }

  async initPlugins() {
    const plugins = await loadPlugins(
      this.script.config.plugins,
      this.script,
      this.opts
    );

    //
    // init plugins
    //
    for (const [name, result] of Object.entries(plugins)) {
      if (result.isLoaded) {
        if (result.version === 3) {
          // TODO: load the plugin, subscribe to events
          // global.artillery.plugins[name] = result.plugin;
        } else {
          //           global.artillery.log(`WARNING: Legacy plugin detected: ${name}
          // See https://artillery.io/docs/resources/core/v2.html for more details.`,
          //                                'warn');

          // NOTE:
          // We are giving v1 and v2 plugins a throw-away script
          // object because we only care about the plugin setting
          // up event handlers here. The plugins will be loaded
          // properly in individual workers where they will have the
          // opportunity to attach custom code, modify the script
          // object etc.
          // If we let a plugin access to the actual script object,
          // and it happens to attach code to it (with a custom
          // processor function for example) - spawning a worker
          // will fail.
          const dummyScript = JSON.parse(JSON.stringify(this.script));
          dummyScript.config = {
            ...dummyScript.config,
            // Load additional plugins configuration from the environment
            plugins: loadPluginsConfig(this.script.config.plugins)
          };

          if (result.version === 1) {
            result.plugin = new result.PluginExport(
              dummyScript.config,
              this.pluginEventsLegacy
            );
            global.artillery.plugins.push(result);
          } else if (result.version === 2) {
            if (result.PluginExport.LEGACY_METRICS_FORMAT === false) {
              result.plugin = new result.PluginExport.Plugin(
                dummyScript,
                this.pluginEvents,
                this.opts
              );
            } else {
              result.plugin = new result.PluginExport.Plugin(
                dummyScript,
                this.pluginEventsLegacy,
                this.opts
              );
            }
            global.artillery.plugins.push(result);
          } else {
            // TODO: print warning
          }
        }
      } else {
        // Cast: the boolean discriminant does not narrow under the
        // loose (non-strictNullChecks) root typecheck.
        const failure = result as PluginLoadFailure;
        global.artillery.log(`WARNING: Could not load plugin: ${name}`, 'warn');
        global.artillery.log(failure.msg, 'warn');
        // global.artillery.log(failure.error, 'warn');
      }
    }
  }

  async handleAllWorkersFinished() {
    const allWorkersDone =
      this.exitedWorkersCount === this.platform.getDesiredWorkerCount();
    if (allWorkersDone) {
      clearInterval(this.i1);
      clearInterval(this.i2);

      // Flush messages from workers
      await this.flushWorkerMessages(0);
      await this.flushIntermediateMetrics(true);

      const pds = Object.keys(this.finalReportsByWorker).map(
        (k) => this.finalReportsByWorker[k]
      );

      const statsByPeriod = Object.values(SSMS.mergeBuckets(pds));
      const stats = SSMS.pack(statsByPeriod);

      stats.summaries = {};
      for (const [name, value] of Object.entries(stats.histograms || {})) {
        const summary = SSMS.summarizeHistogram(value);
        stats.summaries[name] = summary;
      }

      clearInterval(this.workerExitWatcher);

      // Relay event to workers
      this.pluginEvents.emit('done', stats);

      global.artillery.globalEvents.emit('done', stats);
      this.pluginEventsLegacy.emit('done', SSMS.legacyReport(stats));

      this.events.emit('done', stats);
    }
  }

  async flushWorkerMessages(maxAge = 9000) {
    // Collect messages older than maxAge msec and group by log message:
    const now = Date.now();
    const okToPrint = this.workerMessageBuffer.filter(
      (m) => now - m.ts > maxAge
    );
    this.workerMessageBuffer = this.workerMessageBuffer.filter(
      (m) => now - m.ts <= maxAge
    );

    const readyMessages = okToPrint.reduce(
      (acc: Record<string, BufferedWorkerMessage[]>, message) => {
        const { error } = message;
        // TODO: Take event type and level into account
        // NOTE: pre-existing behavior: throws when a buffered message
        // has no error property (only workerError messages are
        // buffered today).
        const key = (error as Error).message;
        if (typeof acc[key] === 'undefined') {
          acc[key] = [];
        }
        acc[key].push(message);
        return acc;
      },
      {}
    );

    for (const [_logMessage, messageObjects] of Object.entries(readyMessages)) {
      if (messageObjects[0].error) {
        global.artillery.log(
          `[${messageObjects[0].id}] ${messageObjects[0].error.message}`,
          messageObjects[0].level
        );
      } else {
        // Expect a msg property:
        global.artillery.log(
          `[${messageObjects[0].id}] ${messageObjects[0].msg}`,
          messageObjects[0].level
        );
      }
    }
  }

  async flushIntermediateMetrics(flushAll = false) {
    if (Object.keys(this.metricsByPeriod).length === 0) {
      debug('No metrics received yet');
      return;
    }

    // We always look at the earliest period available so that reports come in chronological order
    const unreportedPeriods = Object.keys(this.metricsByPeriod)
      .filter((x) => this.periodsReportedFor.indexOf(x) === -1)
      .sort();

    const earliestPeriodAvailable = unreportedPeriods[0];

    // TODO: better name. One above is earliestNotAlreadyReported
    const earliest = Object.keys(this.metricsByPeriod).sort()[0];
    if (this.periodsReportedFor.indexOf(earliest) > -1) {
      global.artillery.log(
        'Warning: multiple batches of metrics for period',
        earliest,
        new Date(Number(earliest))
      );

      delete this.metricsByPeriod[earliest]; // FIXME: need to merge them in for the final report
    }

    // Dynamically adjust the duration we're willing to wait for. This matters on SQS where messages are received
    // in batches of 10 and more workers => need to wait longer.
    const MAX_WAIT_FOR_PERIOD_MS =
      (Math.ceil(this.platform.getDesiredWorkerCount() / 10) * 3 + 30) * 1000;

    debug({
      now: Date.now(),
      count: this.platform.getDesiredWorkerCount(),
      earliestPeriodAvailable,
      earliest,
      MAX_WAIT_FOR_PERIOD_MS,
      numReports: this.metricsByPeriod[earliestPeriodAvailable]?.length,
      periodsReportedFor: this.periodsReportedFor,
      metricsByPeriod: Object.keys(this.metricsByPeriod)
    });

    const allWorkersReportedForPeriod =
      this.metricsByPeriod[earliestPeriodAvailable]?.length ===
      this.platform.getDesiredWorkerCount();
    const waitedLongEnough =
      Date.now() - Number(earliestPeriodAvailable) > MAX_WAIT_FOR_PERIOD_MS;

    if (flushAll) {
      for (const period of unreportedPeriods) {
        this.emitIntermediatesForPeriod(period);
      }
    } else if (
      typeof earliestPeriodAvailable !== 'undefined' &&
      (allWorkersReportedForPeriod || waitedLongEnough)
    ) {
      this.emitIntermediatesForPeriod(earliestPeriodAvailable);
      // TODO: autoscaling. Handle workers that drop off or join, and update count
    } else {
      debug('Waiting for more workerStats before emitting stats event');
    }
  }

  emitIntermediatesForPeriod(period: string) {
    debug(
      'Report @',
      new Date(Number(period)),
      'made up of items:',
      this.metricsByPeriod[String(period)].length
    );

    // TODO: Track how many workers provided metrics in the metrics report
    // summarize histograms for console reporter:
    const merged = SSMS.mergeBuckets(this.metricsByPeriod[String(period)]);
    const stats = merged[String(period)];

    stats.summaries = {};
    for (const [name, value] of Object.entries(stats.histograms || {})) {
      const summary = SSMS.summarizeHistogram(value);
      stats.summaries[name] = summary;
    }

    delete this.metricsByPeriod[String(period)];

    this.periodsReportedFor.push(period);
    this.pluginEvents.emit('stats', stats);
    global.artillery.globalEvents.emit('stats', stats);
    this.pluginEventsLegacy.emit('stats', SSMS.legacyReport(stats));

    this.events.emit('stats', stats);
  }

  async run() {
    await this.initPlugins();

    this.i1 = setInterval(async () => {
      await this.flushWorkerMessages();
    }, 1 * 1000).unref();

    this.i2 = setInterval(async () => {
      this.flushIntermediateMetrics();
    }, 2 * 1000).unref();

    this.workerExitWatcher = setInterval(async () => {
      await this.handleAllWorkersFinished();
    }, 2 * 1000);

    await this.initWorkerEvents(this.platform.events);
    await this.platform.startJob();
    debug('workers running');
  }

  async shutdown() {
    await this.platform.shutdown();

    // TODO: flush worker messages, and intermediate stats

    // Unload plugins
    // TODO: v3 plugins
    if (global.artillery?.plugins) {
      for (const o of global.artillery.plugins) {
        if (o.plugin.cleanup) {
          try {
            await p(o.plugin.cleanup.bind(o.plugin))();
            debug('plugin unloaded:', o.name);
          } catch (cleanupErr) {
            global.artillery.log(cleanupErr, 'error');
          }
        }
      }
    }
  }
}

export default createLauncher;