# Copyright 2020 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.

import("//build_overrides/pigweed.gni")

import("$dir_pw_build/input_group.gni")
import("$dir_pw_build/python_action.gni")

# Python packages provide the following targets as $target_name.$subtarget.
_python_subtargets = [
  "tests",
  "lint",
  "lint.mypy",
  "lint.pylint",
  "install",
  "wheel",
]

# Defines a Python package. GN Python packages contain several GN targets:
#
#   - $name - Provides the Python files in the build, but does not take any
#         actions. All subtargets depend on this target.
#   - $name.lint - Runs static analyis tools on the Python code. This is a group
#     of two subtargets:
#     - $name.lint.mypy - Runs mypy (if enabled).
#     - $name.lint.pylint - Runs pylint (if enabled).
#   - $name.tests - Runs all tests for this package.
#   - $name.install - Installs the package in a venv.
#   - $name.wheel - Builds a Python wheel for the package. (Not implemented.)
#
# All Python packages are instantiated with the default toolchain, regardless of
# the current toolchain.
#
# TODO(pwbug/239): Implement wheel building.
#
# Args:
#   setup: List of setup file paths (setup.py or pyproject.toml & setup.cfg),
#       which must all be in the same directory.
#   sources: Python sources files in the package.
#   tests: Test files for this Python package.
#   python_deps: Dependencies on other pw_python_packages in the GN build.
#   python_test_deps: Test-only pw_python_package dependencies.
#   other_deps: Dependencies on GN targets that are not pw_python_packages.
#   inputs: Other files to track, such as package_data.
#   lint: If true (default), applies mypy and pylint to the package. If false,
#       does not.
#   pylintrc: Optional path to a pylintrc configuration file to use. If not
#       provided, Pylint's default rcfile search is used. Pylint is executed
#       from the package's setup directory, so pylintrc files in that directory
#       will take precedence over others.
#   mypy_ini: Optional path to a mypy configuration file to use. If not
#       provided, mypy's default configuration file search is used. mypy is
#       executed from the package's setup directory, so mypy.ini files in that
#       directory will take precedence over others.
template("pw_python_package") {
  if (defined(invoker.sources)) {
    _all_py_files = invoker.sources
  } else {
    _all_py_files = []
  }

  if (defined(invoker.tests)) {
    _test_sources = invoker.tests
  } else {
    _test_sources = []
  }

  # The Python targets are always instantiated in the default toolchain. Use
  # fully qualified labels so that the toolchain is not lost.
  _other_deps = []
  if (defined(invoker.other_deps)) {
    foreach(dep, invoker.other_deps) {
      _other_deps += [ get_label_info(dep, "label_with_toolchain") ]
    }
  }

  _all_py_files += _test_sources

  # pw_python_script uses pw_python_package, but with a limited set of features.
  # _pw_standalone signals that this target is actually a pw_python_script.
  _is_package = !(defined(invoker._pw_standalone) && invoker._pw_standalone)

  # Some build targets generate Python packages, setting _pw_generated to
  # indicate this.
  _is_generated_package =
      defined(invoker._pw_generated) && invoker._pw_generated

  # Argument: invoker.lint = [true | false]; default = true.
  # Default to false for generated packages, but allow overrides.
  if (defined(invoker.lint)) {
    _should_lint = invoker.lint
  } else {
    _should_lint = !_is_generated_package
  }

  if (_is_package) {
    assert(defined(invoker.setup) && invoker.setup != [],
           "pw_python_package requires 'setup' to point to a setup.py file " +
               "or pyproject.toml and setup.cfg files")

    if (!_is_generated_package) {
      _all_py_files += invoker.setup
    }

    # Get the directories of the setup files. All files must be in the same dir.
    _setup_dirs = get_path_info(invoker.setup, "dir")
    _setup_dir = _setup_dirs[0]

    foreach(dir, _setup_dirs) {
      assert(dir == _setup_dir,
             "All files in 'setup' must be in the same directory")
    }

    # If sources are provided, make sure there is an __init__.py file.
    if (!_is_generated_package && defined(invoker.sources) &&
        invoker.sources != []) {
      assert(filter_include(invoker.sources, [ "*\b__init__.py" ]) != [],
             "Python packages must have at least one __init__.py file")
    }
  }

  _python_deps = []
  if (defined(invoker.python_deps)) {
    foreach(dep, invoker.python_deps) {
      # Use the fully qualified name so the subtarget can be appended as needed.
      _python_deps += [ get_label_info(dep, "label_no_toolchain") ]
    }
  }

  # All dependencies needed for the package and its tests.
  _python_test_deps = _python_deps
  if (defined(invoker.python_test_deps)) {
    foreach(test_dep, invoker.python_test_deps) {
      _python_test_deps += [ get_label_info(test_dep, "label_no_toolchain") ]
    }
  }

  if (_test_sources == []) {
    assert(!defined(invoker.python_test_deps),
           "python_test_deps was provided, but there are no tests in " +
               get_label_info(":$target_name", "label_no_toolchain"))
    not_needed(_python_test_deps)
  }

  _internal_target = "$target_name._internal"

  # Create groups with the public target names ($target_name, $target_name.lint,
  # $target_name.install, etc.). These are actually wrappers around internal
  # Python actions instantiated with the default toolchain. This ensures there
  # is only a single copy of each Python action in the build.
  #
  # The $target_name.tests group is created separately below.
  foreach(subtarget, _python_subtargets - [ "tests" ]) {
    group("$target_name.$subtarget") {
      deps = [ ":$_internal_target.$subtarget($default_toolchain)" ]
    }
  }

  group("$target_name") {
    deps = [ ":$_internal_target($default_toolchain)" ]
  }

  # Declare the main Python package group. This represents the Python files, but
  # does not take any actions. GN targets can depend on the package name to run
  # when any files in the package change.
  pw_input_group("$_internal_target") {
    inputs = _all_py_files
    if (defined(invoker.inputs)) {
      inputs += invoker.inputs
    }

    deps = _python_deps + _other_deps
  }

  if (_is_package) {
    # Install this Python package and its dependencies in the current Python
    # environment.
    pw_python_action("$_internal_target.install") {
      module = "pip"
      args = [ "install" ]

      # Don't install generated packages with --editable, since the build
      # directory is ephemeral.
      if (!_is_generated_package) {
        args += [ "--editable" ]
      }
      args += [ rebase_path(_setup_dir) ]

      stamp = true

      # Parallel pip installations don't work, so serialize pip invocations.
      pool = "$dir_pw_build:pip_pool"

      deps = [ ":$_internal_target" ]
      foreach(dep, _python_deps) {
        _subtarget = get_label_info(dep, "label_no_toolchain") + ".install"
        deps += [ "$_subtarget(" + get_label_info(dep, "toolchain") + ")" ]
      }
    }

    # TODO(pwbug/239): Add support for building groups of wheels. The code below
    #     is incomplete and untested.
    pw_python_action("$_internal_target.wheel") {
      script = "$dir_pw_build/py/pw_build/python_wheels.py"

      args = [
        "--out_dir",
        rebase_path(target_out_dir),
      ]
      args += rebase_path(_all_py_files)

      deps = [ ":$_internal_target.install" ]
      stamp = true
    }
  } else {
    # If this is not a package, install or build wheels for its deps only.
    group("$_internal_target.install") {
      deps = []
      foreach(dep, _python_deps) {
        deps += [ "$dep.install" ]
      }
    }
    group("$_internal_target.wheel") {
      deps = []
      foreach(dep, _python_deps) {
        deps += [ "$dep.wheel" ]
      }
    }
  }

  # Define the static analysis targets for this package.
  group("$_internal_target.lint") {
    deps = [
      ":$_internal_target.lint.mypy",
      ":$_internal_target.lint.pylint",
    ]
  }

  if (_should_lint || _test_sources != []) {
    # Packages that must be installed to use the package or run its tests.
    _test_install_deps = [ ":$_internal_target.install" ]
    foreach(dep, _python_test_deps + [ "$dir_pw_build:python_lint" ]) {
      _test_install_deps += [ "$dep.install" ]
    }
  }

  # For packages that are not generated, create targets to run mypy and pylint.
  # Linting is not performed on generated packages.
  if (_should_lint) {
    # Run lint tools from the setup or target directory so that the tools detect
    # config files (e.g. pylintrc or mypy.ini) in that directory. Config files
    # may be explicitly specified with the pylintrc or mypy_ini arguments.
    if (defined(_setup_dir)) {
      _lint_directory = rebase_path(_setup_dir)
    } else {
      _lint_directory = rebase_path(".")
    }

    pw_python_action("$_internal_target.lint.mypy") {
      module = "mypy"
      args = [
        "--pretty",
        "--show-error-codes",
      ]

      if (defined(invoker.mypy_ini)) {
        args += [ "--config-file=" + rebase_path(invoker.mypy_ini) ]
        inputs = [ invoker.mypy_ini ]
      }

      args += rebase_path(_all_py_files)

      # Use this environment variable to force mypy to colorize output.
      # See https://github.com/python/mypy/issues/7771
      environment = [ "MYPY_FORCE_COLOR=1" ]

      directory = _lint_directory
      stamp = true

      deps = _test_install_deps

      foreach(dep, _python_test_deps) {
        deps += [ "$dep.lint.mypy" ]
      }
    }

    # Create a target to run pylint on each of the Python files in this
    # package and its dependencies.
    pw_python_action_foreach("$_internal_target.lint.pylint") {
      module = "pylint"
      args = [
        rebase_path(".") + "/{{source_target_relative}}",
        "--jobs=1",
        "--output-format=colorized",
      ]

      if (defined(invoker.pylintrc)) {
        args += [ "--rcfile=" + rebase_path(invoker.pylintrc) ]
        inputs = [ invoker.pylintrc ]
      }

      if (host_os == "win") {
        # Allow CRLF on Windows, in case Git is set to switch line endings.
        args += [ "--disable=unexpected-line-ending-format" ]
      }

      sources = _all_py_files

      directory = _lint_directory
      stamp = "$target_gen_dir/{{source_target_relative}}.pylint.passed"

      deps = _test_install_deps

      foreach(dep, _python_test_deps) {
        deps += [ "$dep.lint.pylint" ]
      }
    }
  } else {
    pw_input_group("$_internal_target.lint.mypy") {
      if (defined(invoker.pylintrc)) {
        inputs = [ invoker.pylintrc ]
      }
    }
    pw_input_group("$_internal_target.lint.pylint") {
      if (defined(invoker.mypy_ini)) {
        inputs = [ invoker.mypy_ini ]
      }
    }
  }

  # Create a target for each test file.
  _test_targets = []

  foreach(test, _test_sources) {
    _test_name = string_replace(test, "/", "_")
    _internal_test_target = "$_internal_target.tests.$_test_name"

    pw_python_action(_internal_test_target) {
      script = test
      stamp = true

      deps = _test_install_deps

      foreach(dep, _python_test_deps) {
        deps += [ "$dep.tests" ]
      }
    }

    # Create a public version of each test target, so tests can be executed as
    # //path/to:package.tests.foo.py.
    group("$target_name.tests.$_test_name") {
      deps = [ ":$_internal_test_target" ]
    }

    _test_targets += [ ":$target_name.tests.$_test_name" ]
  }

  group("$target_name.tests") {
    deps = _test_targets
  }
}

