# -*- coding: utf-8 -*-
# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Downloads files upon request in a thread/process safe way.

DEPRECATED: Should be merged into chromite.lib.cache.
"""

from __future__ import print_function

import hashlib
import os
import shutil
import stat
import time

from chromite.lib import cros_logging as logging
from chromite.lib import locking
from chromite.lib import osutils
from chromite.lib.paygen import urilib
from chromite.lib.paygen import utils


FETCH_RETRY_COUNT = 10
DEFAULT_DAYS_TO_KEEP = 1
ONE_DAY = 24 * 60 * 60


class RetriesExhaustedError(Exception):
  """Raised when we make too many attempts to download the same file."""


def _DefaultFetchFunc(uri, cache_file):
  """The default fetch function.

  This simply downloads the uri into the cache file using urilib

  Args:
    uri: The URI to download.
    cache_file: The path to put the downloaded file in.
  """
  urilib.Copy(uri, cache_file)


class DownloadCache(object):
  """This class downloads files into a local directory upon request.

  This classes uses locking to make this safe across processes, and
  threads.

  Examples:
    # This will create the cache dir, and purge old contents.
    cache = DownloadCache('/tmp/my_cache')

    # file is copied into file, blocking for download if needed.
    cache.GetFileCopy('gs://bucket/foo', '/tmp/foo')
  """

  # Name of the purge management lock over the entire cache.
  _CACHE_LOCK = 'cache.lock'
  _FILE_DIR = 'cache'
  _LOCK_DIR = 'lock'

  _GET_FILE_SPIN_DELAY = 2

  def __init__(self, cache_dir, max_age=ONE_DAY, cache_size=None):
    """Create a DownloadCache.

    Since Purging is not performed very often, we can exceed max_age or
    cache_size.

    Args:
      cache_dir: The directory in which to create the cache.
      max_age: Purge files not used for this number of seconds. None for no
               max_age.
      cache_size: Purge the least recently used files until the cache is
                  below this size in bytes. None for no size limit.

      If no condition is provided, we purge all files unused for one full day.
    """
    # One directory for cached files, one for lock files.
    self._cache_dir = os.path.realpath(cache_dir)
    self._file_dir = os.path.join(self._cache_dir, self._FILE_DIR)
    self._lock_dir = os.path.join(self._cache_dir, self._LOCK_DIR)

    self._max_age = max_age
    self._cache_size = cache_size

    self._SetupCache()

  def _SetupCache(self):
    """Make sure that our cache contains only files/directories we expect."""
    try:
      osutils.SafeMakedirs(self._cache_dir)
      # The purge lock ensures nobody else is modifying the cache in any way.
      with self._PurgeLock(blocking=False, shared=False):
        # We have changed the layout of our cache directories over time.
        # Clean up any left over files.
        expected = (self._CACHE_LOCK, self._FILE_DIR, self._LOCK_DIR)
        unexpected = set(os.listdir(self._cache_dir)).difference(expected)

        for name in unexpected:
          filename = os.path.join(self._cache_dir, name)
          if os.path.isdir(filename):
            shutil.rmtree(filename)
          else:
            os.unlink(filename)

        # Create the cache file dir if needed.
        if not os.path.exists(self._file_dir):
          os.makedirs(self._file_dir)

        # Create the lock dir if needed.
        if not os.path.exists(self._lock_dir):
          os.makedirs(self._lock_dir)
    except locking.LockNotAcquiredError:
      # If we can't get an exclusive lock on the cache, someone else set it up.
      pass

  def _UriToCacheFile(self, uri):
    """Convert a URI to an cache file (full path).

    Args:
      uri: The uri of the file to be cached locally.

    Returns:
      The full path file name of the cache file associated with a given URI.
    """
    # We use the md5 hash of the URI as our file name. This allows us to
    # store all cache files in a single directory, which removes race
    # conditions around directories.
    m = hashlib.md5()
    m.update(uri.encode('utf-8'))
    return os.path.join(self._file_dir, m.hexdigest())

  def _PurgeLock(self, blocking=False, shared=False):
    """Acquire a lock on the cache as a whole.

    An exclusive lock proves nobody else will modify anything, and nobody
    else will hold any _CacheFileLocks. A shared lock is required before
    getting any kind of _CacheFileLock.

    Args:
      blocking: Block until the lock is available?
      shared: Get a shared lock, or an exclusive lock?

    Returns:
      Locking.FileLock (acquired)
    """
    lock_file = os.path.join(self._cache_dir, self._CACHE_LOCK)
    lock = locking.FileLock(lock_file, locktype=locking.FLOCK,
                            blocking=blocking)
    return lock.lock(shared)

  def _CacheFileLock(self, cache_file, blocking=False, shared=False):
    """Acquire a lock on a file in the cache.

    A shared lock will ensure no other processes are modifying the file, but
    getting it does not ensure that the file in question actually exists.

    An exclusive lock is required to modify a cache file, this usually means
    downloading it.

    A shared _PurgeLock should be held before trying to acquire any type
    of cache file lock.

    Args:
      cache_file: The full path of file in cache to lock.
      blocking: Block until the lock is available?
      shared: Get a shared lock, or an exclusive lock?

    Returns:
      Locking.FileLock (acquired)
    """
    lock_file = os.path.join(self._lock_dir, os.path.basename(cache_file))
    lock = locking.FileLock(lock_file, locktype=locking.FLOCK,
                            blocking=blocking)
    return lock.lock(shared)

  def Purge(self, max_age=None, cache_size=None, protect=None):
    """Attempts to clean up the cache contents.

    Is a no-op if cache lock is not acquirable.

    Args:
      max_age: Overrides the __init__ max_age for this one
                       purge. Mostly intended for unittests.
      cache_size: Overrides the __init__ cache_size for this one
                       purge. Mostly intended for unittests.
      protect: List of files to definitely not purge.
    """
    max_age = self._max_age if max_age is None else max_age
    cache_size = self._cache_size if cache_size is None else cache_size
    # TODO(crbug/1016555): remove or reduce logging once fixed.
    logging.info('Purge(max_age=%s, cache_size=%s)', max_age, cache_size)

    try:
      # Prevent other changes while we purge the cache.
      with self._PurgeLock(shared=False, blocking=False):

        # Purge files based on age, if specified.
        if max_age is not None:
          now = time.time()
          for f in utils.ListdirFullpath(self._file_dir):
            if (now - os.path.getmtime(f)) > max_age:
              # TODO(crbug/1016555): remove or reduce logging once fixed.
              logging.info(
                  'Purge removing %s. age=%d', f, now - os.path.getmtime(f))
              os.unlink(f)

        # Purge files based on size, if specified.
        if cache_size is not None:
          # Find cache files, and sort them so the oldest are first.
          # This defines which ones we will purge first.
          cache_files = utils.ListdirFullpath(self._file_dir)
          cache_files.sort(key=os.path.getmtime)

          sizes = [os.path.getsize(f) for f in cache_files]
          total_size = sum(sizes)

          # Remove files until we are small enough to fit.
          for f, size in zip(cache_files, sizes):
            if total_size < cache_size:
              break
            if f in protect:
              logging.info('Purge not removing protected %s', f)
              continue
            # TODO(crbug/1016555): remove or reduce logging once fixed.
            logging.info('Purge removing %s. total_size=%d', f, total_size)
            total_size -= size
            os.unlink(f)

        # Leave any lock files in place.  They can be used as is.
        # See crbug.com/1016555.

    except locking.LockNotAcquiredError:
      # If we can't get an exclusive lock on the file, it's in use, leave it.
      pass

  def _FetchIntoCache(self, uri, cache_file, fetch_func=_DefaultFetchFunc):
    """This function downloads the specified file (if not already local).

    You must hold the PurgeLock when calling this method.

    If it can't get an exclusive lock, or if the file is already present,
    it does nothing.

    Args:
      uri: The uri of the file.
      cache_file: The location in the cache to download too.
      fetch_func: Function to get the file.

    Returns:
      True if a file was downloaded, False otherwise. (used in unittests)

    Raises:
      May raise any download error associated with the URI's protocol.
    """
    try:
      # Write protect the file before modifying it.
      with self._CacheFileLock(cache_file, shared=False, blocking=False):
        if os.path.exists(cache_file):
          return False

        try:
          fetch_func(uri, cache_file)
          # Make the file read-only by everyone.
          os.chmod(cache_file, stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH)
        except Exception as e:
          # If there was any error with the download, make sure no partial
          # file was left behind.
          logging.info('Failed to fetch %s to %s', uri, cache_file)
          if os.path.exists(cache_file):
            os.unlink(cache_file)
          elif isinstance(e, OSError) and e.errno == 2:
            # TODO(crbug/1016555): This failure seems to be due to thread
            # exhaustion.  Log the state of tasks on the machine.
            os.execv('/bin/bash', [
                'bash', '-c',
                'exec ls -ld /proc/*/{exe,task/*} 2>/dev/null || true'])
          raise

    except locking.LockNotAcquiredError:
      # In theory, if it's already locked, that either means a download is in
      # progress, or there is a shared lock which means it's already present.
      return False

    # Try to cleanup the cache after we just grew it.
    self.Purge(protect=[cache_file])
    return True

  # TODO: Instead of hooking in fetch functions in the cache here, we could
  # set up protocol handlers which would know how to handle special cases
  # generally, identified by a protocol prefix like "prepimage://" or
  # "decompress://". That would help make sure they're handled consistently.
  def GetFileObject(self, uri, fetch_func=_DefaultFetchFunc):
    """Get an open readonly File object for the file in the cache.

    This method will populate the cache with the requested file if it's
    not already present, and will return an already opened read only file
    object for the cache contents.

    Even if the file is purged, this File object will remain valid until
    closed. Since this method is the only legitimate way to get access to
    a file in the cache, and it returns read only Files, cache files should
    never be modified.

    This method may block while trying to download and/or lock the file.

    Args:
      uri: The uri of the file to access.
      fetch_func: A function to produce the file if it isn't already in the
                  cache.

    Returns:
      File object opened with 'rb' mode.

    Raises:
      Exceptions from a failed download are passed through 'as is' from
      the underlying download mechanism.

      RetriesExhaustedError if we need a large number of attempts to
      download the same file.
    """
    cache_file = self._UriToCacheFile(uri)

    # We keep trying until we succeed, or throw an exception.
    for _ in range(FETCH_RETRY_COUNT):
      with self._PurgeLock(shared=True, blocking=True):
        # Attempt to download the file, if needed.
        self._FetchIntoCache(uri, cache_file, fetch_func)

        # Get a shared lock on the file. This can block if another process
        # has a non-shared lock (ie: they are downloading).
        with self._CacheFileLock(cache_file, shared=True, blocking=True):

          if os.path.exists(cache_file):
            fd = open(cache_file, 'rb')

            # Touch the timestamp on cache file to help purging logic.
            os.utime(cache_file, None)

            return fd
          else:
            # We don't have the file in our cache. There are three ways this
            # can happen:
            #
            # A) Another process was trying to download, blocked our download,
            #    then got a download error.
            # B) Another process removed the file(illegally). We will recover as
            #    soon as all read-only locks are released.
            # C) Our download failed without throwing an exception. We will
            #    block forever if this continues to happen.

            # Sleep so we don't spin too quickly, then try again.
            time.sleep(self._GET_FILE_SPIN_DELAY)

    raise RetriesExhaustedError(uri)

  def GetFileCopy(self, uri, filepath):
    """Copy a cache file into your file (downloading as needed).

    Copy the file into your specified filename (creating or overridding). It
    will be downloaded into the cache first, if needed. It is your
    responsibility to manage filepath after it is populated.

    Args:
      uri: The uri of the file to access.
      filepath: The name of the file to copy uri contents into.

    Raises:
      Exceptions from a failed download are passed through 'as is' from
      the underlying download mechanism.
    """
    with self.GetFileObject(uri) as src:
      with open(filepath, 'w+b') as dest:
        shutil.copyfileobj(src, dest)

  # Cache objects can be used with "with" statements.
  def __enter__(self):
    return self

  def __exit__(self, _type, _value, _traceback):
    self.Purge()