# Copyright 2014 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import absolute_import, print_function

import atexit
import contextlib
import errno
import os
import shutil
import stat
import sys
import tempfile
import threading
import time
import zipfile
from collections import defaultdict
from contextlib import contextmanager
from datetime import datetime
from uuid import uuid4

# We use the start of MS-DOS time, which is what zipfiles use (see section 4.4.6 of
# https://pkware.cachefly.net/webdocs/casestudies/APPNOTE.TXT).
DETERMINISTIC_DATETIME = datetime(
  year=1980, month=1, day=1, hour=0, minute=0, second=0, tzinfo=None
)


def die(msg, exit_code=1):
  print(msg, file=sys.stderr)
  sys.exit(exit_code)


def safe_copy(source, dest, overwrite=False):
  def do_copy():
    temp_dest = dest + uuid4().hex
    shutil.copy(source, temp_dest)
    os.rename(temp_dest, dest)

  # If the platform supports hard-linking, use that and fall back to copying.
  # Windows does not support hard-linking.
  if hasattr(os, 'link'):
    try:
      os.link(source, dest)
    except OSError as e:
      if e.errno == errno.EEXIST:
        # File already exists.  If overwrite=True, write otherwise skip.
        if overwrite:
          do_copy()
      elif e.errno in (errno.EPERM, errno.EXDEV):
        # For a hard link across devices issue, fall back on copying.
        #
        # For a permission issue, the cause could be one of:
        # 1. We can't read source.
        # 2. We can't write dest.
        # 3. We don't own source but can read it.
        # Although we can't do anything about cases 1 and 2, case 3 is due to
        # `protected_hardlinks` (see: https://www.kernel.org/doc/Documentation/sysctl/fs.txt) and
        # we can fall back to copying in that case.
        #
        # See also https://github.com/pantsbuild/pex/issues/850 where this was discovered.
        do_copy()
      else:
        raise
  elif os.path.exists(dest):
    if overwrite:
      do_copy()
  else:
    do_copy()


# See http://stackoverflow.com/questions/2572172/referencing-other-modules-in-atexit
class MktempTeardownRegistry(object):
  def __init__(self):
    self._registry = defaultdict(set)
    self._getpid = os.getpid
    self._lock = threading.RLock()
    self._exists = os.path.exists
    self._getenv = os.getenv
    self._rmtree = shutil.rmtree
    atexit.register(self.teardown)

  def __del__(self):
    self.teardown()

  def register(self, path):
    with self._lock:
      self._registry[self._getpid()].add(path)
    return path

  def teardown(self):
    for td in self._registry.pop(self._getpid(), []):
      if self._exists(td):
        self._rmtree(td)


_MKDTEMP_SINGLETON = MktempTeardownRegistry()


class PermPreservingZipFile(zipfile.ZipFile, object):
  """A ZipFile that works around https://bugs.python.org/issue15795"""

  @classmethod
  def zip_info_from_file(cls, filename, arcname=None, date_time=None):
    """Construct a ZipInfo for a file on the filesystem.

    Usually this is provided directly as a method of ZipInfo, but it is not implemented in Python
    2.7 so we re-implement it here. The main divergance we make from the original is adding a
    parameter for the datetime (a time.struct_time), which allows us to use a deterministic
    timestamp. See https://github.com/python/cpython/blob/master/Lib/zipfile.py#L495."""
    st = os.stat(filename)
    isdir = stat.S_ISDIR(st.st_mode)
    if arcname is None:
      arcname = filename
    arcname = os.path.normpath(os.path.splitdrive(arcname)[1])
    while arcname[0] in (os.sep, os.altsep):
      arcname = arcname[1:]
    if isdir:
      arcname += '/'
    if date_time is None:
      date_time = time.localtime(st.st_mtime)
    zinfo = zipfile.ZipInfo(filename=arcname, date_time=date_time[:6])
    zinfo.external_attr = (st.st_mode & 0xFFFF) << 16  # Unix attributes
    if isdir:
      zinfo.file_size = 0
      zinfo.external_attr |= 0x10  # MS-DOS directory flag
    else:
      zinfo.file_size = st.st_size
    return zinfo

  def _extract_member(self, member, targetpath, pwd):
    result = super(PermPreservingZipFile, self)._extract_member(member, targetpath, pwd)
    info = member if isinstance(member, zipfile.ZipInfo) else self.getinfo(member)
    self._chmod(info, result)
    return result

  def _chmod(self, info, path):
    # This magic works to extract perm bits from the 32 bit external file attributes field for
    # unix-created zip files, for the layout, see:
    #   https://www.forensicswiki.org/wiki/ZIP#External_file_attributes
    attr = info.external_attr >> 16
    os.chmod(path, attr)


