#!/usr/bin/env python3

# This script drives the clean up process of the CCDB backend of the QC. 
#
# It should ideally be ran as a cron on a machine. It uses plugins to implement 
# the actual actions defined in the config file config.yaml. Each item in the 
# config file describes which plugin (by name of the file) should be used for 
# a certain path in the CCDB. 
# 
# If several rules apply to an object, we pick the first one. Thus, mind carefully
# the order of the rules !
#
# The plugins should have a function "process()" that takes 3 arguments : 
# ccdb: Ccdb, object_path: str and delay: int
#
# We depend on requests, yaml, dryable, responses (to mock and test with requests)
#
# Usage
#         # run with debug logs and don't actually touch the database
#         PYTHONPATH=./rules:$PYTHONPATH./o2-qc-repo-cleaner --dry-run --log-level 10

import argparse
import logging
import requests
import re
import sys
import importlib
from typing import List
import tempfile
import socket

import dryable
import yaml
import time
import consul
import multiprocessing as mp
from pathlib import Path
from datetime import datetime

from qcrepocleaner.Ccdb import Ccdb
from qcrepocleaner.pidfile import PIDFile, AlreadyRunningError
import traceback

class Rule:
    """A class to hold information about a "rule" defined in the config file."""

    def __init__(self, object_path=None, delay=None, policy=None,  from_timestamp=None, to_timestamp=None,
                 continue_with_next_rule=None, all_params=None):
        '''
        Constructor.
        :param object_path: path to the object, or pattern, to which a rule will apply.
        :param delay: the grace period during which a new object is never deleted.
        :param policy: which policy to apply in order to clean up. It should correspond to a plugin.
        :param all_params: a map with all the parameters from the config file for this rule. We will keep only the
        extra ones.
        '''
        self.object_path = object_path
        self.delay = delay
        self.policy = policy
        self.from_timestamp = from_timestamp
        self.to_timestamp = to_timestamp
        self.continue_with_next_rule = continue_with_next_rule

        self.extra_params = all_params
        if all_params is not None:
            self.extra_params.pop("object_path")
            self.extra_params.pop("delay")
            self.extra_params.pop("policy")
            self.extra_params.pop("from_timestamp", 0)
            self.extra_params.pop("to_timestamp", 0)
            self.extra_params.pop("continue_with_next_rule", "False")

    def __repr__(self):
        return 'Rule(object_path={.object_path}, delay={.delay}, policy={.policy}, from_timestamp={.from_timestamp}, ' \
               'to_timestamp={.to_timestamp}, continue_with_next_rule={.continue_with_next_rule}, ' \
               'extra_params={.extra_params})'\
            .format(self, self, self, self, self, self, self)


def parseArgs():
    """Parse the arguments passed to the script."""
    logging.info("Parsing arguments")
    parser = argparse.ArgumentParser(description='Clean the QC database.')
    parser.add_argument('--config', dest='config', action='store', default="config.yaml",
                        help='Path to the config file')
    parser.add_argument('--config-git', action='store_true',
                        help='Check out the config file from git (branch repo_cleaner), ignore --config.')
    parser.add_argument('--config-consul', action='store',
                        help='Specify the consul url (without `http[s]://`) and port in the form of <url>:<port>,'
                             ' file must be stored in o2/components/qc/ANY/any/repoCleanerConfig.yaml,'
                             ' if specified ignore both --config and --config-git.')
    parser.add_argument('--log-level', dest='log_level', action='store', default="20",
                        help='Log level (CRITICAL->50, ERROR->40, WARNING->30, INFO->20,DEBUG->10)')
    parser.add_argument('--dry-run', action='store_true',
                        help='Dry run, no actual deletion nor modification to the CCDB.')
    parser.add_argument('--only-path', dest='only_path', action='store', default="",
                        help='Only work on given path (omit the initial slash).')
    parser.add_argument('--workers', dest='workers', action='store', default="1",
                        help='Number of parallel workers.')
    parser.add_argument('--only-path-no-subdir', action='store_true', default=False, help='Set to true if the '
                          'only-path points to an object rather than a folder or if subdirectories must be ignored.')
    args = parser.parse_args()
    dryable.set(args.dry_run)
    logging.info(args)
    return args


