/**
Copyright: 2024 Contrast Security, Inc
  Contact: support@contrastsecurity.com
  License: Commercial

  NOTICE: This Software and the patented inventions embodied within may only be
  used as part of Contrast Security’s commercial offerings. Even though it is
  made available through public repositories, use of this Software is subject to
  the applicable End User Licensing Agreement found at
  https://www.contrastsecurity.com/enduser-terms-0317a or as otherwise agreed
  between Contrast Security and the End User. The Software may not be reverse
  engineered, modified, repackaged, sold, redistributed or otherwise used in a
  way not consistent with the End User License Agreement.
*/
import { fileURLToPath } from 'url';
import { createRequire } from 'module';
import { promises as fs } from 'fs';

const require = createRequire(import.meta.url);
const loader = require('./agent-loader.js');

const { enabled } = await loader.init(process.argv);
if (enabled) {
  await loader.bootstrap(process.argv);
}
await loader.resetArgs(process.argv[0], process.argv[1]);

const agent = require('./lib/agent.js');
const logger = require('./lib/core/logger/index.js')('contrast:esm-loaders');
const rewriter = require('./lib/core/rewrite/index.js')(agent);
const helpers = require('./lib/hooks/module/helpers.js');
const getType = require('./lib/util/get-file-type.js');

const loadedFromCache = new Set();

/**
 * The `getSource` hook is used to provide a custom method for retrieving source
 * code. In our case, we check for previously rewritten ESM files in our cache
 * location and return the string content of said files if they exist. If no
 * cached file is located, we fall back to the default Node method for reading
 * source code.
 *
 * @see https://nodejs.org/docs/latest-v12.x/api/esm.html#esm_code_getsource_code_hook
 * @param {string} url
 * @param {{ format: string }} context
 * @param {Function} defaultGetSource
 * @returns {Promise<{ source: string | SharedArrayBuffer | Uint8Array }>}
 */
export async function getSource(url, context, defaultGetSource) {
  if (!enabled) return defaultGetSource(url, context, defaultGetSource);

  const filename = fileURLToPath(url);
  logger.debug('getSource %s', filename);
  try {
    const cached = helpers.find(agent, filename);
    if (!cached) return defaultGetSource(url, context, defaultGetSource);
    loadedFromCache.add(filename);
    return { source: cached };
  } catch (err) {
    logger.warn(
      'Failed to locate cached rewrite for %s, err: %o',
      filename,
      err
    );

    return defaultGetSource(url, context, defaultGetSource);
  }
}

/**
 * The `transformSource` hook implements our source-code rewriting whenever an
 * ECMAScript Module file is loaded. Before the string content of a loaded ESM
 * file is parsed by Node, we rewrite in a manner similar to our current
 * `Module._compile` instrumentation. If a cached file was returned by the
 * `getSource` hook above, we can skip rewriting by returning the cached result.
 *
 * @see https://nodejs.org/docs/latest-v12.x/api/esm.html#esm_code_transformsource_code_hook
 * @param {string | SharedArrayBuffer | Uint8Array} source
 * @param {{ format: string, url: string }} context
 * @param {Function} defaultTransformSource
 * @returns {Promise<{ source: string | SharedArrayBuffer | Uint8Array }>}
 */
export async function transformSource(source, context, defaultTransformSource) {
  if (!enabled)
    return defaultTransformSource(source, context, defaultTransformSource);

  const filename = fileURLToPath(context.url);
  logger.debug('transformSource %s', filename);
  let result;
  try {
    if (loadedFromCache.has(filename)) return { source };

    result = rewriter.rewriteFile(source, filename, { sourceType: 'module' });
    helpers.cacheWithSourceMap(agent, filename, result);
    return { source: result.code };
  } catch (err) {
    logger.error(
      'Failed to compile rewritten code for %s, err: %o, rewritten code %s, compiling original code.',
      filename,
      err,
      result.code
    );
    return defaultTransformSource(source, context, defaultTransformSource);
  }
}

/**
 * For Node 16 and above, the 'getFormat', 'getSource' and 'transformSource' have been
 * consolidated into one 'load' hook. The logic is similar to that of transformSource
 * except that the source is not provided and must be either read in from the file provided
 * or accessed from the cache.
 *
 * @see https://nodejs.org/dist/latest-v16.x/docs/api/esm.html#loadurl-context-defaultload
 * @param {string} url
 * @param {{ format: string, url: string }} context
 * @param {Function} defaultLoad
 * @returns {Promise<{ format: string, source: string | SharedArrayBuffer | Uint8Array }>}
 */
export async function load(url, context, defaultLoad) {
  if (!enabled) return defaultLoad(url, context, defaultLoad);

  const type = getType(url);

  if (type === 'builtin' || type === 'unknown') {
    logger.debug(
      'Skipping rewrite for %s module %s, loading original code',
      type,
      url
    );
    return defaultLoad(url, context, defaultLoad);
  }

  const filename = fileURLToPath(url);

  try {
    let result;
    const cached = helpers.find(agent, filename);
    if (cached) {
      result = { code: cached };
    } else {
      const source = await fs.readFile(filename, 'utf8');
      result = rewriter.rewriteFile(source, filename, {
        sourceType: type === 'commonjs' ? 'script' : 'module',
      });
      helpers.cacheWithSourceMap(agent, filename, result);
    }
    return { format: type, source: result.code, shortCircuit: true };
  } catch (err) {
    logger.error(
      'Failed to load rewritten code for %s, err: %o, rewritten code %s, loading original code.',
      filename,
      err
    );
    return defaultLoad(url, context, defaultLoad);
  }
}
