#!/usr/bin/env python


#
# =============================================================================
#
#                                   Preamble
#
# =============================================================================
#


from optparse import OptionParser
import sys


from glue.ligolw import ligolw
from glue.ligolw import lsctables
from glue.ligolw import utils as ligolw_utils
from pylal import git_version


lsctables.use_in(ligolw.LIGOLWContentHandler)


__author__ = "Kipp Cannon <kipp.cannon@ligo.org>"
__version__ = "git id %s" % git_version.id
__date__ = git_version.date


#
# =============================================================================
#
#                                 Command Line
#
# =============================================================================
#


def parse_command_line():
	parser = OptionParser(
		version = "Name: %%prog\n%s" % git_version.verbose_msg,
		usage = "%prog [options] [file ...]",
		description = "Fixes XML files generated by ligolw_segments_query and friends so that they are compatible with other segment lists in XML format.  Currently, this means adding domain, jobid, and is_online to the process table and start_time_ns and end_time_ns columns to the segment table.  If XML files are given on the command line they are modified in-place, otherwise input is read from stdin and output is written to stdout."
	)
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
	options, filenames = parser.parse_args()

	return options, (filenames or [None])


#
# =============================================================================
#
#                                     Main
#
# =============================================================================
#


#
# parse command line
#


options, filenames = parse_command_line()


#
# iterate over files
#


for filename in filenames:
	#
	# load file
	#

	xmldoc = ligolw_utils.load_filename(filename, verbose = options.verbose, contenthandler = ligolw.LIGOLWContentHandler)

	#
	# add domain, jobid, and is_online columns to process table
	# add start_time_ns and end_time_ns to segment table
	#

	#
	# tables to fix
	#

	tables = {}
	tables["process"] = lsctables.ProcessTable.tableName
	tables["segment"] = lsctables.SegmentTable.tableName
	tables["segment_definer"] = lsctables.SegmentDefTable.tableName

	#
	# default values for new columns
	#

	values = {}
	values["process"] = {}
	values["process"]["domain"] = None
	values["process"]["jobid"] = None
	values["process"]["is_online"] = False
	values["segment"] = {}
	values["segment"]["start_time_ns"] = 0
	values["segment"]["end_time_ns"] = 0
	values["segment_definer"] = {}

	#
	# columns to remove
	#

	remove = {}
	remove["process"] = []
	remove["segment"] = ["creator_db", "segment_def_cdb"]
	remove["segment_definer"] = ["creator_db"]

	#
	# loop over tables
	#

	for table_name in tables:
		#
		# retrieve table
		#

		try:
			table = lsctables.table.get_table(xmldoc, tables[table_name])
		except ValueError:
			#
			# no table --> no op
			#

			continue

		#
		# loop over columns to add
		#

		for column_name in values[table_name]:
			try:
				column = table.appendColumn(column_name)
			except ValueError:
				#
				# table already has this column
				#

				continue

			#
			# fill with appropriate values
			#

			for i in range(len(table)):
				column[i] = values[table_name][column_name]

		#
		# loop over columns to remove
		#

		for column_name in remove[table_name]:
			try:
				table.removeChild(table.getColumnByName(column_name))
			except KeyError:
				#
				# column is already missing
				#

				pass

	#
	# write file
	#

	ligolw_utils.write_filename(xmldoc, filename, gz = (filename or "stdout").endswith(".gz"), verbose = options.verbose)