def parseConfig(config_file_path):
    """
    Read the config file and prepare a list of rules.

    Return a dictionary containing the list of rules and other config elements from the file.

    :param config_file_path: Path to the config file
    :raises yaml.YAMLError If the config file does not contain a valid yaml.
    """

    logging.info(f"Parsing config file {config_file_path}")
    with open(config_file_path, 'r') as stream:
        config_content = yaml.safe_load(stream)

    # also add something to the important logs file
    message = datetime.today().strftime('%Y-%m-%d - %H:%M:%S')
    storeCrucialLog("\n" + message + " - Start of the cleaner")

    rules = []
    logging.debug("Rules found in the config file:")

    for rule_yaml in config_content["Rules"]:
        logging.debug(f"rule_yaml: {rule_yaml}")
        if "from_timestamp" in rule_yaml:
            from_timestamp = rule_yaml["from_timestamp"]
        else:
            from_timestamp = 1
        if "to_timestamp" in rule_yaml:
            to_timestamp = rule_yaml["to_timestamp"]
        else:
            to_timestamp = 2785655701000  # 2058

        continue_with_next_rule = rule_yaml.get("continue_with_next_rule", False)

        rule = Rule(rule_yaml["object_path"], rule_yaml["delay"], rule_yaml["policy"],
                    from_timestamp, to_timestamp, continue_with_next_rule=continue_with_next_rule,
                    all_params=rule_yaml)
        rules.append(rule)
        logging.debug(f"   * {rule}")
        storeCrucialLog(f"   * {rule}")

    ccdb_url = config_content["Ccdb"]["Url"]

    return {'rules': rules, 'ccdb_url': ccdb_url}


def downloadConfigFromGit():
    """
    Download a config file from git.
    :return: the path to the config file
    """

    logging.debug("Get it from git")
    r = requests.get(
        'https://raw.github.com/AliceO2Group/QualityControl/repo_cleaner/Framework/script/RepoCleaner/config.yaml')
    logging.debug(f"config file from git : \n{r.text}")
    path = "/tmp/config.yaml"
    with open(path, 'w') as f:
        f.write(r.text)
    logging.info(f"Config path : {path}")
    return path


def downloadConfigFromConsul(consul_url, consul_port):
    """
    Download a config file from consul.
    :return: the path to the config file
    """

    logging.debug("Get it from consul")
    consul_server = consul.Consul(host=consul_url, port=consul_port)
    index, data = consul_server.kv.get(key='o2/components/qc/ANY/any/repoCleanerConfig.yaml')
    logging.debug(f"config file from consul : \n{data['Value']}")
    text = data["Value"].decode()
    logging.debug(f"config file from consul : \n{text}")
    path = "/tmp/repoCleanerConfig.yaml"
    with open(path, 'w') as f:
        f.write(text)
    logging.info(f"Config path : {path}")
    return path


def findMatchingRules(rules, object_path):
    """Return a list of all matching rules for the given path."""

    logging.debug(f"findMatchingRules for {object_path}")

    if object_path is None:
        logging.error(f"findMatchingRules: object_path is None")
        return []

    result = []
    for rule in rules:
        pattern = re.compile(rule.object_path)
        matched = pattern.match(object_path)
        if matched is not None:
            logging.debug(f"   Found! {rule}")
            result.append(rule)

    return result


filepath = tempfile.gettempdir() + "/repoCleaner.txt"
currentTimeStamp = int(time.time() * 1000)


def getTimestampLastExecution():
    """
    Returns the timestamp of the last execution.
    It is stored in a file in $TMP/repoCleaner.txt.
    :return: the timestampe of the last execution or 0 if it cannot find it.
    """
    try:
        f = open(filepath, "r")
    except IOError as e:
        logging.info(f"File {filepath} not readable, we return 0 as timestamp.")
        return 0
    timestamp = f.read()
    logging.info(f"Timestamp retrieved from {filepath}: {timestamp}")
    f.close()
    return timestamp


def storeSavedTimestamp():
    """
    Store the timestamp we saved at the beginning of the execution of this script.
    """
    try:
        f = open(filepath, "w+")
    except IOError:
        logging.error(f"Could not write the saved timestamp to {filepath}")
    f.write(str(currentTimeStamp))
    logging.info(f"Stored timestamp {currentTimeStamp} in {filepath}")
    f.close()


def storeMonitoringMetrics(success, duration):
    """
    Store the status and the duration in influxdb via telegraf for monitoring purpose.
    """
    socket_file="/tmp/telegraf.sock"
    if Path(socket_file).exists():
        telegraf = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
        telegraf.connect(socket_file)
        telegraf.send(f"repoCleaner success={success}".encode('utf-8'))
        telegraf.send(f"repoCleaner duration={duration}".encode('utf-8'))
        logging.info(f"Monitoring metrics stored.")
    else:
        logging.warning(f"File {socket_file} does not exist, no monitoring metrics stored.")


def storeCrucialLog(message):
    """
    Store few but very important messages to the file ~/repocleaner_logs.txt
    :param message:
    :return:
    """
    logs_filename = str(Path.home()) + "/repocleaner_logs.txt"   # very limited but important logs
    try:
        f = open(logs_filename, "a")
        f.write(message+"\n")
    except IOError as e:
        logging.error(f"Could not write crucial log to {logs_filename} : {e}")


