#!/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


#
# =============================================================================
#
#                                Configuration
#
# =============================================================================
#


#
# values for new columns
#


columns_to_add = {
	lsctables.ProcessTable: {"domain": None, "jobid": None, "is_online": None},
	lsctables.SegmentTable: {"start_time_ns": 0, "end_time_ns": 0},
}


#
# columns to remove
#


columns_to_remove = {
	lsctables.SegmentTable: ["creator_db", "segment_def_cdb"],
	lsctables.SegmentDefTable: ["creator_db"]
}


#
# =============================================================================
#
#                                 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)

	#
	# loop over tables.  retrieve each one, skip ones that are missing,
	# then loop over columns
	#

	for table_type, columns in columns_to_add.items():
		try:
			table = table_type.get_table(xmldoc)
		except ValueError:
			# no table --> no op
			continue

		for column_name, val in columns.items():
			try:
				# add column
				col = table.appendColumn(column_name)
			except ValueError:
				# already there --> no op
				continue
			# fill
			col[:] = (val,) * len(table)

	for table_type, columns in columns_to_remove.items():
		try:
			table = table_type.get_table(xmldoc)
		except ValueError:
			# no table --> no op
			continue

		for column_name in columns:
			try:
				# remove column
				table.removeChild(table.getColumnByName(column_name))
			except KeyError:
				# not there --> no op
				pass

	#
	# write file
	#

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