@contextlib.contextmanager
def open_zip(path, *args, **kwargs):
  """A contextmanager for zip files. Passes through positional and kwargs to zipfile.ZipFile."""
  with contextlib.closing(PermPreservingZipFile(path, *args, **kwargs)) as zip:
    yield zip


@contextlib.contextmanager
def temporary_dir(cleanup=True):
  td = tempfile.mkdtemp()
  try:
    yield td
  finally:
    if cleanup:
      safe_rmtree(td)


def safe_mkdtemp(**kw):
  """Create a temporary directory that is cleaned up on process exit.

  Takes the same parameters as tempfile.mkdtemp.
  """
  # proper lock sanitation on fork [issue 6721] would be desirable here.
  return _MKDTEMP_SINGLETON.register(tempfile.mkdtemp(**kw))


def register_rmtree(directory):
  """Register an existing directory to be cleaned up at process exit."""
  return _MKDTEMP_SINGLETON.register(directory)


def safe_mkdir(directory, clean=False):
  """Safely create a directory.

  Ensures a directory is present.  If it's not there, it is created.  If it
  is, it's a no-op. If clean is True, ensures the directory is empty.
  """
  if clean:
    safe_rmtree(directory)
  try:
    os.makedirs(directory)
  except OSError as e:
    if e.errno != errno.EEXIST:
      raise


def safe_open(filename, *args, **kwargs):
  """Safely open a file.

  ``safe_open`` ensures that the directory components leading up the
  specified file have been created first.
  """
  safe_mkdir(os.path.dirname(filename))
  return open(filename, *args, **kwargs)  # noqa: T802


def safe_delete(filename):
  """Delete a file safely. If it's not present, no-op."""
  try:
    os.unlink(filename)
  except OSError as e:
    if e.errno != errno.ENOENT:
      raise


def safe_rmtree(directory):
  """Delete a directory if it's present. If it's not present, no-op."""
  if os.path.exists(directory):
    shutil.rmtree(directory, True)


def safe_sleep(seconds):
  """Ensure that the thread sleeps at a minimum the requested seconds.

  Until Python 3.5, there was no guarantee that time.sleep() would actually sleep the requested
  time. See https://docs.python.org/3/library/time.html#time.sleep."""
  if sys.version_info[0:2] >= (3, 5):
    time.sleep(seconds)
  else:
    start_time = current_time = time.time()
    while current_time - start_time < seconds:
      remaining_time = seconds - (current_time - start_time)
      time.sleep(remaining_time)
      current_time = time.time()


class AtomicDirectory(object):
  def __init__(self, target_dir):
    self._target_dir = target_dir
    self._work_dir = '{}.{}'.format(target_dir, uuid4().hex)

  @property
  def work_dir(self):
    return self._work_dir

  @property
  def target_dir(self):
    return self._target_dir

  @property
  def is_finalized(self):
    return os.path.exists(self._target_dir)

  def finalize(self, source=None):
    """Rename `work_dir` to `target_dir` using `os.rename()`.

    :param str source: An optional source offset into the `work_dir`` to use for the atomic
                       update of `target_dir`. By default the whole `work_dir` is used.

    If a race is lost and `target_dir` already exists, the `target_dir` dir is left unchanged and
    the `work_dir` directory will simply be removed.
    """
    if self.is_finalized:
      return

    source = os.path.join(self._work_dir, source) if source else self._work_dir
    try:
      # Perform an atomic rename.
      #
      # Per the docs: https://docs.python.org/2.7/library/os.html#os.rename
      #
      #   The operation may fail on some Unix flavors if src and dst are on different filesystems.
      #   If successful, the renaming will be an atomic operation (this is a POSIX requirement).
      #
      # We have satisfied the single filesystem constraint by arranging the `work_dir` to be a
      # sibling of the `target_dir`.
      os.rename(source, self._target_dir)
    except OSError as e:
      if e.errno not in (errno.EEXIST, errno.ENOTEMPTY):
        raise e
    finally:
      self.cleanup()

  def cleanup(self):
    safe_rmtree(self._work_dir)


