#!/usr/bin/env python

import optparse, os, sys
import datetime, time, re

p = optparse.OptionParser(description = """Part of the MetagenomeDB toolkit.
Remove objects from the database.""")

g = optparse.OptionGroup(p, "Target")

g.add_option("--all", dest = "remove_all", action = "store_true", default = False,
	help = "Remove all objects.")

g.add_option("--type", dest = "remove_type", metavar = "STRING",
	help = "Remove all objects of a given type.")

g.add_option("--since", dest = "remove_since", metavar = "TIMESTAMP",
	help = """Remove all objects created since TIMESTAMP. Timestamp is expressed
as a number of seconds, minutes, hours, days and/or weeks before the time
mdb-remove is executed. Format is the following: [0-9][wdhms]; several arguments
can be concatenated. E.g., 2h30m will remove all objects created in the last 2
hours and 30 minutes.""")

g.add_option("--between", dest = "remove_between", nargs = 2, metavar = "TIMESTAMP1 TIMESTAMP2",
	help = """Remove all objects created between (and including) TIMESTAMP1 and
TIMESTAMP2. One of the two timestamps can be 'now' to indicate the current time.""")

p.add_option_group(g)

p.add_option("-v", "--verbose", dest = "verbose", action = "store_true", default = False)
p.add_option("--dry-run", dest = "dry_run", action = "store_true", default = False)

g = optparse.OptionGroup(p, "Connection")

g.add_option("--host", dest = "connection_host", metavar = "HOSTNAME", default = "localhost",
	help = "Host name or IP address of the MongoDB server (optional). Default: %default")

g.add_option("--port", dest = "connection_port", metavar = "INTEGER", default = 27017,
	help = "Port of the MongoDB server (optional). Default: %default")

g.add_option("--db", dest = "connection_db", metavar = "STRING", default = "MetagenomeDB",
	help = "Name of the database in the MongoDB server (optional). Default: '%default'")

g.add_option("--user", dest = "connection_user", metavar = "STRING", default = '',
	help = "User for the MongoDB server connection (optional). Default: '%default'")

g.add_option("--password", dest = "connection_password", metavar = "STRING", default = '',
	help = "Password for the MongoDB server connection (optional). Default: '%default'")

p.add_option_group(g)

(p, a) = p.parse_args()

def error (msg):
	if str(msg).endswith('.'):
		msg = str(msg[:-1])
	print >>sys.stderr, "ERROR: %s." % msg
	sys.exit(1)

TIMESTAMP = re.compile("([0-9]+)([ywdhms])")

def convert_timestamp (timestamp):
	if (timestamp.lower() == "now"):
		return datetime.datetime.utcnow()

	m = list(TIMESTAMP.finditer(timestamp.lower()))
	if (len(m) == 0):
		raise Exception("Invalid timestamp '%s'" % timestamp)

	delta = {}
	for item in m:
		value, key = item.groups()
		key = {
			"w": "weeks",
			"d": "days",
			"h": "hours",
			"m": "minutes",
			"s": "seconds",
		}[key]
		delta[key] = int(value)

	timestamp = datetime.datetime.utcnow() - datetime.timedelta(**delta)
	return timestamp

if (p.remove_all) and (p.remove_type or p.remove_since or p.remove_between):
	error("--all cannot be used with other objects selectors")

if (p.remove_type) and (p.remove_since or p.remove_between):
	error("--type cannot be used with --since or --between")

if (p.remove_since and p.remove_between):
	error("--since and --between cannot be used simultaneously")

try:
	if (p.remove_since != None):
		p.remove_since = convert_timestamp(p.remove_since)

	if (p.remove_between):
		a, b = p.remove_between
		p.remove_between = sorted([convert_timestamp(a), convert_timestamp(b)])

except Exception as msg:
	error(msg)

#:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

import MetagenomeDB as mdb

if (p.verbose):
	mdb.max_verbosity()

if (p.connection_host or p.connection_port or p.connection_db or p.connection_user or p.connection_password):
	try:
		mdb.connect(p.connection_host, p.connection_port, p.connection_db, p.connection_user, p.connection_password)
	except Exception as msg:
		error(msg)

def plural (n, text):
	if (n > 1):
		return text % n + "s"
	else:
		return text % n

try:
	if (p.remove_all):
		for collection_name in mdb.backend.list_collections():
			print "Removing all objects of type '%s' (%s)" % (collection_name.lower(), plural(mdb.backend.count(collection_name, {}), "%s object"))

			if (not p.dry_run):
				mdb.backend.drop_collection(collection_name)

	elif (p.remove_type):
		mapping = {}
		for collection_name in mdb.backend.list_collections():
			mapping[collection_name.lower()] = collection_name

		p.remove_type = p.remove_type.lower()
		if (not p.remove_type in mapping):
			raise Exception("Unknown object type '%s'" % p.remove_type)

		collection_name = mapping[p.remove_type]
		print "Removing all objects of type '%s' (%s)" % (p.remove_type, plural(mdb.backend.count(collection_name, {}), "%s object"))

		if (not p.dry_run):
			mdb.backend.drop_collection(collection_name)

	elif (p.remove_since):
		timestamp = p.remove_since
		query = {"_creation_time": {"$gte": timestamp}}
		timestamp -= datetime.timedelta(seconds = time.timezone)

		print "Removing all objects created since %s" % timestamp.strftime("%c")

		for (collection_name, collection) in mdb.backend.list_collections(True):
			print "  %s: %s" % (collection_name.lower(), plural(collection.count(query), "%s object"))

			for object in collection.find(query):
				if (p.dry_run):
					print "    %s" % object
				else:
					mdb.backend.remove_object(object)

	elif (p.remove_between):
		timestamp1, timestamp2 = p.remove_between
		query = {"_creation_time": {"$gte": timestamp1, "$lte": timestamp2}}
		timestamp1 -= datetime.timedelta(seconds = time.timezone)
		timestamp2 -= datetime.timedelta(seconds = time.timezone)

		print "Removing all objects created between %s and %s" % (
			timestamp1.strftime("%c"),
			timestamp2.strftime("%c")
		)

		for (collection_name, collection) in mdb.backend.list_collections(True):
			print "  %s: %s" % (collection_name.lower(), plural(collection.count(query), "%s object"))

			for object in collection.find(query):
				if (p.dry_run):
					print "    %s" % object
				else:
					mdb.backend.remove_object(object)

except Exception as msg:
	error(msg)

print "Done."

if (p.dry_run):
	print "(dry run)"
