import sys

from utils.helpers import *

cassandra_domain_tables = ["domain_activity", "domain_user", "domain_host", "domain_group", "domain_expired_user",
                           "domain_locked_user", "domain_last_seen", "domain_activity_event_head",
                           "domain_activity_event_head_meta", "domain_event"]
cassandra_asset_tables = ["asset_activity", "asset_shared_folder",
                          "asset_drive", "asset_monitored_folder", "asset_last_seen"]

cassandra_network_tables = ["network_alert", "network_file", "network_application", "network_website",
                            "network_mac_last_seen", "network_ip_last_seen"]
cassandra_resource_tables = ["resource_activity", "resource_last_seen"]
cassandra_all_tables = {
    "network": cassandra_network_tables,
    "domain": cassandra_domain_tables,
    "asset": cassandra_asset_tables,
    "resource": cassandra_resource_tables
}

memsql_domain_tables = ["domain_expired_user", "domain_group", "domain_host", "domain_last_seen",
                        "domain_locked_user", "domain_month_cache", "domain_month_last_seen_cache", "domain_month_timestamp_cache"]
memsql_network_tables = ["network_alert_month_last_seen_cache", "network_alert_month_timestamp_cache", "network_application_month_last_seen_cache", "network_application_month_timestamp_cache",
                         "network_file_month_timestamp_cache", "network_file_month_last_seen_cache",
                         "network_ip_last_seen", "network_mac_last_seen",
                         "network_website_month_timestamp_cache", "network_website_month_last_seen_cache"]
memsql_asset_tables = ['asset_drive', 'asset_last_seen', 'asset_monitored_folder', 'asset_shared_folder',
                       'asset_month_last_seen_cache', 'asset_month_timestamp_cache']

memsql_resource_tables = ['resource_last_seen',
                          'resource_month_timestamp_cache', 'resource_month_last_seen_cache']
memsql_all_tables = {
    'network': memsql_network_tables,
    'domain': memsql_domain_tables,
    'asset': memsql_asset_tables,
    'resource': memsql_resource_tables
}

DELETE_QUERY_FORMAT = lambda table, workspace_uuid, sensor_uuid: "DELETE FROM %s where workspace_uuid = '%s' and sensor_uuid = '%s'" % (
    table,
    workspace_uuid,
    sensor_uuid)
COUNT_QUERY_FORMAT = lambda table, workspace_uuid, sensor_uuid: "SELECT count(*) FROM %s where workspace_uuid = '%s' and sensor_uuid = '%s'" % (
    table,
    workspace_uuid,
    sensor_uuid)


def delete_cassandra(workspace_uuid, sensor_uuid):
    print "Cassandra cleaning..."
    session = create_cassandra_connection()
    for keyspace in cassandra_all_tables.keys():
        session.set_keyspace(keyspace)
        tables = cassandra_all_tables[keyspace]
        for table in tables:
            result = session.execute(COUNT_QUERY_FORMAT(
                table, workspace_uuid, sensor_uuid))
            row_counts = result.current_rows[0].count
            try:
                session.execute(DELETE_QUERY_FORMAT(
                    table, workspace_uuid, sensor_uuid))
                print "Deleted %s from %s.%s" % (row_counts, keyspace, table)
            except Exception as e:
                print "Error on deleting %s.%s" % (keyspace, table)
    session.shutdown()


def delete_memsql(workspace_uuid, sensor_uuid):
    print "MemSQL cleaning..."
    for keyspace in memsql_all_tables.keys():
        session = create_memsql_connection(db_name=keyspace)
        tables = memsql_all_tables[keyspace]
        for table in tables:
            result = session.query(COUNT_QUERY_FORMAT(
                table, workspace_uuid, sensor_uuid))
            row_counts = result[0].get("count(*)", 0)
            try:
                session.execute(DELETE_QUERY_FORMAT(
                    table, workspace_uuid, sensor_uuid))
                print "Deleted %s from %s.%s" % (row_counts,keyspace, table)
            except Exception as e:
                print "Error on deleting %s.%s" % (keyspace, table)

if __name__ == "__main__":
    workspace_uuid = sys.argv[1]
    sensor_uuid = sys.argv[2]
    delete_cassandra(workspace_uuid, sensor_uuid)
    delete_memsql(workspace_uuid, sensor_uuid)