#!/usr/bin/env python3
# Copyright 2021 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Builds the Rust toolchain used to compile and run TockOS."""

# pylint: disable=cros-logging-import

import argparse
import logging
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import tempfile
from typing import Iterable, List, NamedTuple, Optional

# The target triples that we want this toolchain to be able to build binaries
# for. ATM this consists of the host and one target. If you update this,
# remember to update |write_config_toml| with info about the compilers to use.
#
# _A_ rust compiler is needed for the host because some packages have build
# scripts (e.g., build.rs), and Ti50 has previously expressed interest in
# running _some_ tests in a test environment that runs natively on the host.
# It's unclear how difficult it'd be to use the CrOS Rust toolchain for both of
# these, given that said toolchain is locked to the stable channel, and is
# likely to be a very different version of Rust than Ti50's toolchain.
TARGET_TRIPLES_TO_SHIP = (
    'riscv32imc-unknown-none-elf',
    'x86_64-unknown-linux-gnu',
)


class RustComponent(NamedTuple):
  """A component of Rust that we want to build and install."""
  # The name of the component that |./x.py dist| understands.
  dist_build_name: str
  # A regex pattern that matches the tarball produced for this component under
  # dist/.
  dist_tarball_pattern: str
  # Whether the component's dist tarballs are per-triple.
  dist_tarballs_are_per_triple: bool


# Components we ship as a part of our toolchain.
COMPONENTS_TO_SHIP = (
    # cargo, which is shipped with |rustup| toolchains by default.
    RustComponent(
        dist_build_name='cargo',
        dist_tarball_pattern=r'^cargo-',
        dist_tarballs_are_per_triple=False,
    ),
    # clippy, which was requested by TockOS.
    RustComponent(
        dist_build_name='clippy',
        dist_tarball_pattern=r'^clippy-',
        dist_tarballs_are_per_triple=False,
    ),
    # llvm-tools-preview, which is apparently required by TockOS itself.
    RustComponent(
        dist_build_name='llvm-tools',
        dist_tarball_pattern=r'^llvm-tools-',
        dist_tarballs_are_per_triple=False,
    ),
    # src, which is required when cross-compiling in some contexts.
    RustComponent(
        dist_build_name='src',
        dist_tarball_pattern=r'^rust-src-',
        dist_tarballs_are_per_triple=False,
    ),
    # src/librustc, which is the actual |rustc| compiler.
    RustComponent(
        dist_build_name='src/librustc',
        dist_tarball_pattern=r'^rustc-[^-]+-dev-(?!src)',
        dist_tarballs_are_per_triple=False,
    ),
    # rustfmt, the standard formatter for rust code.
    RustComponent(
        dist_build_name='rustfmt',
        dist_tarball_pattern=r'^rustfmt-',
        dist_tarballs_are_per_triple=False,
    ),
    # library/std, which consists of stdlib |rlib|s.
    RustComponent(
        dist_build_name='library/std',
        dist_tarball_pattern=r'^rust-std-',
        dist_tarballs_are_per_triple=True,
    ),
)


def must_getenv(key: str) -> str:
  """Gets the value of an environment variable; raises if unset or empty."""
  x = os.getenv(key)
  if not x:
    raise ValueError(f'No value found for {key}; this env var is required')
  return x


def write_config_toml(rust_src: Path, rv_clang_bin: Path, install_prefix: Path,
                      rustc: Optional[Path], cargo: Optional[Path],
                      rustfmt: Optional[Path]) -> None:
  """Writes Rust's config.toml with the given parameters."""
  # We need a Pythony list, but formatted with double-quotes.
  assert not any('"' in x or '\\' in x
                 for x in TARGET_TRIPLES_TO_SHIP), TARGET_TRIPLES_TO_SHIP
  target_triples_str = ', '.join(f'"{x}"' for x in TARGET_TRIPLES_TO_SHIP)

  component_locations = (
      ('rustc', rustc),
      ('cargo', cargo),
      ('rustfmt', rustfmt),
  )
  optional_bootstrap_configuration = '\n'.join(
      f'{key} = "{exe}"' for key, exe in component_locations if exe is not None)

  toml_contents = f"""
[build]
target = [{target_triples_str}]
extended = true
vendor = true
python = "{must_getenv('EPYTHON')}"
submodules = false
{optional_bootstrap_configuration}

[llvm]
use-libcxx = true
ninja = true

[install]
prefix = "{install_prefix}"

[rust]
# Provide llvm-tools like llvm-objcopy for the llvm-tools package.
llvm-tools = true
codegen-units = 0
codegen-tests = false

lld = true
default-linker = "{must_getenv('CBUILD')}-clang"
use-lld = false

[target.riscv32imc-unknown-none-elf]
cc = "{rv_clang_bin}/clang"
cxx = "{rv_clang_bin}/clang++"
linker = "{rv_clang_bin}/clang++"

[target.x86_64-unknown-linux-gnu]
cc = "x86_64-pc-linux-gnu-clang"
cxx = "x86_64-pc-linux-gnu-clang++"
linker = "x86_64-pc-linux-gnu-clang++"
"""

  (rust_src / 'config.toml').write_text(toml_contents, encoding='utf-8')


def compile_rust_src(rust_src: Path) -> None:
  """Compiles all targets specified by COMPONENTS_TO_SHIP."""
  targets = [x.dist_build_name for x in COMPONENTS_TO_SHIP]
  subprocess.check_call(['./x.py', 'dist'] + targets, cwd=rust_src)


