#!/usr/bin/python3
# -*- coding: utf-8 -*-
#
# (c) 2008 Mandriva, http://www.mandriva.com/
#
# $Id$
#
# This file is part of Pulse 2, http://pulse2.mandriva.org
#
# Pulse 2 is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# Pulse 2 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 Pulse 2; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
# MA 02110-1301, USA.

import sys
from sqlalchemy import *
import sqlalchemy
import logging
import datetime
import getopt

STATES_TO_CLEANUP = ['done', 'failed', 'stop', 'stopped', 'over_timed']

logger = logging.getLogger(sys.argv[0])
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
logger.addHandler(handler)
logger.setLevel(logging.DEBUG)

def query_bundles_to_delete(commands, commands_on_host, creation_timestamp):
    """
    Get bundle and related commands to delete.
    A bundle is deleted only if all its commands can be deleted

    @rtype: list
    @returns: list of bundle ids and command ids to remove
    """
    if sqlalchemy.__version__.startswith('0.3'):
        if msc_dbversion > 13:
            # The not finished bundles are those that own commands where all
            # commands_on_host state are not in the done or failed state
            not_finished_bundles = select([commands.c.fk_bundle], and_(commands_on_host.c.fk_commands == commands.c.id, not_(commands_on_host.c.current_state.in_(STATES_TO_CLEANUP)), commands.c.fk_bundle != None), distinct = True)
            # Get the finished bundles with commands that haven't expired
            finished_bundles_not_expired = select([commands.c.fk_bundle], and_(not_(commands.c.fk_bundle.in_(not_finished_bundles)), commands.c.creation_date >= creation_timestamp), distinct = True)
            # The finished and expired bundles are those that are finished, and
            # don't have commands that haven't expired
            to_delete = select([commands.c.fk_bundle, commands.c.id], and_(not_(commands.c.fk_bundle.in_(finished_bundles_not_expired)), not_(commands.c.fk_bundle.in_(not_finished_bundles)), commands.c.fk_bundle != None), distinct = True)
        else:
            # The not finished bundles are those that own commands where all
            # commands_on_host state are not in the done or failed state
            not_finished_bundles = select([commands.c.bundle_id], and_(commands_on_host.c.fk_commands == commands.c.id, not_(commands_on_host.c.current_state.in_(STATES_TO_CLEANUP)), commands.c.bundle_id != None), distinct = True)
            # Get the finished bundles with commands that haven't expired
            finished_bundles_not_expired = select([commands.c.bundle_id], and_(not_(commands.c.bundle_id.in_(not_finished_bundles)), commands.c.creation_date >= creation_timestamp), distinct = True)
            # The finished and expired bundles are those that are finished, and
            # don't have commands that haven't expired
            to_delete = select([commands.c.bundle_id, commands.c.id], and_(not_(commands.c.bundle_id.in_(finished_bundles_not_expired)), not_(commands.c.bundle_id.in_(not_finished_bundles)), commands.c.bundle_id != None), distinct = True)
    else:
        if msc_dbversion > 13:
            # The not finished bundles are those that own commands where all
            # commands_on_host state are not in the done or failed state
            not_finished_bundles = select([commands.c.fk_bundle]).select_from(commands.join(commands_on_host)).where(and_(not_(commands_on_host.c.current_state.in_(STATES_TO_CLEANUP)), commands.c.fk_bundle != None)).distinct()
            # Get the finished bundles with commands that haven't expired
            finished_bundles_not_expired = select([commands.c.fk_bundle]).where(and_(commands.c.creation_date >= creation_timestamp, commands.c.fk_bundle != None)).distinct()
            # The finished and expired bundles are those that are finished, and
            # don't have commands that haven't expired
            to_delete = select([commands.c.fk_bundle, commands.c.id]).where(and_(not_(commands.c.fk_bundle.in_(finished_bundles_not_expired)), not_(commands.c.fk_bundle.in_(not_finished_bundles)), commands.c.fk_bundle != None)).distinct()
        else:
            # The not finished bundles are those that own commands where all
            # commands_on_host state are not in the done or failed state
            not_finished_bundles = select([commands.c.bundle_id]).select_from(commands.join(commands_on_host)).where(and_(not_(commands_on_host.c.current_state.in_(STATES_TO_CLEANUP)), commands.c.bundle_id != None)).distinct()
            # Get the finished bundles with commands that haven't expired
            finished_bundles_not_expired = select([commands.c.bundle_id]).where(and_(commands.c.creation_date >= creation_timestamp, commands.c.bundle_id != None)).distinct()
            # The finished and expired bundles are those that are finished, and
            # don't have commands that haven't expired
            to_delete = select([commands.c.bundle_id, commands.c.id]).where(and_(not_(commands.c.bundle_id.in_(finished_bundles_not_expired)), not_(commands.c.bundle_id.in_(not_finished_bundles)), commands.c.bundle_id != None)).distinct()
    ret = to_delete.execute().fetchall()
    return ret