@contextmanager
def atomic_directory(target_dir, source=None):
  """A context manager that yields a new empty work directory path it will move to `target_dir`.

  :param str target_dir: The target directory to atomically update.
  :param str source: An optional source offset into the work directory to use for the atomic update
                     of the target directory. By default the whole work directory is used.

  If the `target_dir` already exists the enclosed block will be yielded `None` to signal there is
  no work to do.

  If the enclosed block fails the `target_dir` will be undisturbed.

  The new work directory will be cleaned up regardless of whether or not the enclosed block
  succeeds.
  """
  atomic_dir = AtomicDirectory(target_dir=target_dir)
  if atomic_dir.is_finalized:
    yield None
    return

  safe_mkdir(atomic_dir.work_dir)
  try:
    yield atomic_dir.work_dir
    atomic_dir.finalize(source=source)
  finally:
    atomic_dir.cleanup()


def chmod_plus_x(path):
  """Equivalent of unix `chmod a+x path`"""
  path_mode = os.stat(path).st_mode
  path_mode &= int('777', 8)
  if path_mode & stat.S_IRUSR:
    path_mode |= stat.S_IXUSR
  if path_mode & stat.S_IRGRP:
    path_mode |= stat.S_IXGRP
  if path_mode & stat.S_IROTH:
    path_mode |= stat.S_IXOTH
  os.chmod(path, path_mode)


def chmod_plus_w(path):
  """Equivalent of unix `chmod +w path`"""
  path_mode = os.stat(path).st_mode
  path_mode &= int('777', 8)
  path_mode |= stat.S_IWRITE
  os.chmod(path, path_mode)


def can_write_dir(path):
  """Determines if the directory at path can be written to by the current process.

  If the directory doesn't exist, determines if it can be created and thus written to.

  N.B.: This is a best-effort check only that uses permission heuristics and does not actually test
  that the directory can be written to with and writes.

  :param str path: The directory path to test.
  :return: `True` if the given path is a directory that can be written to by the current process.
  :rtype boo:
  """
  while not os.access(path, os.F_OK):
    parent_path = os.path.dirname(path)
    if not parent_path or (parent_path == path):
      # We've recursed up to the root without success, which shouldn't happen,
      return False
    path = parent_path
  return os.path.isdir(path) and os.access(path, os.R_OK | os.W_OK | os.X_OK)


def touch(file, times=None):
  """Equivalent of unix `touch path`.

  :file The file to touch.
  :times Either a tuple of (atime, mtime) or else a single time to use for both.  If not
  specified both atime and mtime are updated to the current time.
  """
  if times:
    if len(times) > 2:
      raise ValueError('times must either be a tuple of (atime, mtime) or else a single time value '
                       'to use for both.')

    if len(times) == 1:
      times = (times, times)

  with safe_open(file, 'a'):
    os.utime(file, times)