def copy_tree(input_path: Path, output_path: Path,
              ignore_file_names: Iterable[str]):
  """shutil.copytree with dirs_exist_ok=True.

  dirs_exist_ok=True is only available in py3.8+.
  """
  output_path.mkdir(parents=True, exist_ok=True)
  for root, _, files in os.walk(input_path):
    root = Path(root).relative_to(input_path)
    input_root = input_path / root
    output_root = output_path / root
    output_root.mkdir(exist_ok=True)
    for f in files:
      if f in ignore_file_names:
        continue

      output_file = output_root / f
      if output_file.exists():
        raise ValueError(f'File at {output_file} already exists; refusing to '
                         'overwrite')
      logging.debug('Copying %s => %s', input_root / f, output_root / f)
      shutil.copy2(input_root / f, output_root / f)


def package_rust_src(rust_src: Path, install_dir: Path):
  """Installs all components built in |rust_src| into |install_dir|."""
  dist_path = rust_src / 'build' / 'dist'
  all_tarballs = [x for x in dist_path.iterdir() if x.name.endswith('.tar.gz')]
  temp_dir = Path(tempfile.mkdtemp(prefix='build_rust_package'))

  def install_component(tarball: Path) -> None:
    logging.info('Untarring %s into %s', tarball, temp_dir)

    # All tarballs unpack to a single directory, which is the tarball without
    # the |.tar.${comp}| suffix.
    subprocess.check_call(['tar', 'xaf', str(tarball)], cwd=temp_dir)

    subdir_name = tarball.stem
    tar_suffix = '.tar'
    assert tarball.stem.endswith(tar_suffix), subdir_name
    subdir_name = subdir_name[:-len(tar_suffix)]
    base_dir = temp_dir / subdir_name

    # ...And this directory itself should contain a |components| file
    # describing what components this tarball has.
    components = [
        x.strip()
        for x in (base_dir /
                  'components').read_text(encoding='utf-8').splitlines()
    ]

    for component in components:
      input_path = base_dir / component
      logging.info('Installing component at %s into %s', input_path,
                   install_dir)
      copy_tree(
          input_path,
          install_dir,
          ignore_file_names=('manifest.in',),
      )

      # All inputs have a manifest.in at their root, which describes the files
      # installed by the given component. |rustup| likes to include those in
      # the final installation directory, and that seems like a generally okay
      # idea. These are copied _purely_ for informational purposes, rather than
      # serving some functional goal.
      manifest_dir = install_dir / 'lib' / 'rustlib'
      manifest_dir.mkdir(parents=True, exist_ok=True)
      manifest_name = f'manifest-{input_path.name}'
      shutil.copyfile(input_path / 'manifest.in', manifest_dir / manifest_name)

  try:
    install_dir.mkdir(parents=True, exist_ok=True)

    for component in COMPONENTS_TO_SHIP:
      regex = re.compile(component.dist_tarball_pattern)
      targets = [x for x in all_tarballs if regex.search(x.name)]
      if component.dist_tarballs_are_per_triple:
        for target_triple in TARGET_TRIPLES_TO_SHIP:
          ts = [x for x in targets if target_triple in x.name]
          if len(ts) != 1:
            raise ValueError(f'Expected exactly one match for {component} with '
                             f'triple {target_triple}; got {ts}')
          install_component(ts[0])
      else:
        if len(targets) != 1:
          raise ValueError(f'Expected exactly one match for {component}; '
                           f'got {targets}')
        install_component(targets[0])
  except:
    logging.error('Packaging failed; leaving tempdir around at %s', temp_dir)
    raise
  else:
    logging.info('Packaging succeeded; cleaning up tempdir. Results are in %s',
                 install_dir)
    shutil.rmtree(temp_dir)


def get_parser():
  """Creates a parser for commandline args."""
  parser = argparse.ArgumentParser(
      description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
  parser.add_argument(
      '--rust-src',
      required=True,
      type=Path,
      help="Path to rust's source code",
  )
  parser.add_argument(
      '--install-dir',
      required=True,
      type=Path,
      help='Path to the directory in which we should put installation '
      'artifacts',
  )
  parser.add_argument(
      '--install-prefix',
      required=True,
      type=Path,
      help='Prefix under which installed artifacts will sit on the system',
  )
  parser.add_argument(
      '--rv-clang-bin',
      required=True,
      type=Path,
      help='Path the bin/ directory of a riscv-enabled clang',
  )
  parser.add_argument(
      '--rustc',
      type=Path,
      help='Optional path to the rustc we should use to bootstrap the '
      'compiler',
  )
  parser.add_argument(
      '--cargo',
      type=Path,
      help='Optional path to the cargo we should use to bootstrap the '
      'compiler',
  )
  parser.add_argument(
      '--rustfmt',
      type=Path,
      help='Optional path to the rustfmt we should use to bootstrap the '
      'compiler',
  )
  return parser


def main(argv: List[str]):
  logging.basicConfig(
      format='%(asctime)s: %(levelname)s: %(filename)s:%(lineno)d: %(message)s',
      level=logging.INFO,
  )

  parser = get_parser()
  opts = parser.parse_args(argv)

  install_dir = opts.install_dir.resolve()
  rust_src = opts.rust_src.resolve()
  rv_clang_bin = opts.rv_clang_bin.resolve()
  # This might not exist yet, so we don't want to resolve() it.
  install_prefix = opts.install_prefix
  if not install_prefix.is_absolute():
    parser.error('--install_prefix should be an absolute path')

  cargo = Path(opts.cargo).resolve() if opts.cargo else None
  rustc = Path(opts.rustc).resolve() if opts.rustc else None
  rustfmt = Path(opts.rustfmt).resolve() if opts.rustfmt else None

  write_config_toml(rust_src, rv_clang_bin, install_prefix, rustc, cargo,
                    rustfmt)
  compile_rust_src(rust_src)
  package_rust_src(rust_src, install_dir)


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