#!/usr/bin/env python3
# Copyright 2023 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import itertools
import json
import os
import glob
import platform
import re
import shutil
import sys

_HERE_PATH = os.path.dirname(__file__)
_SRC_PATH = os.path.normpath(os.path.join(_HERE_PATH, '..', '..', '..', '..'))
_CWD = os.getcwd()  # NOTE(dbeam): this is typically out/<gn_name>/.

sys.path.append(os.path.join(_SRC_PATH, 'third_party', 'node'))
import node
import node_modules

def _request_list_path(out_folder, target_name):
  # Using |target_name| as a prefix which is guaranteed to be unique within the
  # same folder, to avoid problems when multiple bundle_js() targets in the
  # same BUILD.gn file exist.
  return os.path.join(out_folder, target_name + '_requestlist.txt')


def _get_dep_path(dep, host_url, out_folder):
  if dep.startswith(host_url):
    return dep.replace(host_url, os.path.relpath(out_folder, _CWD))
  elif not (dep.startswith('chrome://') or dep.startswith('//')):
    return os.path.relpath(out_folder, _CWD) + '/' + dep
  return dep


# Get a list of all files that were bundled with rollup and update the
# depfile accordingly such that Ninja knows when to re-trigger.
def _update_dep_file(in_folder, args, out_file_path, manifest):
  in_path = os.path.join(_CWD, in_folder)

  # Gather the dependencies of all bundled root files.
  request_list = []
  for out_file in manifest:
    request_list += manifest[out_file]

  # Add a slash in front of every dependency that is not a chrome:// URL, so
  # that we can map it to the correct source file path below.
  request_list = map(
      lambda dep: _get_dep_path(dep, args.host_url, args.out_folder),
      request_list)

  deps = map(os.path.normpath, request_list)

  with open(
      os.path.join(_CWD, args.depfile), 'w', newline='', encoding='utf-8') as f:
    f.write(out_file_path + ': ' + ' '.join(deps))


# Autogenerate a rollup config file so that we can import the plugin and
# pass it information about the location of the directories and files to
# exclude from the bundle.
# Arguments:
# out_dir: The root directory for the output (i.e. corresponding to
#          host_url at runtime).
# in_path: Root directory for the input files.
# bundle_dir_path: Path to the directory holding the bundled output files
#                  relative to the root output directory. E.g. if bundle is
#                  chrome://<blah>/foo/bundle.js, this is |foo|.
# host_url: URL of the host. Usually something like "chrome://settings".
# excludes: Imports to exclude from the bundle.
# external_paths: Path mappings for import paths that are outside of
#                 |in_path|. For example:
#                 chrome://resources/|gen/ui/webui/resources/tsc
def _generate_rollup_config(out_dir, in_path, bundle_dir_path, host_url,
                            excludes, external_paths):
  rollup_config_file = os.path.join(out_dir, 'rollup.config.mjs')
  path_to_plugin = os.path.join(
      os.path.relpath(_HERE_PATH, out_dir), 'rollup_plugin.mjs')
  config_content = r'''
    import plugin from '{plugin_path}';
    export default ({{
      plugins: [
        plugin('{in_path}', '{bundle_dir_path}', '{host_url}', {exclude_list},
               {external_path_list}) ]
    }});
    '''.format(
      plugin_path=path_to_plugin.replace('\\', '/'),
      in_path=in_path.replace('\\', '/'),
      bundle_dir_path=bundle_dir_path.replace('\\', '/'),
      host_url=host_url,
      exclude_list=json.dumps(excludes),
      external_path_list=json.dumps(external_paths))
  with open(rollup_config_file, 'w', newline='', encoding='utf-8') as f:
    f.write(config_content)
  return rollup_config_file


# Create the manifest file from the sourcemap generated by rollup and return the
# list of bundles.
def _generate_manifest_file(out_dir, bundled_paths, bundle_dir_path,
                            manifest_out_path):
  manifest = {}
  for bundled_path in bundled_paths:
    sourcemap_file = bundled_path + '.map'
    with open(sourcemap_file, 'r', encoding='utf-8') as f:
      sourcemap = json.loads(f.read())
      if not 'sources' in sourcemap:
        raise Exception('rollup could not construct source map')
      sources = sourcemap['sources']
      replaced_sources = []
      # Normalize everything to be relative to the output directory. This is
      # where the conversion to a dependency file expects it to be.
      bundle_to_output = os.path.relpath(out_dir,
                                         os.path.join(out_dir, bundle_dir_path))
      for source in sources:
        if bundle_to_output != ".":
          replaced_sources.append(source.replace(bundle_to_output + "/", "", 1))
        else:
          replaced_sources.append(source)
      filepath = os.path.join(bundle_dir_path,
                              os.path.basename(bundled_path)).replace(
                                  '\\', '/')
      manifest[filepath] = replaced_sources

  with open(manifest_out_path, 'w', newline='', encoding='utf-8') as f:
    f.write(json.dumps(manifest))


