# Copyright 2022 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/python.gni")
import("$dir_pw_build/python_action.gni")

# Defines and creates a Python virtualenv. This template is used by Pigweed in
# https://cs.pigweed.dev/pigweed/+/main:pw_env_setup/BUILD.gn to
# create a virtualenv for use within the GN build that all Python actions will
# run in.
#
# Example:
#
#   pw_python_venv("test_venv") {
#     path = "test-venv"
#     constraints = [ "//tools/constraints.list" ]
#     requirements = [ "//tools/requirements.txt" ]
#     source_packages = [
#       "$dir_pw_cli/py",
#       "$dir_pw_console/py",
#       "//tools:another_pw_python_package",
#     ]
#   }
#
# Args:
#   path: The directory where the virtualenv will be created. This is relative
#     to the GN root and must begin with "$root_build_dir/" if it lives in the
#     output directory or "//" if it lives in elsewhere.
#
#   constraints: A list of constraint files used when performing pip install
#     into this virtualenv. By default this is set to pw_build_PIP_CONSTRAINTS
#
#   requirements: A list of requirements files to install into this virtualenv
#     on creation. By default this is set to pw_build_PIP_REQUIREMENTS
#
#   pip_generate_hashes: (Default: false) Use --generate-hashes When
#     running pip-compile to compute the final requirements.txt
#
#   source_packages: A list of in-tree pw_python_package targets that will be
#     checked for external third_party pip dependencies to install into this
#     virtualenv. Note this list of targets isn't actually installed into the
#     virtualenv. Only packages defined inside the [options] install_requires
#     section of each pw_python_package's setup.cfg will be pip installed. See
#     this page for a setup.cfg example:
#     https://setuptools.pypa.io/en/latest/userguide/declarative_config.html
#
#   output_logs: (Default: true) Commands will output logs.
#
template("pw_python_venv") {
  assert(defined(invoker.path), "pw_python_venv requires a 'path'")

  _path = invoker.path

  _generated_requirements_file =
      "$target_gen_dir/$target_name/generated_requirements.txt"

  _compiled_requirements_file =
      "$target_gen_dir/$target_name/compiled_requirements.txt"

  _source_packages = []
  if (defined(invoker.source_packages)) {
    _source_packages += invoker.source_packages
  } else {
    not_needed([
                 "_source_packages",
                 "_generated_requirements_file",
               ])
  }
  _output_logs = true
  if (defined(invoker.output_logs)) {
    _output_logs = invoker.output_logs
  }
  if (!defined(invoker.output_logs) ||
      current_toolchain != pw_build_PYTHON_TOOLCHAIN) {
    not_needed([ "_output_logs" ])
  }

  _source_package_labels = []
  foreach(pkg, _source_packages) {
    _source_package_labels += [ get_label_info(pkg, "label_no_toolchain") ]
  }

  if (defined(invoker.requirements)) {
    _requirements = invoker.requirements
  } else {
    _requirements = pw_build_PIP_REQUIREMENTS
  }

  if (defined(invoker.constraints)) {
    _constraints = invoker.constraints
  } else {
    _constraints = pw_build_PIP_CONSTRAINTS
  }

  _python_interpreter = _path + "/bin/python"
  if (host_os == "win") {
    _python_interpreter = _path + "/Scripts/python.exe"
  }

  _venv_metadata_json_file = "$target_gen_dir/$target_name/venv_metadata.json"
  _venv_metadata = {
    gn_target_name =
        get_label_info(":${invoker.target_name}", "label_no_toolchain")
    path = rebase_path(_path, root_build_dir)
    generated_requirements =
        rebase_path(_generated_requirements_file, root_build_dir)
    compiled_requirements =
        rebase_path(_compiled_requirements_file, root_build_dir)
    requirements = rebase_path(_requirements, root_build_dir)
    constraints = rebase_path(_constraints, root_build_dir)
    interpreter = rebase_path(_python_interpreter, root_build_dir)
    source_packages = _source_package_labels
  }
  write_file(_venv_metadata_json_file, _venv_metadata, "json")

  pw_python_action("${target_name}._create_virtualenv") {
    _pw_internal_run_in_venv = false

    # Note: The if the venv isn't in the out dir then we can't declare
    # outputs and must stamp instead.
    stamp = true

    # The virtualenv should depend on the version of Python currently in use.
    stampfile = "$target_gen_dir/$target_name.pw_pystamp"
    depfile = "$target_gen_dir/$target_name.d"
    script = "$dir_pw_build/py/pw_build/create_gn_venv.py"
    args = [
      "--depfile",
      rebase_path(depfile, root_build_dir),
      "--destination-dir",
      rebase_path(_path, root_build_dir),
      "--stampfile",
      rebase_path(stampfile, root_build_dir),
    ]
  }

  if (defined(invoker.source_packages) &&
      current_toolchain == pw_build_PYTHON_TOOLCHAIN) {
    pw_python_action("${target_name}._generate_3p_requirements") {
      inputs = _requirements + _constraints

      _pw_internal_run_in_venv = false
      _forward_python_metadata_deps = true

      script = "$dir_pw_build/py/pw_build/generate_python_requirements.py"

      _pkg_gn_labels = []
      foreach(pkg, _source_packages) {
        _pkg_gn_labels += [ get_label_info(pkg, "label_no_toolchain") +
                            "($pw_build_PYTHON_TOOLCHAIN)" ]
      }

      # Add target packages to the python_metadata_deps. This will let
      # GN expand the transitive pw_python_package deps which are read
      # by generate_python_requirements.py
      python_metadata_deps = _pkg_gn_labels

      args = [
        "--gn-root-build-dir",
        rebase_path(root_build_dir, root_build_dir),
        "--output-requirement-file",
        rebase_path(_generated_requirements_file, root_build_dir),
      ]

      if (_constraints != []) {
        args += [ "--constraint-files" ]
      }
      foreach(_constraints_file, _constraints) {
        args += [ rebase_path(_constraints_file, root_build_dir) ]
      }

      args += [
        "--gn-packages",
        string_join(",", _pkg_gn_labels),
      ]

      outputs = [ _generated_requirements_file ]
      deps = [ ":${invoker.target_name}._create_virtualenv($pw_build_PYTHON_TOOLCHAIN)" ]
    }
  } else {
    group("${target_name}._generate_3p_requirements") {
    }
  }

  _pip_generate_hashes = false
  if (defined(invoker.pip_generate_hashes)) {
    _pip_generate_hashes = invoker.pip_generate_hashes
  } else {
    not_needed([ "_pip_generate_hashes" ])
  }

  if (defined(invoker.source_packages) || defined(invoker.requirements)) {
    if (current_toolchain == pw_build_PYTHON_TOOLCHAIN) {
      # Compile requirements with hashes
      pw_python_action("${target_name}._compile_requirements") {
        module = "piptools"

        # Set the venv to run this pip install in.
        _pw_internal_run_in_venv = true
        _skip_installing_external_python_deps = true
        venv = get_label_info(":${invoker.target_name}", "label_no_toolchain")

        _pip_args = []
        if (pw_build_PYTHON_PIP_INSTALL_OFFLINE) {
          _pip_args += [ "--no-index" ]
        }
        if (pw_build_PYTHON_PIP_INSTALL_DISABLE_CACHE) {
          _pip_args += [ "--no-cache-dir" ]
        }
        if (pw_build_PYTHON_PIP_INSTALL_FIND_LINKS != []) {
          foreach(link_dir, pw_build_PYTHON_PIP_INSTALL_FIND_LINKS) {
            _pip_args +=
                [ "--find-links=" + rebase_path(link_dir, root_build_dir) ]
          }
        }

        args = [ "compile" ]

        if (_pip_generate_hashes) {
          args += [
            "--generate-hashes",
            "--reuse-hashes",
          ]
        }

        args += [
          "--resolver=backtracking",

          # --allow-unsafe will force pinning pip and setuptools.
          "--allow-unsafe",
          "--output-file",
          rebase_path(_compiled_requirements_file, root_build_dir),

          # Input requirements file:
          rebase_path(_generated_requirements_file, root_build_dir),
        ]

        # Pass offline related pip args through the pip-compile command.
        if (_pip_args != []) {
          args += [
            "--pip-args",
            string_join(" ", _pip_args),
          ]
        }

        # Extra requirements files
        foreach(_requirements_file, _requirements) {
          args += [ rebase_path(_requirements_file, root_build_dir) ]
        }

        inputs = []

        # NOTE: constraint files are included in the content of the
        # _generated_requirements_file. This occurs in the
        # ._generate_3p_requirements target.
        inputs += _constraints
        inputs += _requirements
        inputs += [ _generated_requirements_file ]

        deps = [
          ":${invoker.target_name}._generate_3p_requirements($pw_build_PYTHON_TOOLCHAIN)",
          ":${invoker.target_name}._install_base_3p_deps($pw_build_PYTHON_TOOLCHAIN)",
        ]
        outputs = [ _compiled_requirements_file ]
      }

      # This target will run 'pip install' in the venv to provide base
      # dependencies needed to run piptools commands. That is required for the
      # _compile_requirements sub target.
      pw_python_action("${target_name}._install_base_3p_deps") {
        module = "pip"

        # Set the venv to run this pip install in.
        _pw_internal_run_in_venv = true
        _skip_installing_external_python_deps = true
        venv = get_label_info(":${invoker.target_name}", "label_no_toolchain")

        _base_requirement_file = "$dir_pw_env_setup/py/pw_env_setup/virtualenv_setup/python_base_requirements.txt"

        args = [
          "install",
          "--requirement",
          rebase_path(_base_requirement_file, root_build_dir),
        ]
        if (_output_logs) {
          _pip_install_log_file =
              "$target_gen_dir/$target_name/pip_install_log.txt"
          args += [
            "--log",
            rebase_path(_pip_install_log_file, root_build_dir),
          ]
          outputs = [ _pip_install_log_file ]
        }

        # NOTE: Constraints should be ignored for this step.

        if (pw_build_PYTHON_PIP_INSTALL_OFFLINE) {
          args += [ "--no-index" ]
        }
        if (pw_build_PYTHON_PIP_INSTALL_DISABLE_CACHE) {
          args += [ "--no-cache-dir" ]
        }
        if (pw_build_PYTHON_PIP_INSTALL_FIND_LINKS != []) {
          foreach(link_dir, pw_build_PYTHON_PIP_INSTALL_FIND_LINKS) {
            args += [
              "--find-links",
              rebase_path(link_dir, root_build_dir),
            ]
          }
        }

        deps = [ ":${invoker.target_name}._create_virtualenv($pw_build_PYTHON_TOOLCHAIN)" ]
        stamp = true
        pool = "$dir_pw_build/pool:pip($default_toolchain)"
      }

      # Install all 3rd party Python dependencies.
      pw_python_action("${target_name}._install_3p_deps") {
        module = "pip"

        # Set the venv to run this pip install in.
        _pw_internal_run_in_venv = true
        _skip_installing_external_python_deps = true
        venv = get_label_info(":${invoker.target_name}", "label_no_toolchain")

        args = pw_build_PYTHON_PIP_DEFAULT_OPTIONS
        args += [
          "install",
          "--upgrade",
        ]

        if (_output_logs) {
          _pip_install_log_file =
              "$target_gen_dir/$target_name/pip_install_log.txt"
          args += [
            "--log",
            rebase_path(_pip_install_log_file, root_build_dir),
          ]
        }

        if (_pip_generate_hashes) {
          args += [ "--require-hashes" ]
        }

        if (pw_build_PYTHON_PIP_INSTALL_OFFLINE) {
          args += [ "--no-index" ]
        }
        if (pw_build_PYTHON_PIP_INSTALL_DISABLE_CACHE) {
          args += [ "--no-cache-dir" ]
        }
        if (pw_build_PYTHON_PIP_INSTALL_FIND_LINKS != []) {
          foreach(link_dir, pw_build_PYTHON_PIP_INSTALL_FIND_LINKS) {
            args += [
              "--find-links",
              rebase_path(link_dir, root_build_dir),
            ]
          }
        }

        # Note: --no-build-isolation should be avoided for installing 3rd party
        # Python packages that use C/C++ extension modules.
        # https://setuptools.pypa.io/en/latest/userguide/ext_modules.html
        inputs = _constraints + _requirements + [ _compiled_requirements_file ]

        # Use the pip-tools compiled requiremets file. This contains the fully
        # expanded list of deps with constraints applied.
        if (defined(invoker.source_packages)) {
          inputs += [ _compiled_requirements_file ]
          args += [
            "--requirement",
            rebase_path(_compiled_requirements_file, root_build_dir),
          ]
        }

        deps = [
          ":${invoker.target_name}._compile_requirements($pw_build_PYTHON_TOOLCHAIN)",
          ":${invoker.target_name}._generate_3p_requirements($pw_build_PYTHON_TOOLCHAIN)",
          ":${invoker.target_name}._install_base_3p_deps($pw_build_PYTHON_TOOLCHAIN)",
        ]
        stamp = true
        pool = "$dir_pw_build/pool:pip($default_toolchain)"
      }

      # Target to create a Python package cache for this pw_python_venv.
      pw_python_action("${target_name}.vendor_wheels") {
        _wheel_output_dir = "$target_gen_dir/$target_name/wheels"
        _pip_download_logfile =
            "$target_gen_dir/$target_name/pip_download_log.txt"
        _pip_wheel_logfile = "$target_gen_dir/$target_name/pip_wheel_log.txt"
        metadata = {
          pw_python_package_wheels = [ _wheel_output_dir ]
        }

        script = "$dir_pw_build/py/pw_build/generate_python_wheel_cache.py"

        # Set the venv to run this pip install in.
        _pw_internal_run_in_venv = true
        _skip_installing_external_python_deps = true
        venv = get_label_info(":${invoker.target_name}", "label_no_toolchain")

        args = [
          "--pip-download-log",
          rebase_path(_pip_download_logfile, root_build_dir),
          "--wheel-dir",
          rebase_path(_wheel_output_dir, root_build_dir),
          "-r",
          rebase_path(_compiled_requirements_file, root_build_dir),
        ]

        if (pw_build_PYTHON_PIP_DOWNLOAD_ALL_PLATFORMS) {
          args += [ "--download-all-platforms" ]
        }

        deps = [
          ":${invoker.target_name}._compile_requirements($pw_build_PYTHON_TOOLCHAIN)",
          ":${invoker.target_name}._generate_3p_requirements($pw_build_PYTHON_TOOLCHAIN)",
        ]

        outputs = [
          _wheel_output_dir,
          _pip_wheel_logfile,
          _pip_download_logfile,
        ]
        pool = "$dir_pw_build/pool:pip($default_toolchain)"
      }

      # End pw_build_PYTHON_TOOLCHAIN check
    } else {
      group("${target_name}._compile_requirements") {
        public_deps = [ ":${target_name}($pw_build_PYTHON_TOOLCHAIN)" ]
      }
      group("${target_name}._install_3p_deps") {
        public_deps = [ ":${target_name}($pw_build_PYTHON_TOOLCHAIN)" ]
      }
      group("${target_name}.vendor_wheels") {
        public_deps = [ ":${target_name}($pw_build_PYTHON_TOOLCHAIN)" ]
      }
    }
  } else {
    group("${target_name}._compile_requirements") {
    }
    group("${target_name}._install_3p_deps") {
    }
    group("${target_name}.vendor_wheels") {
    }
  }

  # Have this target directly depend on _install_3p_deps
  group("$target_name") {
    public_deps =
        [ ":${target_name}._install_3p_deps($pw_build_PYTHON_TOOLCHAIN)" ]
  }
}