/*
 * (c) Copyright IBM Corp. 2021
 * (c) Copyright Instana Inc. and contributors 2018
 */

'use strict';

/**
 * @typedef {import('@supertenant/superbrain/types/binding.cjs').SpanLabels} SpanLabels
 * @typedef {import('@supertenant/superbrain/types/binding.cjs').OpenSpanResult} OpenSpanResult
 * @typedef {import('@supertenant/superbrain/types/binding.cjs').JSOpenSpanResult &
*  {canceled?: boolean, executed?: boolean}
* } JSOpenSpanResult
*/

const shimmer = require('shimmer');

const requireHook = require('../../../util/requireHook');
const tracingUtil = require('../../tracingUtil');
const constants = require('../../constants');
const cls = require('../../cls');

const { superbrain } = require('@supertenant/superbrain');
const superconsts = require('@supertenant/superconsts');
const { getOrCreateTask } = require('../../taskManager');

let isActive = true;

exports.spanName = 'redis';
exports.batchable = true;

exports.activate = function activate() {
  isActive = true;
};

exports.deactivate = function deactivate() {
  isActive = false;
};

exports.init = function init() {
  requireHook.onModuleLoad('ioredis', instrument);
};

function instrument(ioredis) {
  shimmer.wrap(ioredis.prototype, 'sendCommand', instrumentSendCommand);
  shimmer.wrap(ioredis.prototype, 'multi', instrumentMultiCommand);
  shimmer.wrap(ioredis.prototype, 'pipeline', instrumentPipelineCommand);
}

function instrumentSendCommand(original) {
  return function wrappedInternalSendCommand(command) {
    const client = this;
    const parentSpan = cls.getCurrentSpan();
    let callback;

    if (
      command.promise == null ||
      typeof command.name !== 'string' ||
      cls.skipExitTracing({ isActive, skipParentSpanCheck: true })
    ) {
      return original.apply(this, arguments);
    }

    if (
      // I wanted to keep instana parentSpan for its functionality here, but we don't have root spans always so it can be undefined
      parentSpan &&
      parentSpan.n === exports.spanName &&
      (parentSpan.data.redis.command === 'multi' || parentSpan.data.redis.command === 'pipeline') &&
      // the multi call is handled in instrumentMultiCommand but since multi is also send to Redis it will also
      // trigger instrumentSendCommand, which is why we filter it out.
      command.name !== 'multi'
    ) {
      // multi commands could actually be recorded as multiple spans, but we only want to record one
      // batched span considering that a multi call represents a transaction.
      // The same is true for pipeline calls, but they have a slightly different semantic.
      const parentSpanSubCommands = (parentSpan.data.redis.subCommands = parentSpan.data.redis.subCommands || []);
      parentSpanSubCommands.push(command.name);
    } else if (constants.isExitSpan(parentSpan)) {
      // Apart from the special case of multi/pipeline calls, redis exits can't be child spans of other exits.
      return original.apply(this, arguments);
    }

    const argsForOriginal = arguments;
    return cls.ns.runAndReturn(() => {
      const span = cls.startSpan(exports.spanName, constants.EXIT);
      span.data.redis = {
        connection: `${client.options.host}:${client.options.port}`,
        command: command.name.toLowerCase()
      };

      /** @type {JSOpenSpanResult} */
      let openSpanResult = null;
      const stTaskId = getOrCreateTask();
      if (stTaskId !== 0) {
        /** @type {SpanLabels} */
        const stSpanData = {};
        // TODO: Add integration to differentiate between redis and ioredis
        stSpanData[superconsts.Label.SupertenantResourceType] = superconsts.ResourceType.Redis;
        stSpanData[superconsts.Label.IntegrationModuleResourceId] = client.options.host;
        stSpanData[superconsts.Label.DbDatabase] = client.options.db;
        stSpanData[superconsts.Label.DbHost] = client.options.host;
        stSpanData[superconsts.Label.DbPort] = client.options.port;
        stSpanData[superconsts.Label.DbCommand] = span.data.redis.command;
        openSpanResult = superbrain.openSpan(stTaskId, superconsts.SpanType.ClientRequest, stSpanData);
      }

      callback = cls.ns.bind(onResult);
      command.promise.then(
        // make sure that the first parameter is never truthy
        callback.bind(null, null),
        callback
      );

      return original.apply(client, argsForOriginal);

      function onResult(error) {
        /** @type {SpanLabels} */
        const stSpanData = {};

        // multi commands are ended by exec. Wait for the exec result
        if (command.name === 'multi') {
          return;
        }

        span.d = Date.now() - span.ts;

        if (error) {
          span.ec = 1;
          span.data.redis.error = error.message;
          stSpanData[superconsts.Label.SupertenantError] = 'true';
        }

        if (openSpanResult != null) {
          // TODO: When adding actions, check canceled here.
          superbrain.closeSpan(openSpanResult.spanId, stSpanData);
        }

        span.transmit();
      }
    });
  };
}

function instrumentMultiCommand(original) {
  return instrumentMultiOrPipelineCommand('multi', original);
}

function instrumentPipelineCommand(original) {
  return instrumentMultiOrPipelineCommand('pipeline', original);
}