def _bundle(out_folder, in_path, manifest_out_path, js_module_in_files,
            rollup_config_file):
  bundle_dir_path = os.path.dirname(js_module_in_files[0])
  out_dir = out_folder if not bundle_dir_path else os.path.join(
      out_folder, bundle_dir_path)
  if not os.path.exists(out_dir):
    os.makedirs(out_dir)

  rollup_args = [os.path.join(in_path, f) for f in js_module_in_files]

  # Confirm names are as expected. This is necessary to avoid having to replace
  # import statements in the generated output files.
  # TODO(rbpotter): Is it worth adding import statement replacement to support
  # arbitrary names?
  bundled_paths = []
  bundle_names = []

  assert len(js_module_in_files) < 3, '3+ input files not supported'

  for index, js_file in enumerate(js_module_in_files):
    bundle_name = '%s.rollup.js' % js_file[:-len('.js')]
    assert os.path.dirname(js_file) == bundle_dir_path, \
           'All input files must be in the same directory.'
    bundled_paths.append(os.path.join(out_folder, bundle_name))
    bundle_names.append(bundle_name)

  # This indicates that rollup is expected to generate a shared chunk file as
  # well as one file per module. Set its name using --chunkFileNames. Note:
  # Currently, this only supports 2 entry points, which generate 2 corresponding
  # outputs and 1 shared output.
  if (len(js_module_in_files) == 2):
    shared_file_name = 'shared.rollup.js'
    rollup_args += ['--chunkFileNames', shared_file_name]
    bundled_paths.append(os.path.join(out_dir, shared_file_name))
    bundle_names.append(os.path.join(bundle_dir_path, shared_file_name))

  node.RunNode([node_modules.PathToRollup()] + rollup_args + [
      '--format',
      'esm',
      '--dir',
      out_dir,
      '--entryFileNames',
      '[name].rollup.js',
      '--sourcemap',
      '--sourcemapExcludeSources',
      '--config',
      rollup_config_file,
  ])

  # Create the manifest file from the sourcemaps generated by rollup.
  _generate_manifest_file(out_folder, bundled_paths, bundle_dir_path,
                          manifest_out_path)

  for bundled_file in bundled_paths:
    with open(bundled_file, 'r', encoding='utf-8') as f:
      output = f.read()
      assert "<if expr" not in output, \
          'Unexpected <if expr> found in bundled output. Check that all ' + \
          'input files using such expressions are preprocessed.'

  return bundle_names


def _optimize(in_folder, args):
  in_path = os.path.normpath(os.path.join(_CWD, in_folder)).replace('\\', '/')
  out_path = os.path.join(_CWD, args.out_folder).replace('\\', '/')
  manifest_out_path = _request_list_path(out_path, args.target_name)

  excludes = [
      # This file is dynamically created by C++. Should always be imported with
      # a relative path.
      'strings.m.js',
  ]
  excludes.extend(args.exclude or [])

  for exclude in excludes:
    extension = os.path.splitext(exclude)[1]
    assert extension == '.js', f'Unexpected |excludes| entry: {exclude}.' + \
        ' Only .js files can appear in |excludes|.'

  external_paths = args.external_paths or []

  if args.rollup_config:
    # Use configuration provided from the caller.
    rollup_config_file = args.rollup_config
  else:
    # Use default configuration.
    bundle_dir_path = os.path.dirname(args.js_module_in_files[0])
    rollup_config_file = _generate_rollup_config(out_path, in_path,
                                                 bundle_dir_path, args.host_url,
                                                 excludes, external_paths)

  js_module_out_files = _bundle(out_path, in_path, manifest_out_path,
                                args.js_module_in_files, rollup_config_file)
  return {
      'manifest_out_path': manifest_out_path,
      'js_module_out_files': js_module_out_files,
  }


def main(argv):
  parser = argparse.ArgumentParser()
  parser.add_argument('--depfile', required=True)
  parser.add_argument('--target_name', required=True)
  parser.add_argument('--exclude', nargs='*')
  parser.add_argument('--external_paths', nargs='*')
  parser.add_argument('--host', required=True)
  parser.add_argument('--input', required=True)
  parser.add_argument('--out_folder', required=True)
  parser.add_argument('--js_module_in_files', nargs='*', required=True)
  parser.add_argument('--out-manifest')
  parser.add_argument('--rollup_config')
  args = parser.parse_args(argv)

  # NOTE(dbeam): on Windows, GN can send dirs/like/this. When joined, you might
  # get dirs/like/this\file.txt. This looks odd to windows. Normalize to right
  # the slashes.
  args.depfile = os.path.normpath(args.depfile)
  args.input = os.path.normpath(args.input)
  args.out_folder = os.path.normpath(args.out_folder)
  scheme_end_index = args.host.find('://')
  if (scheme_end_index == -1):
    args.host_url = 'chrome://%s/' % args.host
  else:
    args.host_url = args.host

  optimize_output = _optimize(args.input, args)

  # Prior call to _optimize() generated an output manifest file, containing
  # information about all files that were bundled. Grab it from there.
  with open(optimize_output['manifest_out_path'], 'r', encoding='utf-8') as f:
    manifest = json.loads(f.read())

    # Output a manifest file that will be used to auto-generate a grd file
    # later.
    if args.out_manifest:
      manifest_data = {
          'base_dir': args.out_folder.replace('\\', '/'),
          'files': list(manifest.keys()),
      }
      with open(
          os.path.normpath(os.path.join(_CWD, args.out_manifest)),
          'w',
          newline='',
          encoding='utf-8') as manifest_file:
        json.dump(manifest_data, manifest_file)

    dep_file_header = os.path.join(args.out_folder,
                                   optimize_output['js_module_out_files'][0])
    _update_dep_file(args.input, args, dep_file_header, manifest)


if __name__ == '__main__':
  main(sys.argv[1:])