"""Script to help with the creation of python virtual envs."""
import glob
import logging
import os
from pathlib import Path
import sys
from typing import Iterable

DEFAULT_LOCAL_PACKAGES = 'python_packages'
VENV_RECIPE_PY2 = '~/venv_recipe/bin'


def _get_pip(venv_root):
  return os.path.join(venv_root, 'bin', 'pip')


def add_import_paths(venv_root: Path, paths: Iterable[Path]) -> None:
  """Add an entry into the custom import locations for a venv.

  This will create a python path file in the provided virtual
  environment such that python can import modules from
  the provided path.

  Args:
    venv_root: The root directory of a virtualenv to install into.
    paths: The paths to add
  """
  venv_lib = os.path.join(venv_root, 'lib', '*')
  for version_lib in glob.glob(venv_lib):
    pathfile = os.path.join(version_lib, 'site-packages', 'google.pth')
    try:
      with open(pathfile, 'r') as f:
        content = f.read()
      old_paths = [Path(p) for p in content.split('\n') if p]
      new_paths = sorted(old_paths + paths)
    except IOError:
      new_paths = paths
    with open(pathfile, 'w') as f:
      f.write('\n'.join(str(p) for p in new_paths))
    logging.info('Wrote the following paths to %s: %s', pathfile, paths)


def create_virtual_env(
    executor, directory, gcs_dir, local_packages=None, py3=False):
  """Creates a python virtual env.

  Args:
    executor: Executes the subprocess.
    directory: Directory to execute the subprocess in.
    gcs_dir: The directory that mirrors the gcs bucket
    local_packages: Relative path to python packages under ~/gcs-mirror
    py3: if True use python3

  Returns:
    The path to the new virtualenv directory root.
  """
  python_sandbox_dir = os.path.join(directory, 'python-sandbox')
  env = {}
  if py3:
    cmd = ['python3', '-m', 'venv', python_sandbox_dir]
  else:
    cmd = ['python2', '-m', 'virtualenv',
           '--extra-search-dir={}'.format(
               _get_find_links_path(gcs_dir, local_packages)),
           '--no-download',
           python_sandbox_dir]
    env['PATH'] = os.pathsep.join([
        os.path.expanduser(VENV_RECIPE_PY2), os.getenv('PATH')])
  executor.exec_subprocess(cmd, env=env, check_output=True)
  return python_sandbox_dir


def install_package(executor, venv_root, path, gcs_dir, local_packages=None):
  """Install a local package from disk.

  Args:
    executor: Executes the subprocess.
    venv_root: The root directory of a virtualenv to install into.
    path: The path to the package on local disk
    gcs_dir: The path to the gcs mirror directory.
    local_packages: Relative path to python packages under ~/gcs-mirror
  """
  # Perform necessary default package upgrades.
  _update_default_deps(
      executor,
      venv_root,
      ['pip', 'wheel', 'setuptools'],
      local_packages,
      gcs_dir
  )

  venv_pip = _get_pip(venv_root)

  executor.exec_subprocess(
      [venv_pip, 'install', '-e', path],
      check_output=True)


def install_py_deps(
    executor, venv_root, requirements, gcs_dir, local_packages=None):
  """Install the python dependencies.

  Args:
    executor: Executes the subprocess.
    venv_root: The root directory of a virtualenv to install |requirements| in.
    requirements: The path to the requirements.txt to install.
    gcs_dir: The path to the gcs mirror directory.
    local_packages: Relative path to python packages under ~/gcs-mirror
  """
  if not os.path.isfile(requirements):
    logging.warning('No requirements file found: %s', requirements)
    return

  # Perform necessary default package upgrades.
  _update_default_deps(
      executor,
      venv_root,
      ['pip', 'wheel', 'setuptools'],
      local_packages,
      gcs_dir
  )

  venv_pip = _get_pip(venv_root)

  executor.exec_subprocess(
      [venv_pip, 'install', '--no-index',
       '--find-links=file://{}'.format(
           _get_find_links_path(gcs_dir, local_packages)),
       '-r', requirements],
      check_output=True)


def _update_default_deps(executor, venv_root, deps, local_packages, gcs_dir):
  """Update the pip packages required to install dependencies.

  Args:
    executor: Executes the subprocess.
    venv_root: The root directory of the virtualenv.
    deps: List of pip packages to update.
    local_packages: Relative path to python packages under ~/gcs-mirror.
    gcs_dir: The path to the gcs mirror directory.
  """
  venv_pip = _get_pip(venv_root)
  for dep in deps:
    executor.exec_subprocess(
        [venv_pip, 'install', '--no-index',
         '--find-links=file://{}'.format(
             _get_find_links_path(gcs_dir, local_packages)),
         '--upgrade', dep],
        check_output=True)


def exec_within_virtual_env(
    executor, venv_root, cmd, env=None):
  """Executes |cmd| within a python virtualenv.

  Args:
    executor: Executes the subprocess.
    venv_root: The root directory of a virtualenv to run in.
    cmd: The command list to pass to virtualenv python interpreter. This will
        be prepended with virtualenv python, so it should not start with python.
        (e.g. If you want to run `python my_python_script.py --flag1 --arg=3`
        then pass cmd = ['my_python_script.py', '--flag1', '--arg=3'].)
    env: A dictionary with additional environment variables.

  Returns:
    returncode, stdout, stderr from the subprocess.
  """
  venv_python = os.path.join(venv_root, 'bin', 'python')
  full_cmd = [venv_python] + cmd
  return executor.exec_subprocess(full_cmd, env=env)


def _get_find_links_path(gcs_dir: str, local_packages=None):
  if not local_packages:
    local_packages = DEFAULT_LOCAL_PACKAGES
  return os.path.expanduser(os.path.join(gcs_dir, local_packages))


class ActivateVenv(object):
  """Context manage to activate python virtual environment.

  Will deactive on exit, and revert back to the previous venv.
  This is designed to run from inside another venv. If not, it
  will raise KeyError since sys.real_prefix only exists
  when Python is run in a venv. Since our infras always run
  from a venv, this is an acceptable assumption.

  Args:
    venv_root: Path to the python virtual env root.
  """
  def __init__(self, venv_root):
    self._venv_root = venv_root
    self._prev_os_path = None
    self._prev_sys_path = None
    self._prev_prefix = None
    self._prev_real_prefix = None
    self._save_current_env()

  def __enter__(self):
    env_activation_file = os.path.join(
        self._venv_root, 'bin', 'activate_this.py')
    exec(compile(open(env_activation_file, "rb").read(),
                 env_activation_file, 'exec'),
                 dict(__file__=env_activation_file))

  def __exit__(self, *args):
    sys.path[:0] = self._prev_sys_path
    sys.prefix = self._prev_prefix
    os.environ['PATH'] = self._prev_os_path
    if self._prev_real_prefix is not None:
      sys.real_prefix = self._prev_real_prefix

  def _save_current_env(self):
    self._prev_os_path = os.environ['PATH']
    self._prev_sys_path = list(sys.path)
    self._prev_prefix = sys.prefix
    # real_prefix attribute only exists if Python is run from a venv.
    if hasattr(sys, 'real_prefix'):
      self._prev_real_prefix = sys.real_prefix