#!/usr/bin/env python
# Copyright European Organization for Nuclear Research (CERN)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# You may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#                       http://www.apache.org/licenses/LICENSE-2.0
#
# Authors:
# - Ralph Vigne, <ralph.vigne@cern.ch> 2015
#
# This script enforces the retention of the Rucio indices stored in Elasticsearch (ES).
# It also maintaines the aliases used by Rucio to be up to date.
#
# This script is supposed to run every day shortly after midnight (e.g. 00:02) (UTC).


import re
import traceback
import socket
import sys

from datetime import datetime, timedelta
from elasticsearch import Elasticsearch

# Uses the HTTP proxy running on any rucio logger nodes
ES_URL = 'rucio-logger-prod-01.cern.ch:80/elastic'

# Retention for non-debug entries in days
RETENTION = 30

# Retention for debug entries in days
DEBUG_RETENTION = 5

if __name__ == '__main__':
    es = Elasticsearch(ES_URL)
    today = datetime.today()
    general_cutoff = datetime.now() - timedelta(days=RETENTION)
    debug_cutoff = datetime.now() - timedelta(days=DEBUG_RETENTION)
    tbd = list()  # List with to be deleted indices
    log_index = 'logs-rucio-residue-%04d.%02d.%02d' % (today.year, today.month, today.day)

    try:
        for line in es.cat.indices(index='logs-rucio-*').split('\n'):
            # Lines look like this: green open logs-rucio-residue-2015.09.30       5 2  9883699 0    3.8gb   1.2gb
            general = re.match(r"(\w+)\s+(\w+)\s+(?P<index>\S+\-(?P<date>(?P<year>\d{4}).(?P<month>\d{2}).(?P<day>\d{2})))\s.*", line)
            debug = re.match(r"(\w+)\s+(\w+)\s+(?P<index>\S+-debug-(?P<date>(?P<year>\d{4}).(?P<month>\d{2}).(?P<day>\d{2})))\s.*", line)
            if debug and ("%s-%s-%s" % (debug.group('year'), debug.group('month'), debug.group('day')) < str(debug_cutoff.date())):
                tbd.append(debug.group('index'))
            elif general and ("%s-%s-%s" % (general.group('year'), general.group('month'), general.group('day')) < str(general_cutoff.date())):
                tbd.append(general.group('index'))
    except Exception as e:
        print 'Failed rerquesting index list from ES: %s' % str(e)
        traceback.print_exc()
        sys.exit(2)

    if not len(tbd):
        try:
            es.index(index=log_index,
                     doc_type='rucio-residue',
                     body={'message': 'No indices qualify for deletion with given retention policy (%s days / %s days (debug))' % (RETENTION, DEBUG_RETENTION),
                           'application': 'check_indices',
                           'severity': 6,
                           'severity_label': 'info',
                           'facility': 1,
                           'host': socket.getfqdn(),
                           '@timestamp': datetime.utcnow()})
        except Exception:
            traceback.print_exc()
            sys.exit(1)  # In Nagios 1 is  WARNING
    else:
        for index in tbd:
            try:
                es.indices.delete(index)
                es.index(index=log_index,
                         doc_type='rucio-residue',
                         body={'message': 'Successfully deleted index: %s' % index,
                               'application': 'check_indices',
                               'severity': 6,
                               'severity_label': 'info',
                               'facility': 1,
                               'host': socket.getfqdn(),
                               '@timestamp': datetime.utcnow()})
            except Exception as e:
                try:
                    es.index(index=log_index,
                             doc_type='rucio-residue',
                             body={'message': 'Failed deleting index: %s' % str(e),
                                   'application': 'check_indices',
                                   'severity': 3,
                                   'severity_label': 'error',
                                   'facility': 1,
                                   'host': socket.getfqdn(),
                                   '@timestamp': datetime.utcnow()})
                except Exception:
                    traceback.print_exc()
                    sys.exit(2)  # In Nagi os 2 is CRITICAL

    # Updating alias to include latest indices

    try:
        es.indices.put_alias(index='logs-rucio-daemons-*', name='rucio-daemon-logs')
        es.index(index=log_index,
                 doc_type='rucio-residue',
                 body={'message': 'Successfully update alias \"rucio-daemons-logs\"',
                       'application': 'check_indices',
                       'severity': 6,
                       'severity_label': 'info',
                       'facility': 1,
                       'host': socket.getfqdn(),
                       '@timestamp': datetime.utcnow()})
    except Exception as e:
        traceback.print_exc()
        try:
            es.index(index=log_index,
                     doc_type='rucio-residue',
                     body={'message': 'Failed updating alias: %s' % str(e),
                           'application': 'check_indices',
                           'severity': 2,
                           'severity_label': 'critical',
                           'facility': 1,
                           '@timestamp': datetime.utcnow()})
        except Exception as arg:
            traceback.print_exc()
            sys.exit(2)  # In Nagios 2 is CRITICAL

    sys.exit(0)  # Nagios for 'All Good'
