#!/usr/bin/env python

import optparse, os, sys
import datetime, time, re
import MetagenomeDB as mdb

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)
p.add_option("--version", dest = "display_version", action = "store_true", default = False)

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

connection_parameters = {}
def declare_connection_parameter (option, opt, value, parser):
	connection_parameters[opt[2:]] = value

g.add_option("--host", dest = "connection_host", metavar = "HOSTNAME",
	type = "string", action = "callback", callback = declare_connection_parameter,
	help = """Host name or IP address of the MongoDB server (optional). Default:
'host' property in ~/.MetagenomeDB, or 'localhost' if not found.""")

g.add_option("--port", dest = "connection_port", metavar = "INTEGER",
	type = "string", action = "callback", callback = declare_connection_parameter,
	help = """Port of the MongoDB server (optional). Default: 'port' property
in ~/.MetagenomeDB, or 27017 if not found.""")

g.add_option("--db", dest = "connection_db", metavar = "STRING",
	type = "string", action = "callback", callback = declare_connection_parameter,
	help = """Name of the database in the MongoDB server (optional). Default:
'db' property in ~/.MetagenomeDB, or 'MetagenomeDB' if not found.""")

g.add_option("--user", dest = "connection_user", metavar = "STRING",
	type = "string", action = "callback", callback = declare_connection_parameter,
	help = """User for the MongoDB server connection (optional). Default:
'user' property in ~/.MetagenomeDB, or none if not found.""")

g.add_option("--password", dest = "connection_password", metavar = "STRING",
	type = "string", action = "callback", callback = declare_connection_parameter,
	help = """Password for the MongoDB server connection (optional). Default:
'password' property in ~/.MetagenomeDB, or none if not found.""")

p.add_option_group(g)

(p, a) = p.parse_args()

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

if (p.display_version):
	print mdb.version
	sys.exit(0)

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)

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

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

try:
	mdb.connect(**connection_parameters)
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.title(), 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 "  objects of type %s: %s removed" % (collection_name.title(), 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 "  objects of type %s: %s removed" % (collection_name.title(), 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 (mdb.errors.DBConnectionError, mdb.errors.DBOperationError) as msg:
	error(msg)

print "Done."

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