#!/usr/bin/env python2.7
# s3cfg.cli
# Command Line Interface to s3 cfg

# builtin
from datetime import datetime
import json
import re

# vendor
import boto
import yaml

class S3Cfg:
  '''
  A class to connect to S3 to manage configuration files. It has 3 main methods:

  get - Get a configuration file at a certain version
  put - Put a Configuration file with a given version (in the file)
  merge - Get multiple configuration files and merge them
  list - List all configuration files OR list all versions of one file

  S3Cfg works with multiple formats of configuration files and implies their use through extensions.
  Supported extensions are .json, .yaml, and .yml

  Internally, S3Cfg stores data as dicts and pickles to s3.
  '''

  def __init__(self, bucket='bucket', default_ext='json', version_field='version'):
    self.bucket = boto.connect_s3().create_bucket(bucket)
    self.default_ext = default_ext
    self.version_field = version_field
    self.name_rgx = re.compile('^(\w+)(\.(json|yaml|yml))?(@([\d\.\w]+))?$')

  def get(self, *names):
    '''
    Get configuration files from s3.
    Filenames look like as follows and are written to the current working directory.

    my-project
    my-project.ext
    my-project@version
    my-project.ext@version

    ext can be any supported extension.
    version can be any version already put or 'latest'
    '''
    for name in names:
      base, ext, version = self._parse(name)
      content = self._retrieve(base, version if version else 'latest')
      raw = self._encode(content, ext)

      with open('%s.%s' % (base, ext), 'w') as f:
        f.write(raw)
        print 'Wrote %d characters to %s.%s' % (len(raw), base, ext)

  def put(self, *names):
    '''
    Put configuration files in s3.
    Filenames must be local files in the current working directory.
    The extensions must be supported types with a version specifier in the file.

    Examples
    json: {"version": "1.2"}
    yaml: version: 1.2
    '''
    for name in names:
      base, ext, _ = self._parse(name)
      with open(name, 'r') as f:
        raw = f.read()
      content, version = self._decode(raw, ext)

      if version is None:
        raise VersionRequiredInConfigurationItem()

      try:
        # Should throw if no version exists
        self._retrieve(base, version)
        raise VersionAlreadyExists()
      except ConfigNotFound:
        self._push(base, version, content)
        print content

  def merge(self, *names, **kwargs):
    '''
    Get configuration files from s3 and merge them together.
    Filenames must follow get documentation. Later names take precedence.
    Merging is shallow. Output to the s3cfg.timestamp.json or an output file given.
    '''
    output = kwargs.get('output', None)

    def contents(*names):
      for name in names:
        base, ext, version = self._parse(name)
        yield self._retrieve(base, version if version else 'latest')

    merge = {}
    for content in contents(*names):
      merge.update(content)

    if output:
      base, ext, _ = self._parse(output)
    else:
      base, ext = 's3cfg.%s' % datetime.now().strftime('%Y-%m-%d.%H-%M-%S'), 'json'

    raw = self._encode(content, ext)


    with open('%s.%s' % (base, ext), 'w') as f:
      f.write(raw)
      print 'Wrote %d characters to %s.%s' % (len(raw), base, ext)

  def list(self, *names):
    '''
    If names were passed in, list all versions of it.
    If no name is given, list all configuration files in the bucket.
    '''
    if names:
      for name in names:
        base, _, _ = self._parse(name)
        keys = (key.name for key in self.bucket.list(prefix=base))
        versions = (key.split('@')[1] for key in keys)
        versions = [version for version in versions if version != 'latest']

        key_latest = self.bucket.new_key(base + '@latest')
        key_latest.get_contents_as_string() # retrieves remote metadata too
        latest_version = key_latest.get_metadata('version')

        print 'Listing all versions of %s (%d items). latest at %s' % (base, len(versions), latest_version)
        for version in versions:
          print version
        print ''

    else:
      keys = (key.name for key in self.bucket.list())
      names = [key.split('@')[0] for key in keys if key.split('@')[1] == 'latest']
      print 'Listing all stored configurations (%d items)' % len(names)
      for name in names:
        print name

  #### Everything below is helpers

  def _parse(self, name):
    '''Parse a config name. See get for allowable docs'''
    m = self.name_rgx.match(name)
    print  m.group(1, 3, 5)
    return m.group(1, 3, 5)

  def _retrieve(self, basename, version):
    '''Retrieve the object contents from s3 and decode into a python dict'''
    key = self.bucket.new_key(basename + '@' + version)

    try:
      raw = key.get_contents_as_string()
    except boto.exception.S3ResponseError as e:
      if e.message == 'The specified key does not exist.':
        raise ConfigNotFound()
      raise
    return json.loads(raw)

  def _push(self, basename, version, content):
    '''Encoding a python dict to a string and push it to s3 under the basename and version as well as latest.'''
    key = self.bucket.new_key(basename + '@' + version)
    key_latest = self.bucket.new_key(basename + '@latest')
    raw = json.dumps(content)
    key.set_contents_from_string(raw)
    key_latest.set_metadata('version', version)
    key_latest.set_contents_from_string(raw)
    print 'Pushed %d characters as %s version %s to s3' % (len(raw), basename, version)

  def _encode(self, data, ext):
    '''Encode a python dict (data) into some encoding defined by ext. See classdoc for allowable ext'''
    if ext in ('json',):
      return json.dumps(data, indent=2)
    elif ext in ('yaml', 'yml'):
      return yaml.dump(data, indent=2)
    else:
      raise ExtensionNotKnown(ext)

  def _decode(self, raw, ext):
    '''Decode raw content into a python dict as defined by ext. See classdoc for allowable ext'''
    if ext in ('json',):
      obj = json.loads(raw)
    elif ext in ('yaml', 'yml'):
      obj = yaml.load(raw)
    else:
      raise ExtensionNotKnown(ext)

    return obj, obj.get(self.version_field, None)


class ConfigNotFound(Exception):
  pass

class VersionAlreadyExists(Exception):
  pass

class VersionNotAllowed(Exception):
  pass

class ExtensionNotKnown(Exception):
  pass

class VersionRequiredInConfigurationItem(Exception):
  pass

if __name__ == '__main__':
  from argparse import ArgumentParser
  parser = ArgumentParser(description='Manage configuration files in s3')
  parser.add_argument('action', choices=('get', 'put', 'merge', 'list'), help='Action to perform')
  parser.add_argument('items', type=str, nargs='*', help='Items to operate on: my-project[.json][@version|latest]')
  parser.add_argument('--output', '-o', type=str, default=None, help='Output file for merge. Otherwise will be s3cfg.<date>.json')
  parser.add_argument('--bucket', type=str, required=True, help='Use something other than default bucket')
  parser.add_argument('--version-field', type=str, default='version')
  args = parser.parse_args()

  s3 = S3Cfg(bucket=args.bucket, version_field=args.version_field)

  if args.action == 'get':
    s3.get(*args.items)
  elif args.action == 'put':
    s3.put(*args.items)
  elif args.action == 'merge':
    s3.merge(*args.items, output=args.output)
  elif args.action == 'list':
    s3.list(*args.items)