# Declares a group of Python packages or other Python groups. pw_python_groups
# expose the same set of subtargets as pw_python_package (e.g.
# "$group_name.lint" and "$group_name.tests"), but these apply to all packages
# in deps and their dependencies.
template("pw_python_group") {
  if (defined(invoker.python_deps)) {
    _python_deps = invoker.python_deps
  } else {
    _python_deps = []
  }

  group(target_name) {
    deps = _python_deps
  }

  foreach(subtarget, _python_subtargets) {
    group("$target_name.$subtarget") {
      deps = []
      foreach(dep, _python_deps) {
        # Split out the toolchain to support deps with a toolchain specified.
        _target = get_label_info(dep, "label_no_toolchain")
        _toolchain = get_label_info(dep, "toolchain")
        deps += [ "$_target.$subtarget($_toolchain)" ]
      }
    }
  }
}

# Declares Python scripts or tests that are not part of a Python package.
# Similar to pw_python_package, but only supports a subset of its features.
#
# pw_python_script accepts the same arguments as pw_python_package, except
# `setup` cannot be provided.
#
# pw_python_script provides the same subtargets as pw_python_package, but
# $target_name.install and $target_name.wheel only affect the python_deps of
# this GN target, not the target itself.
template("pw_python_script") {
  _supported_variables = [
    "sources",
    "tests",
    "python_deps",
    "other_deps",
    "inputs",
    "pylintrc",
  ]

  pw_python_package(target_name) {
    _pw_standalone = true
    forward_variables_from(invoker, _supported_variables)
  }
}

