#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright (C) 2015 Xyne
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# (version 2) as published by the Free Software Foundation.
#
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import calendar
import hashlib
import io
import itertools
import json
import logging
import os
import pwd
import shlex
import shutil
import time
import urllib.request
import warnings

################################## Constants ###################################

DISPLAY_TIME_FORMAT = '%Y-%m-%d %H:%M:%S %Z'
RFC_2822_TIME_FORMAT = '%a, %d %b %Y %H:%M:%S %Z'

SIG_EXT = '.sig'



##################################### HTTP #####################################

def text_from_url(url):
  '''
  Retrieve text from a URL.
  '''
  logging.debug('retrieving data from {}'.format(url))
  with urllib.request.urlopen(url) as f:
    return (f.read().decode(),)



##################################### JSON #####################################

def load_json(path):
  '''
  Simple wrapper function to open a file and load the JSON in it.
  '''
  with open(path, 'r') as f:
    return json.load(f)



def load_json_from_url(url):
  '''
  Simple wrapper to load a JSON object from a remote URL.
  '''
  logging.debug('retrieving JSON data from {}'.format(url))
  with urllib.request.urlopen(url) as f:
    return json.loads(f.read().decode())



################################### Network ####################################

def mirror(url, path, cache_time=-1):
  '''
  Download a file at the given URL to a local path. The 'If-Modified-Since'
  header is used to avoid unnecessary downloads when supported by the server. A
  further option is provided to cache local files for a given interval before
  updating them.

  If the local file does not exist, the remote file will always be downloaded.

  url:
    The remote URL.

  path:
    The local path.

  cache_time:
    If negative, the file will be downloaded only if the server
    reports that a newer version is available by comparing the "Last-Modified"
    to the local file's last modification time.

    If 0, the file will be downloaded regardless of the last modification times.

    If positive, the file will only be downloaded if the local file's last
    modification time is older than this number of seconds AND the server has
    a newer version as determined by the "If-Modified-Since" header.
  '''
  headers = dict()

  if cache_time is None:
    cache_time = -1

  if cache_time != 0:
    try:
      mtime = os.path.getmtime(path)
    except FileNotFoundError:
      pass
    else:
      headers['If-Modified-Since'] = time.strftime(RFC_2822_TIME_FORMAT, time.gmtime(mtime))
      # The "> 0" check should be unnecessary but the user could mess with the mtime.
      if cache_time > 0 and (time.time() - mtime) < cache_time:
          return

  req = urllib.request.Request(url, headers=headers)
  logging.debug('requesting {}'.format(url))
  with urllib.request.urlopen(req) as f:
    # The file has not been modified.
    if f.code == 304:
      logging.debug('unmodified')
      return
    else:
      logging.debug('saving {} to {}'.format(url, path))
      # Download the file.
      with open(path, 'wb') as g:
        buf = f.read(io.DEFAULT_BUFFER_SIZE)
        g.write(buf)
        while(buf):
          buf = f.read(io.DEFAULT_BUFFER_SIZE)
          g.write(buf)
      # Synchronize the modification time with the server.
      try:
        server_mtime = f.headers['Last-Modified']
      except KeyError:
        pass
      else:
        if server_mtime is not None:
          mtime = calendar.timegm(time.strptime(server_mtime, RFC_2822_TIME_FORMAT))
          os.utime(path, (mtime, mtime))



#################################### Pushd #####################################

class Pushd(object):
  '''
  Temporarily change the working directory.
  '''
  def __init__(self, path=None):
    if path is None:
      self.path = os.getcwd()
    else:
      self.path = os.path.abspath(path)
    self.orig = None

  def __enter__(self):
    self.orig = os.getcwd()
    os.chdir(self.path)
    return self.path

  def __exit__(self, typ, val, traceback):
    os.chdir(self.orig)



################################## Lock File ###################################

class LockError(Exception):
  '''
  Errors raised by the lock file class.
  '''
  def __init__(self, path, msg, error=None):
    self.path = path
    self.msg = msg
    self.error = error

  def __str__(self):
    s = '{}({}): {}'.format(self.__class__.__name__, self.path, self.msg)
    if self.error:
      s += ' [{}]'.format(self.error)
    return s