def prepare_main_logger():
    logger = logging.getLogger()
    # Logging (split between stderr and stdout)
    formatter = logging.Formatter(fmt='%(asctime)s - %(levelname)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S')
    h1 = logging.StreamHandler(sys.stdout)
    h1.setLevel(logging.DEBUG)
    h1.addFilter(lambda record: record.levelno <= logging.INFO) # filter out everything that is above INFO level
    h1.setFormatter(formatter)
    logger.addHandler(h1)
    h2 = logging.StreamHandler(sys.stderr)
    h2.setLevel(logging.WARNING)     # take only warnings and error logs
    h2.setFormatter(formatter)
    logger.addHandler(h2)


def create_parallel_logger():
    # TODO merge with prepare_main_logger. It creates problems though.
    logger = mp.get_logger()
    logger.setLevel(logging.INFO)
    formatter = logging.Formatter(fmt='%(asctime)s - %(levelname)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S')
    h1 = logging.StreamHandler(sys.stdout)
    h1.setLevel(logging.DEBUG)
    h1.addFilter(lambda record: record.levelno <= logging.INFO) # filter out everything that is above INFO level
    h1.setFormatter(formatter)
    h2 = logging.StreamHandler(sys.stderr)
    h2.setLevel(logging.WARNING)     # take only warnings and error logs
    h2.setFormatter(formatter)
    if not len(logger.handlers):
        logger.addHandler(h2)
        logger.addHandler(h1)
    return logger


def read_config(args):
    path = args.config
    if args.config_consul:
        items = args.config_consul.split(':')
        path = downloadConfigFromConsul(items[0], items[1])
    elif args.config_git:
        path = downloadConfigFromGit()
    config = parseConfig(path)
    rules: List[Rule] = config['rules']
    ccdb_url = config['ccdb_url']
    return ccdb_url, rules


def process_object_wrapped(object_path, rules, ccdb, args):
    '''avoid getting blocked in parallel processing due to exceptions'''
    try:
        process_object(object_path, rules, ccdb, args)
    except:
        logging.error(f"Exception in process_object: {traceback.format_exc()}")


# @log_sparse
def process_object(object_path, rules, ccdb, args):
    logger = create_parallel_logger()
    logger.setLevel(int(args.log_level))
    logger.info(f"Processing {object_path}")

    # Take the first matching rule, if any
    rules = findMatchingRules(rules, object_path)
    logger.debug(f"Found {len(rules)} rules")

    if len(rules) == 0:
        logger.info(f"    no matching rule")
        return

    for rule in rules:
        logger.info(f"rule: {rule}")
        # Apply rule on object (find the plug-in script and apply)
        try:
            module = importlib.import_module('qcrepocleaner.rules.' + rule.policy)
            module.logger = logger
        except ModuleNotFoundError as err:
            logger.error(f"could not load module {rule.policy}")
            return
        try:
            stats = module.process(ccdb, object_path, int(rule.delay),  rule.from_timestamp, rule.to_timestamp,
                                   rule.extra_params)
        except Exception as e:
            logger.error(f"processing error:  {e}")

        logger.info(f"{rule.policy} applied on {object_path}: {stats}")

        if not rule.continue_with_next_rule:
            break


def run(args, ccdb_url, rules):

    # Get list of objects from CCDB
    ccdb = Ccdb(ccdb_url)
    ccdb.logger = logging.getLogger
    paths = ccdb.getObjectsList(getTimestampLastExecution(), args.only_path, args.only_path_no_subdir)
    if args.only_path != '':
        paths = [item for item in paths if item is not None and item.startswith(args.only_path)]
    logging.debug(paths)
    # For each object call the first matching rule
    logging.info("Loop through the objects and apply first matching rule.")

    logging.info(f"workers: {args.workers}")
    pool = mp.Pool(processes=int(args.workers))
    [pool.apply_async(process_object_wrapped, args=(object_path, rules, ccdb, args)) for object_path in paths]
    pool.close()
    pool.join()

    logging.info(f" *** DONE *** (total deleted: {ccdb.counter_deleted}, total updated: {ccdb.counter_validity_updated})")
    message = datetime.today().strftime('%Y-%m-%d-%H:%M:%S')
    storeCrucialLog(message + f" - End of the cleaner (total deleted: {ccdb.counter_deleted}, total updated: {ccdb.counter_validity_updated})")
    if not args.dry_run:
        storeSavedTimestamp()

# ****************
# We start here !
# ****************

def main():
    start_time = time.time()
    prepare_main_logger()

    # Parse arguments
    args = parseArgs()
    logging.getLogger().setLevel(int(args.log_level))

    try:
        with PIDFile(filename='o2-qc-repo-cleaner.pid'):
            ccdb_url, rules = read_config(args)
            run(args, ccdb_url, rules)
    except AlreadyRunningError:
        print('Already running. Exiting.')
    except:
        storeMonitoringMetrics(success=0, duration=time.time()-start_time)
        raise

    storeMonitoringMetrics(success=1, duration=time.time()-start_time)


if __name__ == "__main__":  # to be able to run the test code above when not imported.
    main()
