#!/usr/bin/env python

import optparse, sys, os, pprint
import MetagenomeDB as mdb

p = optparse.OptionParser(description = """Part of the MetagenomeDB toolkit.
List collections and sequences in the database.""")

p.add_option("-o", "--output", dest = "output_fn", metavar = "FILENAME",
	help = "Destination for the output (optional). Default: standard output.")

p.add_option("-p", "--with-properties", dest = "display_properties", action = "store_true", default = False,
	help = """If set, collections' properties will be displayed as well.""")

p.add_option("-v", "--verbose", dest = "verbose", 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):
	if str(msg).endswith('.'):
		msg = str(msg)[:-1]
	print >>sys.stderr, "ERROR: %s." % msg
	sys.exit(1)

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

if (p.output_fn == None):
	output_fh = sys.stdout
else:
	output_fh = open(p.output_fn, 'w')

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

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

try:
	mdb.connect(**connection_parameters)
except Exception as msg:
	error(msg)

def collection_title (collection):
	n = collection.count_sequences()
	title = collection["name"]
	if (n > 1):
		title += " (%s sequences)" % n
	elif (n > 0):
		title += " (%s sequence)" % n
	return title

def traverse (map, prefix = None):
	for key in sorted(map.keys()):
		if (key.startswith('_')):
			continue

		if (prefix == None):
			prefix_ = [key]
		else:
			prefix_ = prefix + [key]

		if (type(map[key]) == dict):
			for item in traverse(map[key], prefix_):
				yield item
		else:
			yield (prefix_, map[key])

sorted_ = lambda x: sorted(x, lambda x, y: cmp(x["name"], y["name"]))

try:
	for collection in sorted_(mdb.Collection.find()):
		print >>sys.stdout, collection_title(collection)

		if (p.display_properties):
			for (key, value) in traverse(collection.get_properties()):
				if (type(value) == unicode):
					value = "'%s'" % value
				print "	%s = %s" % ('.'.join(key), value)

		for supercollection in sorted_(collection.list_super_collections()):
			print >>sys.stdout, "	(super-collection: '%s')" % supercollection["name"]

		for subcollection in sorted_(collection.list_sub_collections()):
			print >>sys.stdout, "	(sub-collection: '%s')" % subcollection["name"]

		print >>sys.stdout

except mdb.errors.DBConnectionError as msg:
	error(msg)