class Lockfile(object):
  '''
  Lock file class.
  '''
  def __init__(self, path, name, mode=0o644):
    self.path = os.path.abspath(path)
    self.locked = False
    self.name = name
    self.mode = mode

  def error(self, locking, error=None):
    if locking:
      prefix = ''
    else:
      prefix = 'un'
    msg = 'failed to {}lock {}'.format(prefix, self.name)
    le = LockError(self.path, msg, error=error)
    logging.debug(str(le))
    raise le

  def __enter__(self):
    try:
      # Use O_EXCL to ensure that this will only succeed if no previous file
      # exists.
      flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
      fd = os.open(self.path, flags, mode=self.mode)
      with os.fdopen(fd, 'w') as f:
        f.write(str(os.getpid()))
    except OSError as e:
      self.error(True, error=e)
    else:
      self.locked = True
      return self.path

  def __exit__(self, typ, val, traceback):
    if self.locked:
      try:
        os.unlink(self.path)
      except OSError as e:
        self.error(False, error=e)
      else:
        self.locked = False



################################ File Functions ################################

def get_checksum(path, typ):
  '''
  Return a checksum of a local file as a hexadecimal digest.
  '''
  h = hashlib.new(typ)
  b = h.block_size
  try:
    with open(path, 'rb') as f:
      buf = f.read(b)
      while buf:
        h.update(buf)
        buf = f.read(b)
    return h.hexdigest()
  except FileNotFoundError:
    return None



def copy_file_and_maybe_sig(from_path, to_path, sig=False):
  '''
  Copy a file and its signature, if it exists.
  '''
  shutil.copyfile(from_path, to_path)
  if sig:
    a = from_path + SIG_EXT
    b = to_path + SIG_EXT
    try:
      shutil.copyfile(a, b)
    except FileNotFoundError as e:
      logging.debug('no matching sig {}'.format(from_path))
      return False
    else:
      return True
  else:
    return False



############################### Argument Parsing ###############################

def expand_short_args(args):
  '''
  Expand short arguments.

  Example: ["-Syu"] -> ["-S", "-y", "-u"]

  Arguments following '--' will not be expanded.
  '''
  expand = True
  for arg in args:
    if expand:
      try:
        if arg[0] == '-':
          if arg[1] == '-':
            expand = bool(arg[2:])
            yield arg
          else:
            for c in arg[1:]:
              yield '-' + c
          continue
      except IndexError:
        pass
    yield arg




def filter_arguments(args, remove):
  '''
  Iterate over arguments while omitting some.

  remove:
    A dictionary of the arguments to recognize and the number of parameters to
    discard. If the value is an interger then that number of arguments will be
    discarded after the recognized argument. If the argument is a string then
    arguments will be discarded until the next one that begins with that string.
  '''
  itr = iter(args)
  prefix = None
  for a in itr:
    if prefix is not None:
      if not a.startswith(prefix):
        continue
      else:
        prefix = None

    try:
      v = remove[a]
    except KeyError:
      yield a
    else:
      if isinstance(v, int):
        for i in range(v):
          a = next(itr)
      elif isinstance(v, str):
        prefix = v
      else:
        raise ValueError('filter_arguments: values in "remove" argument must be of type int or str, not {}'.format(type(v)))




############################ Configuration Files ##############################

def get_nested_key_or_none(obj, ks):
  '''
  Attempt to retrieve a value from an object using a list of nested keys.
  Return None if any KeyError occurs. This is useful for JSON objects with
  multiple nested mappings that may not exist.
  '''
  target = obj
  for k in ks:
    try:
      target = target[k]
    except KeyError:
      return None
  return target



##################################### Bash #####################################

def sh_comment_header(title):
  '''
  Section header comment.
  '''
  return '{:#^80s}'.format('# {} #'.format(title))