class Chroot(object):
  """A chroot of files overlayed from one directory to another directory.

  Files may be tagged when added in order to keep track of multiple overlays
  in the chroot.
  """
  class Error(Exception): pass
  class ChrootTaggingException(Error):
    def __init__(self, filename, orig_tag, new_tag):
      super(Chroot.ChrootTaggingException, self).__init__(  # noqa: T800
        "Trying to add %s to fileset(%s) but already in fileset(%s)!" % (
          filename, new_tag, orig_tag))

  def __init__(self, chroot_base):
    """Create the chroot.

    :chroot_base Directory for the creation of the target chroot.
    """
    try:
      safe_mkdir(chroot_base)
    except OSError as e:
      raise self.ChrootException('Unable to create chroot in %s: %s' % (chroot_base, e))
    self.chroot = chroot_base
    self.filesets = defaultdict(set)

  def clone(self, into=None):
    """Clone this chroot.

    :keyword into: (optional) An optional destination directory to clone the
      Chroot into.  If not specified, a temporary directory will be created.

    .. versionchanged:: 0.8
      The temporary directory created when ``into`` is not specified is now garbage collected on
      interpreter exit.
    """
    into = into or safe_mkdtemp()
    new_chroot = Chroot(into)
    for label, fileset in self.filesets.items():
      for fn in fileset:
        new_chroot.link(os.path.join(self.chroot, fn), fn, label=label)
    return new_chroot

  def path(self):
    """The path of the chroot."""
    return self.chroot

  def _normalize(self, dst):
    dst = os.path.normpath(dst)
    if dst.startswith(os.sep) or dst.startswith('..'):
      raise self.Error('Destination path is not a relative path!')
    return dst

  def _check_tag(self, fn, label):
    for fs_label, fs in self.filesets.items():
      if fn in fs and fs_label != label:
        raise self.ChrootTaggingException(fn, fs_label, label)

  def _tag(self, fn, label):
    self._check_tag(fn, label)
    self.filesets[label].add(fn)

  def _ensure_parent(self, path):
    safe_mkdir(os.path.dirname(os.path.join(self.chroot, path)))

  def copy(self, src, dst, label=None):
    """Copy file ``src`` to ``chroot/dst`` with optional label.

    May raise anything shutil.copy can raise, e.g.
      IOError(Errno 21 'EISDIR')

    May raise ChrootTaggingException if dst is already in a fileset
    but with a different label.
    """
    dst = self._normalize(dst)
    self._tag(dst, label)
    self._ensure_parent(dst)
    shutil.copy(src, os.path.join(self.chroot, dst))

  def link(self, src, dst, label=None):
    """Hard link file from ``src`` to ``chroot/dst`` with optional label.

    May raise anything os.link can raise, e.g.
      IOError(Errno 21 'EISDIR')

    May raise ChrootTaggingException if dst is already in a fileset
    but with a different label.
    """
    dst = self._normalize(dst)
    self._tag(dst, label)
    self._ensure_parent(dst)
    abs_src = src
    abs_dst = os.path.join(self.chroot, dst)
    safe_copy(abs_src, abs_dst, overwrite=False)
    # TODO: Ensure the target and dest are the same if the file already exists.

  def write(self, data, dst, label=None, mode='wb'):
    """Write data to ``chroot/dst`` with optional label.

    Has similar exceptional cases as ``Chroot.copy``
    """
    dst = self._normalize(dst)
    self._tag(dst, label)
    self._ensure_parent(dst)
    with open(os.path.join(self.chroot, dst), mode) as wp:
      wp.write(data)

  def touch(self, dst, label=None):
    """Perform 'touch' on ``chroot/dst`` with optional label.

    Has similar exceptional cases as Chroot.copy
    """
    dst = self._normalize(dst)
    self._tag(dst, label)
    touch(os.path.join(self.chroot, dst))

  def get(self, label):
    """Get all files labeled with ``label``"""
    return self.filesets.get(label, set())

  def files(self):
    """Get all files in the chroot."""
    all_files = set()
    for label in self.filesets:
      all_files.update(self.filesets[label])
    return all_files

  def labels(self):
    return self.filesets.keys()

  def __str__(self):
    return 'Chroot(%s {fs:%s})' % (self.chroot,
      ' '.join('%s' % foo for foo in self.filesets.keys()))

  def delete(self):
    shutil.rmtree(self.chroot)

  def zip(self, filename, mode='w', deterministic_timestamp=False):
    with open_zip(filename, mode) as zf:
      for f in sorted(self.files()):
        full_path = os.path.join(self.chroot, f)
        zinfo = zf.zip_info_from_file(
            filename=full_path,
            arcname=f,
            date_time=DETERMINISTIC_DATETIME.timetuple() if deterministic_timestamp else None
        )
        with open(full_path, 'rb') as open_f:
          data = open_f.read()
        zf.writestr(zinfo, data, compress_type=zipfile.ZIP_DEFLATED)