def query_commands_to_delete(commands, commands_on_host, creation_timestamp):
    """
    Get commands to delete, not part of a bundle

    @rtype: list
    @returns: list of command ids to delete
    """
    if sqlalchemy.__version__.startswith('0.3'):
        if msc_dbversion > 13:
            not_finished_commands = select([commands.c.id], and_(commands_on_host.c.fk_commands == commands.c.id, not_(commands_on_host.c.current_state.in_(STATES_TO_CLEANUP)), commands.c.fk_bundle == None, commands.c.creation_date < creation_timestamp), distinct = True)
            to_delete = select([commands.c.id], and_(not_(commands.c.id.in_(not_finished_commands)), commands.c.fk_bundle == None, commands.c.creation_date < creation_timestamp), distinct = True)
        else:
            not_finished_commands = select([commands.c.id], and_(commands_on_host.c.fk_commands == commands.c.id, not_(commands_on_host.c.current_state.in_(STATES_TO_CLEANUP)), commands.c.bundle_id == None, commands.c.creation_date < creation_timestamp), distinct = True)
            to_delete = select([commands.c.id], and_(not_(commands.c.id.in_(not_finished_commands)), commands.c.bundle_id == None, commands.c.creation_date < creation_timestamp), distinct = True)
    else:
        if msc_dbversion > 13:
            not_finished_commands = select([commands.c.id]).select_from(commands.join(commands_on_host)).where(and_(not_(commands_on_host.c.current_state.in_(STATES_TO_CLEANUP)), commands.c.fk_bundle == None, commands.c.creation_date < creation_timestamp)).distinct()
            to_delete = select([commands.c.id]).where(and_(not_(commands.c.id.in_(not_finished_commands)), commands.c.fk_bundle == None, commands.c.creation_date < creation_timestamp)).distinct()
        else:
            not_finished_commands = select([commands.c.id]).select_from(commands.join(commands_on_host)).where(and_(not_(commands_on_host.c.current_state.in_(STATES_TO_CLEANUP)), commands.c.bundle_id == None, commands.c.creation_date < creation_timestamp)).distinct()
            to_delete = select([commands.c.id]).where(and_(not_(commands.c.id.in_(not_finished_commands)), commands.c.bundle_id == None, commands.c.creation_date < creation_timestamp)).distinct()
    ret = to_delete.execute().fetchall()
    return ret

def run_query(msg, connection, query):
    logger.info(msg)
    ret = connection.execute(query)
    logger.info('Count of deleted rows: %ld' % ret.rowcount)

def msc_delete(connection, bundles_to_delete, commands_to_delete, bundle, commands, target, commands_on_host, commands_history):
    in_coh_op = commands_on_host.c.fk_commands.in_(commands_to_delete)
    in_c_op = commands.c.id.in_(commands_to_delete)
    in_b_op = bundle.c.id.in_(bundles_to_delete)

    query = commands_history.delete(commands_history.c.fk_commands_on_host.in_(select([commands_on_host.c.id], in_coh_op)))
    run_query('Purging commands_history table', connection, query)

    query = target.delete(target.c.id.in_(select([commands_on_host.c.fk_target], in_coh_op)))
    run_query('Purging target table', connection, query)

    query = commands_on_host.delete(in_coh_op)
    run_query('Purging commands_on_host table', connection, query)

    query = commands.delete(in_c_op)
    run_query('Purging commands table', connection, query)

    if bundles_to_delete:
        query = bundle.delete(in_b_op)
        run_query('Purging bundle table', connection, query)