def sh_prompt_continue(question, expected_char):
  '''
  Prompt the user with a question and exit if the expected character is not
  given.
  '''

  c = expected_char[0]
  ans='{}{}'.format(c.upper(), c.lower())

  return '''{header}
read -p {question} -n 1 -r ans
echo
if [[ ! $ans =~ ^[{ans}]$ ]]
then
  exit 0
fi
'''.format(
    header=sh_comment_header('Prompt user before continuing.'),
    question=shlex.quote(question + ' '),
    ans=ans
  )



def interpolate_words(words, kwargs=None):
  '''
  Insert keyword arguments into given words using the string format method.
  '''
  if kwargs:
    for w in words:
      if isinstance(w, str):
        yield w.format(**kwargs)
      else:
        yield w
  else:
    for w in words:
      yield w


def ensure_strs(words):
  '''
  Ensure that all words in the iterable are strings.
  '''
  for w in words:
    if isinstance(w, str):
      yield w
    else:
      yield str(w)



def sh_quote_words(words, kwargs=None):
  '''
  Convert a command to a shell-escaped string.
  '''
  if kwargs:
    words = list(interpolate_words(words, kwargs))
  return ' '.join((shlex.quote(x) for x in ensure_strs(words)))



################################### Logging ####################################

def configure_logging(level=None, log=None):
  if level is None:
    level = logging.WARNING
  logging.basicConfig(
    format='%(levelname)s: %(message)s',
    filename=log,
    level=level
  )



############################# Deprecation Warnings #############################

def warn_deprecated(msg):
  '''
  Issue deprecation warnings.
  '''
  warnings.warn(
    msg,
    category=DeprecationWarning,
  )



def deprecated(f):
  '''
  Function decorator to issue deprecation warnings.
  '''
  def new_f(*args, **kwargs):
    warn_deprecated('function {} is deprecated'.format(f.__name__))
    return f(*args, **kwargs)
  new_f.__name__ = f.__name__
  new_f.__doc__ = f.__doc__
  new_f.__dict__.update(f.__dict__)
  return new_f



################################# Environment ##################################

def get_sudo_user_info():
  '''
  Get the name, uid and gid of the sudo user, if there is one.
  '''
  usr = os.getenv('SUDO_USER')
  if usr is not None:
    usrpwd = pwd.getpwnam(usr)
    uid = usrpwd[2]
    gid = usrpwd[3]
  else:
    uid = None
    gid = None
  return usr, uid, gid



############################## Number Formatting ###############################

def decimal_offsets(ss):
  '''
  Determine the offsets relative to the decimal point of the first and last
  digits of a number represented in scientific format ("{:e}").
  '''
  for s in ss:
    s = s.lower()
    n, e = s.rsplit('e', 1)
    try:
      b,a = n.split('.',1)
      a = a.rstrip('0')
    except ValueError:
      b = n
      a = ''
    e = int(e)
    f = e - len(a)
    yield e+1, f


def human_format_for_set_of_floats(xs, pad=True):
  '''
  Return a format string that will produce strings of equal width in the
  shortest and most natural format (which I have arbitrarily chosen).
  '''
  if pad:
    pad_char = '0'
  else:
    pad_char = ''

  n = len(xs)
  p = 0
  fmt = None
  ss = None
  while True:
    fmt = '{{:0.{:d}e}}'.format(p)
    ss = list(s for s in set(fmt.format(x) for x in xs) if float(s) in xs)
    if len(ss) < n:
      p += 1
    else:
      break

  offsets = list(decimal_offsets(ss))
  max_len = max(len(s) for s in ss)
  max_off = max(o[0] for o in offsets)
  min_off = min(o[1] for o in offsets)
  max_span = max_off - min(min_off, 0)
  if min_off < 0:
    max_span += 1
    if max_off < 0:
      max_span -= max_off - 1

  if max_span <= max_len:
    if min_off >= 0:
      fmt = '{{:{}{:d}.0f}}'.format(pad_char, max_span)
    else:
      fmt = '{{:{}{:d}.{:d}f}}'.format(pad_char, max_span, abs(min_off))

  return fmt