# Represents a list of Python requirements, as in a requirements.txt.
#
# Args:
#  files: One or more requirements.txt files.
#  requirements: A list of requirements.txt-style requirements.
template("pw_python_requirements") {
  assert(defined(invoker.files) || defined(invoker.requirements),
         "pw_python_requirements requires a list of requirements.txt files " +
             "in the 'files' arg or requirements in 'requirements'")

  _requirements_files = []

  if (defined(invoker.files)) {
    _requirements_files += invoker.files
  }

  if (defined(invoker.requirements)) {
    _requirements_file = "$target_gen_dir/$target_name.requirements.txt"
    write_file(_requirements_file, invoker.requirements)
    _requirements_files += [ _requirements_file ]
  }

  # The default target represents the requirements themselves.
  pw_input_group(target_name) {
    inputs = _requirements_files
  }

  # Use the same subtargets as pw_python_package so these targets can be listed
  # as python_deps of pw_python_packages.
  pw_python_action("$target_name.install") {
    inputs = _requirements_files

    module = "pip"
    args = [ "install" ]

    foreach(_requirements_file, inputs) {
      args += [
        "--requirement",
        rebase_path(_requirements_file),
      ]
    }

    pool = "$dir_pw_build:pip_pool"
    stamp = true
  }

  # Create stubs for the unused subtargets so that pw_python_requirements can be
  # used as python_deps.
  foreach(subtarget, _python_subtargets - [ "install" ]) {
    group("$target_name.$subtarget") {
    }
  }
}