def optimize_tables():

    for table in ['commands_history', 'target', 'commands_on_host', 'commands', 'bundle']:
        logger.info("Optimising table %s" % (table))
        res = text('OPTIMIZE TABLE %s' % table, bind=mysql_db).execute().fetchone()
        if sqlalchemy.__version__.startswith('0.3'):
            logger.info("Table %s : optimisation result is '%s'" % (table, res[3]))
        else:
            logger.info("Table %s : optimisation done" % (table)) # FIXME: should check result ...

def usage(argv):
    print >> sys.stderr, 'Usage: %s --uri=<db_conn_string> --days=<days> [--force]' % argv[0]
    print >> sys.stderr, 'Purge command older than <day> and whom members\' current state belong to %s' % ', '.join(STATES_TO_CLEANUP)
    print >> sys.stderr, 'Where db_conn_string is a SQLAlchemy connection string, e.g. mysql://user:password@host/database'
    print >> sys.stderr, 'Where days is a positive integer specifying from which age the commands should be purged, e.g. 30 for 30 days'
    print >> sys.stderr, 'If --force flag is given, purge commands regardless their current state'
    return 1

if __name__ == "__main__":
    try:
        opts, args = getopt.getopt(sys.argv[1:], "", ["uri=", "force", "days="])
    except getopt.GetoptError, err:
        print >> sys.stderr, err
        sys.exit(usage(sys.argv))

    days = 0
    uri = ''
    force = False

    for o, a in opts:
        if o == '--days':
            days = int(a)
        if o == '--uri':
            uri = a
        if o == '--force':
            STATES_TO_CLEANUP = ['wol_in_progress','wol_done','wol_failed','upload_in_progress','upload_done','upload_failed','execution_in_progress','execution_done','execution_failed','delete_in_progress','delete_done','delete_failed','inventory_in_progress','inventory_done','inventory_failed','reboot_in_progress','reboot_done','reboot_failed','halt_in_progress','halt_done','halt_failed','not_reachable','done','pause','paused','stop','stopped','scheduled','re_scheduled','failed','over_timed']

    if days == 0:
        sys.exit(usage(sys.argv))

    if uri == '':
        sys.exit(usage(sys.argv))

    mysql_db = create_engine(uri)
    metadata = MetaData(mysql_db)

    version = Table('version',
                   metadata,
                   autoload = True)
    msc_dbversion =  version.select('Number').execute().fetchone()[0]

    if msc_dbversion > 13:
        commands = Table('commands',
                         metadata,
                         Column('fk_bundle', Integer, ForeignKey('bundle.id')),
                         autoload = True)
    else:
        commands = Table('commands',
                         metadata,
                         Column('bundle_id', Integer, ForeignKey('bundle.id')),
                         autoload = True)

    bundle = Table('bundle',
                   metadata,
                   autoload = True)

    target = Table("target",
                   metadata,
                   autoload = True)
    commands_on_host = Table('commands_on_host',
                             metadata,
                             Column('fk_commands', Integer, ForeignKey('commands.id')),
                             Column('fk_target', Integer, ForeignKey('target.id')),
                             autoload = True)
    commands_history = Table("commands_history",
                             metadata,
                             Column('fk_commands_on_host', Integer, ForeignKey('commands_on_host.id')),
                             autoload = True)

    creation_timestamp =  datetime.datetime.now() - datetime.timedelta(days)
    cids = query_commands_to_delete(commands, commands_on_host, creation_timestamp)
    commands_to_delete = map(lambda x: x[0], cids)

    bcids = query_bundles_to_delete(commands, commands_on_host, creation_timestamp)
    commands_to_delete.extend(map(lambda x: x[1], bcids))
    tmp = map(lambda x: x[0], bcids)
    # Remove duplicates of bundle ids
    bundles_to_delete = []
    for bundle_id in tmp:
        if bundle_id not in bundles_to_delete:
            bundles_to_delete.append(bundle_id)

    logger.info("Number of commands to delete: %d", len(commands_to_delete))
    logger.info("Number of bundles to delete: %d", len(bundles_to_delete))

    if commands_to_delete:
        # Start transaction
        connection = mysql_db.connect()
        trans = connection.begin()
        msc_delete(connection, bundles_to_delete, commands_to_delete, bundle, commands, target, commands_on_host, commands_history)
        trans.commit()
        optimize_tables()
    else:
        logger.info("No purge to do")
