"""Elastic Index Retention lambda index module.

Registers backup repository and archives snapshots of logstash indexes
from AWS ElasticSearch to AWS S3.
"""

import datetime
import json
import os
import re
from boto.connection import AWSAuthConnection

HOST = os.environ.get('es_host',
                      'search-logging-ficffvusafv4pptagdqbb7gxvu'
                      '.us-east-1.es.amazonaws.com')
REPO_NAME = os.environ.get('es_repo_name', 'index_backups')
BUCKET = os.environ.get('backup_bucket', 'orcd-logging-backups')
ROLE_ARN = os.environ.get('role_arn',
                          'arn:aws:iam::103233932089:role/aws_es_assume_role')
REGION = os.environ.get('AWS_DEFAULT_REGION', 'us-east-1')
RETAIN_DAYS = os.environ.get('days_to_keep', '30')

LS_RE = re.compile(r'logstash-(?P<year>\d{4})\.(?P<month>\d{2})\.'
                   '(?P<day>\d{2})')
_CUT_DATE = datetime.datetime.now() - datetime.timedelta(days=int(RETAIN_DAYS))


class ESConnection(AWSAuthConnection):
    """Auth wrapper for requests to AWS ES."""

    def __init__(self, **kwargs):
        """Construct custom connection."""
        super(ESConnection, self).__init__(**kwargs)
        self._set_auth_service_name('es')

    def _required_auth_capability(self):
        return ['hmac-v4']


def check_snapshot_repository(client, reponame):
    """Check if snapshot repo exists."""
    print('Checking if snapshot repository exists')
    resp = client.make_request(method='GET',
                               path='/_snapshot/{}'.format(reponame))
    return bool(resp.status == 200)


def register_snapshot_repository(client, reponame):
    """Register index snapshot repo."""
    print('Registering Snapshot Repository')
    data = {'type': 's3',
            'settings': {'bucket': BUCKET,
                         'region': REGION,
                         'role_arn': ROLE_ARN}}
    resp = client.make_request(method='POST',
                               path='/_snapshot/{}'.format(reponame),
                               data=json.dumps(data))
    if resp.status == 200:
        print('Snapshot repository %s created' % reponame)
    else:
        print(resp.read())


def verify_snapshot_repository(client, reponame):
    """Verify if index snapshot repo works."""
    resp = client.make_request(method='POST',
                               path='/_snapshot/{}/_verify'.format(reponame))
    return bool(resp.status == 200)


def snapshot_old_indices(client, reponame, indices):
    """Snapshot old indexes."""
    snapshot_path = ('_snapshot/{}'
                     '/lambda_{}'.format(reponame,
                                         datetime.date.today().isoformat()))
    resp = client.make_request(method='GET',
                               path=snapshot_path)
    if resp.status == 200:
        print('Snapshot lambda_{} already '
              'exists'.format(datetime.date.today().isoformat()))
        return True
    data = {'indices': ','.join(indices),
            'ignore_unavailable': 'true',
            'include_global_state': 'false'}
    resp = client.make_request(method='PUT',
                               path=snapshot_path,
                               data=json.dumps(data),
                               params={'wait_for_completion': 'true'})
    print(resp.read())
    if resp.status == 200:
        print('Snapshot lambda_{} '
              'created'.format(datetime.date.today().isoformat()))
    return bool(resp.status == 200)


def get_indices_in_all_snapshots(client, reponame):
    """Get a list of already snapshotted indexes."""
    resp = client.make_request(method='GET',
                               path='_snapshot/{}/_all'.format(reponame))
    indices = []
    if resp.status == 200:
        body = json.loads(resp.read())
        for snapshot in body.get('snapshots'):
            if snapshot.get('state') == 'SUCCESS':
                indices.extend(snapshot.get('indices'))
    return list(set(indices))


def get_indices_in_snapshot(client, reponame, snapshot):
    """Get a list of indexes in snapshot."""
    resp = client.make_request(method='GET',
                               path='_snapshot/{}/{}'.format(reponame,
                                                             snapshot))
    if resp.status != 200:
        return []
    body = json.loads(resp.read())
    snapshot = body['snapshots'][0]
    if snapshot.get('state') == 'SUCCESS':
        return snapshot.get('indeces')
    else:
        return []


def delete_old_indices(client, indices):
    """Clean up indexes from ES."""
    for index in indices:
        resp = client.make_request(method='DELETE', path='/{}'.format(index))
        if resp.status != 200:
            print('Unable to delete {}'.format(index))
        else:
            print('{} deleted'.format(index))


def get_all_logstash_indices(client):
    """List all logstash name formatted indexes."""
    resp = client.make_request(method='GET', path='/*/_stats/store')
    body = json.loads(resp.read())
    if body.get('indices'):
        return sorted([i for i in body.get('indices').keys()
                       if LS_RE.search(i)])


def logstash_date_compare(index_name):
    """Compare index name to cut out date."""
    m_obj = LS_RE.match(index_name)
    idx_date_str = '{year}-{month}-{day}'.format(**m_obj.groupdict())
    idx_date = datetime.datetime.strptime(idx_date_str, '%Y-%m-%d')
    return idx_date < _CUT_DATE


def handler(event=None, context=None):
    """Lambda entry point."""
    client = ESConnection(host=HOST)
    if not check_snapshot_repository(client, REPO_NAME):
        register_snapshot_repository(client, REPO_NAME)
    verify_snapshot_repository(client, REPO_NAME)
    old_indices = list(filter(logstash_date_compare,
                              get_all_logstash_indices(client)))
    print('Snapshotting indices: {}'.format(', '.join(old_indices)))
    snapshotted_indices = get_indices_in_all_snapshots(client, REPO_NAME)
    old_snapshotted_indices = \
        list(set(snapshotted_indices).intersection(set(old_indices)))
    if snapshot_old_indices(client, REPO_NAME, old_indices):
        print('Will delete these indices: '
              '{}'.format(', '.join(old_snapshotted_indices)))
        delete_old_indices(client, old_snapshotted_indices)


if __name__ == '__main__':
    handler()