function instrumentMultiOrPipelineCommand(commandName, original) {
  return function wrappedInternalMultiOrPipelineCommand() {
    const client = this;
    const parentSpan = cls.getCurrentSpan();

    // NOTE: multiple redis transaction can have a parent ioredis call
    if (cls.skipExitTracing({ isActive, skipParentSpanCheck: true }) || constants.isExitSpan(parentSpan)) {
      return original.apply(this, arguments);
    }

    // create a new cls context parent to track the multi/pipeline child calls
    const clsContextForMultiOrPipeline = cls.ns.createContext();
    cls.ns.enter(clsContextForMultiOrPipeline);
    const span = cls.startSpan(exports.spanName, constants.EXIT);
    span.data.redis = {
      connection: `${client.options.host}:${client.options.port}`,
      command: commandName
    };

      /** @type {JSOpenSpanResult} */
      let openSpanResult = null;
      const stTaskId = getOrCreateTask();
      if (stTaskId !== 0) {
        /** @type {SpanLabels} */
        const stSpanData = {};
        // TODO: Add integration to differentiate between redis and ioredis
        stSpanData[superconsts.Label.SupertenantResourceType] = superconsts.ResourceType.Redis;
        stSpanData[superconsts.Label.IntegrationModuleResourceId] = client.options.host;
        stSpanData[superconsts.Label.DbDatabase] = client.options.db;
        stSpanData[superconsts.Label.DbHost] = client.options.host;
        stSpanData[superconsts.Label.DbPort] = client.options.port;
        stSpanData[superconsts.Label.DbCommand] = span.data.redis.command;
        // stSpanData[superconsts.Label.DbSubCommand] = span.data.redis.subCommands.toString();
        openSpanResult = superbrain.openSpan(stTaskId, superconsts.SpanType.ClientRequest, stSpanData);
      }

    const multiOrPipeline = original.apply(this, arguments);
    shimmer.wrap(
      multiOrPipeline,
      'exec',
      instrumentMultiOrPipelineExec.bind(null, clsContextForMultiOrPipeline, commandName, span, openSpanResult)
    );
    return multiOrPipeline;
  };
}

function instrumentMultiOrPipelineExec(clsContextForMultiOrPipeline, commandName, span, original, openSpanResult) {
  const endCallback = commandName === 'pipeline' ? pipelineCommandEndCallback : multiCommandEndCallback;
  return function instrumentedExec() {
    // the exec call is actually when the transmission of these commands to
    // redis is happening
    span.ts = Date.now();

    const result = original.apply(this, arguments);
    if (result.then) {
      result.then(
        results => {
          endCallback.call(null, clsContextForMultiOrPipeline, span, null, results, openSpanResult);
        },
        error => {
          endCallback.call(null, clsContextForMultiOrPipeline, span, error, [], openSpanResult);
        }
      );
    }
    return result;
  };
}

function multiCommandEndCallback(clsContextForMultiOrPipeline, span, error, openSpanResult) {
  /** @type {SpanLabels} */
  const stSpanData = {};

  span.d = Date.now() - span.ts;

  const subCommands = span.data.redis.subCommands;
  let commandCount = 1;
  if (subCommands) {
    // remove exec call
    subCommands.pop();
    commandCount = subCommands.length;
  }

  span.b = {
    s: commandCount
  };

  if (error) {
    span.ec = commandCount;
    span.data.redis.error = error.message;
    stSpanData[superconsts.Label.SupertenantError] = 'true';
  }

  if (openSpanResult != null) {
    // TODO: When adding actions, check canceled here.
    superbrain.closeSpan(openSpanResult.spanId, stSpanData);
  }

  span.transmit();
  cls.ns.exit(clsContextForMultiOrPipeline);
}

function pipelineCommandEndCallback(clsContextForMultiOrPipeline, span, error, results, openSpanResult) {
  /** @type {SpanLabels} */
  const stSpanData = {};

  span.d = Date.now() - span.ts;

  const subCommands = span.data.redis.subCommands;
  const commandCount = subCommands ? subCommands.length : 1;

  span.b = {
    s: commandCount
  };

  if (error) {
    // ioredis docs mention that this should never be possible, but better be safe than sorry
    span.ec = commandCount;
    span.data.redis.error = tracingUtil.getErrorDetails(error);
  } else {
    let numberOfErrors = 0;
    let sampledError;

    // results is an array of the form
    // [[?Error, ?Response]]
    for (let i = 0; i < results.length; i++) {
      if (results[i][0]) {
        numberOfErrors += 1;
        sampledError = sampledError || results[i][0];
      }
    }

    if (numberOfErrors > 0) {
      span.ec = numberOfErrors;
      span.data.redis.error = tracingUtil.getErrorDetails(sampledError);
    }
  }

  if (span.data.redis.error) {
    // We do this to mark error if subcommand had error but batch succeeded
    stSpanData[superconsts.Label.SupertenantError] = 'true';
  }

  if (openSpanResult != null) {
    // TODO: When adding actions, check canceled here.
    superbrain.closeSpan(openSpanResult.spanId, stSpanData);
  }

  span.transmit();
  cls.ns.exit(clsContextForMultiOrPipeline);
}