# Copyright 2019 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.
"""Environment utility functions.

Usage:
api.environment.initialize(checkout_root=...)
with api.environment():
  ...
"""

import attr
from recipe_engine import recipe_api


@attr.s
class Package(object):
  name = attr.ib(type=bytes)
  version = attr.ib(type=bytes)


class EnvironmentApi(recipe_api.RecipeApi):
  """Environment utility functions."""

  def __init__(self, *args, **kwargs):
    super(EnvironmentApi, self).__init__(*args, **kwargs)
    self._cipd_dir = None
    self._path_prefixes = []
    self._ldpath_prefixes = []
    self._env = {}
    self._initialized = False

  def _init_cipd(self, checkout_root):
    """Install CIPD packages."""

    with self.m.step.nest('setup cipd'):
      json_paths = [
          checkout_root.join('env_setup', 'cipd', 'pigweed.json'),
          # TODO(mohrr) handle multiple files in non-predetermined places
      ]

      ensure_file = self.m.cipd.EnsureFile()
      for json_path in json_paths:
        packages = self.m.file.read_json(
            'read {}'.format(self.m.path.basename(json_path)), json_path)
        for pkg in packages:
          # JSON files are read as unicode, need to encode to get str.
          ensure_file.add_package(pkg['path'].encode(),
                                  ' '.join(x.encode() for x in pkg['tags']))

      self._cipd_dir = self.m.path['start_dir'].join('cipd')
      self.m.cipd.max_threads = 0
      self.m.cipd.ensure(self._cipd_dir, ensure_file)

      self._path_prefixes.append(self._cipd_dir)
      self._path_prefixes.append(self._cipd_dir.join('bin'))
      self._path_prefixes.append(self._cipd_dir.join('mingw64', 'bin'))
      self._ldpath_prefixes.append(self._cipd_dir)
      self._ldpath_prefixes.append(self._cipd_dir.join('lib'))
      self._env['PW_CIPD_INSTALL_DIR'] = self._cipd_dir
      self._env['PW_PIGWEED_CIPD_INSTALL_DIR'] = self._cipd_dir

  def _all_python_packages(self, checkout_root):
    """Return all folders with setup.py entries.

    Note: this function should not be called from within a virtualenv. It screws
    with how vpython works.

    Args:
      checkout_root(Path): root of source tree

    Returns:
      A list of package paths.
    """
    files = self.m.file.listdir('ls **/setup.py', checkout_root, recursive=True)
    matches = [
        self.m.path.dirname(x)
        for x in files
        if self.m.path.basename(x) == 'setup.py'
    ]
    with self.m.step.nest('packages') as step:
      step.logs['matches'] = [str(x) for x in matches]
    return matches

  def _init_python(self, checkout_root):
    """Initialize the Python environment. (Specifically, install 'pw'.)"""

    with self.m.step.nest('setup python'):
      packages = self._all_python_packages(checkout_root)

      with self.m.step.nest('setup virtualenv'), self():
        venv_dir = self.m.path['start_dir'].join('venv')
        cipd_python = 'python3'

        # Work around weird bug where venv creation requires the name of the
        # executable to be python.exe and not python3.exe on Windows.
        if self.m.platform.is_win:
          cipd_python = self._cipd_dir.join('bin', 'python3.exe')
          new_cipd_python = self._cipd_dir.join('bin', 'python.exe')
          self.m.file.copy(
              'cp python3.exe python.exe',
              cipd_python,
              new_cipd_python,
          )
          cipd_python = new_cipd_python

        self.m.step('create venv', [cipd_python, '-m', 'venv', venv_dir])

        venv_bin = venv_dir.join('bin')
        if not self.m.path.exists(venv_bin):
          alt_venv_bin = venv_dir.join('Scripts')
          # TODO(mohrr) figure out how to cover this.
          if self.m.path.exists(alt_venv_bin):  # pragma: no cover
            venv_bin = alt_venv_bin

        python = venv_bin.join('python')
        self.m.step('upgrade pip',
                    [python, '-m', 'pip', 'install', '--upgrade', 'pip'])

        # Need to insert at beginning because venv python should trump cipd
        # python.
        self._path_prefixes.insert(0, venv_bin)
        self._env['VIRTUAL_ENV'] = venv_dir

      # Need to exit and reenter 'with self()' to include new context from
      # creating the virtualenv.

      with self.m.step.nest('install packages'), self():
        pip_install_prefix = (python, '-m', 'pip', 'install')
        pw_cmd = list(pip_install_prefix)
        for package in packages:
          pw_cmd.append('--editable={}'.format(package))
        self.m.step('pigweed tools', pw_cmd)

        requirements = checkout_root.join(
            'env_setup/virtualenv/requirements.txt')
        req_cmd = list(pip_install_prefix)
        req_cmd.extend(('-r', requirements))
        self.m.step('build requirements', req_cmd)

  def init(self, checkout_root):
    if not self._initialized:
      with self.m.step.nest('environment'):
        # Setting _initialized immediately because some setup steps need to
        # use the context of previous steps, and invoking self() is the
        # easiest way to do so.
        self._initialized = True
        self._env['PW_ROOT'] = checkout_root
        self._init_cipd(checkout_root)
        self._init_python(checkout_root)

  def __call__(self):
    assert self._initialized

    env_prefixes = {}
    if self._path_prefixes:
      env_prefixes['PATH'] = self._path_prefixes
    if self._ldpath_prefixes:
      env_prefixes['LD_LIBRARY_PATH'] = self._ldpath_prefixes

    return self.m.context(env_prefixes=env_prefixes, env=self._env)