#!/usr/bin/env python

import os
import re
import sys
import logging
import optparse
import commands
import subprocess
import shutil
from datetime import timedelta
import traceback

# Initialize logging object
logger = logging.getLogger()
# Set default level to INFO
logger.setLevel(logging.INFO)

# Use pegasus-config to find our Pegasus environment
bin_dir = os.path.normpath(os.path.join(os.path.dirname(sys.argv[0])))
pegasus_config = os.path.join(bin_dir, "pegasus-config") + " --python-dump"
config = subprocess.Popen(pegasus_config, stdout=subprocess.PIPE, shell=True).communicate()[0]
exec config

# We also need the externals lib 
pegasus_config = os.path.join(bin_dir, "pegasus-config") + " --noeoln --python-externals"
lib_ext_dir = subprocess.Popen(pegasus_config, stdout=subprocess.PIPE, shell=True).communicate()[0]

# Insert the lib directories in our search path
os.sys.path.insert(0, lib_ext_dir)

import Pegasus.common
from Pegasus.plots_stats import utils as plot_utils
from Pegasus.plots_stats.plots import populate
from Pegasus.plots_stats.plots import pegasus_gantt
from Pegasus.plots_stats.plots import pegasus_host_over_time
from Pegasus.plots_stats.plots import pegasus_breakdown
from Pegasus.plots_stats.plots import pegasus_time
from Pegasus.tools import utils


# Regular expressions
re_parse_property = re.compile(r'([^:= \t]+)\s*[:=]?\s*(.*)')

# Global variables----
prog_base = os.path.split(sys.argv[0])[1]	# Name of this program
NEW_LINE_STR ="\n"
MAX_GRAPH_NODES = 100
calc_dax_graph = False
calc_dag_graph = False
calc_gantt_chart = False
calc_host_chart = False
calc_time_chart = False
calc_breakdown_chart = False
max_graph_nodes = MAX_GRAPH_NODES
DEFAULT_OUTPUT_DIR = "plots"
pegasus_env_path = {
'pegasus_bin_dir'       : pegasus_bin_dir,
'pegasus_conf_dir'      : pegasus_conf_dir,
'pegasus_java_dir'      : pegasus_java_dir,
'pegasus_perl_dir'      : pegasus_perl_dir,
'pegasus_python_dir'    : pegasus_python_dir,
'pegasus_php_dir'       : pegasus_php_dir,
'pegasus_javascript_dir': pegasus_javascript_dir,
'pegasus_share_dir'     : pegasus_share_dir
}
	
def setup_logger(level_str):
	"""
	Sets the logging level
	@param level_str logging level
	"""
	level_str = level_str.lower()
	if level_str == "debug":
		logger.setLevel(logging.DEBUG)
	if level_str == "warning":
		logger.setLevel(logging.WARNING)
	if level_str == "error":
		logger.setLevel(logging.ERROR)
	if level_str == "info":
		logger.setLevel(logging.INFO)
	populate.setup_logger(level_str)
	return

def setup(output_dir):
	"""
	Set up the pegaus plots. 
	@param output_dir output directory path
	"""
	src_img_path = os.path.join(pegasus_share_dir, "plots/images/common/")
	dest_img_path = os.path.join(output_dir, "images/")
	if os.path.isdir(dest_img_path):
		logger.warning("Image directory exists. Deleting... " + dest_img_path)
		try:
			shutil.rmtree(dest_img_path)
		except:
			logger.error("Unable to remove image directory." + dest_img_path)
			sys.exit(1) 	 	
        # Need to create the path in case it doesn't exist yet when running under Python < 2.5
        if sys.version_info < (2, 5):
                if not os.path.isdir(output_dir):
                        os.makedirs(output_dir)
	shutil.copytree (src_img_path, dest_img_path)	

def generate_dag_graph(wf_info, output_dir):
	"""
	Generates the DAG graph
	@wf_info WorkflowInfo object reference
	@output_dir the output directory path
	"""
	dag_file_path = wf_info.dag_file_path
	title = str(wf_info.wf_uuid) + " (" + str(wf_info.dax_label) +")"
	if dag_file_path is not None:
		dag2dot_file_path = os.path.join(pegasus_share_dir, "visualize/dag2dot")
		dot_file_path = os.path.join(output_dir, wf_info.wf_uuid + ".dot")
		dag_cmd = dag2dot_file_path 
		dag_cmd += " --output " + dot_file_path
		dag_cmd += " " + dag_file_path
		logger.debug("Executing command :\n" + dag_cmd)
		status, output = commands.getstatusoutput(dag_cmd)
		if status == 0:
			logger.info("Finished executing command." )
		else:
			logger.warn("Failed to generate dag graph for workflow " + title)
			logger.debug("%s: %d:%s" % (dag_cmd, status, output))
			return None
		png_file_path = os.path.join(output_dir, wf_info.wf_uuid + ".png")
		dot_png_cmd = utils.find_exec("dot")
		if dot_png_cmd is None:
			logger.warn("dot is not present . Unable to create chart in png format. ")
			return None
		dot_png_cmd += " -Tpng -o" + png_file_path
		dot_png_cmd += " " + dot_file_path
		logger.debug("Executing command :\n" + dot_png_cmd)
		status, output = commands.getstatusoutput(dot_png_cmd)
		if status == 0:
			logger.debug("Finished executing command." )
			return status
		else:
			logger.warn("%s: %d:%s" % (dot_png_cmd, status, output))
	else:
		logger.warn("Unable to find the dag file for workflow  " + title)
	return None
	
	
def generate_dax_graph(wf_info, output_dir):
	"""
	Generates the DAX graph
	@wf_info WorkflowInfo object reference
	@output_dir the output directory path
	"""
	dax_file_path = wf_info.dax_file_path
	title =  str(wf_info.wf_uuid) + " (" + str(wf_info.dax_label) +")"
	if dax_file_path is not None:
		dax2dot_file_path = os.path.join(pegasus_share_dir, "visualize/dax2dot")
		dot_file_path = os.path.join(output_dir, wf_info.wf_uuid + ".dot")
		dax_cmd = dax2dot_file_path 
		dax_cmd +=" --output "+ dot_file_path
		dax_cmd += " "+ dax_file_path
		logger.debug("Executing command :\n"  + dax_cmd)
		status, output = commands.getstatusoutput(dax_cmd)
		if status == 0:
			logger.debug("Finished executing command." )
		else:
			logger.warn("Failed to generate dax graph for workflow " + title)
			logger.debug("%s: %d:%s" % (dax_cmd, status, output))
			return None
		# Find dot command
		dot_png_cmd = utils.find_exec("dot")
		if dot_png_cmd is None:
			logger.warn("dot is not present . Unable to create chart in png format. ")
			return None
		png_file_path = os.path.join(output_dir, wf_info.wf_uuid + ".png")
		dot_png_cmd +=" -Tpng -o" + png_file_path
		dot_png_cmd += " "+ dot_file_path
		logger.debug("Executing command :\n"  + dot_png_cmd)
		status, output = commands.getstatusoutput(dot_png_cmd)
		if status == 0:
			logger.debug("Finished executing command." )
			return status
		else:
			logger.warn("Failed to generate dax graph in png format for workflow " + wf_info.wf_uuid)
			logger.debug("%s: %d:%s" % (dot_png_cmd, status, output))
	else:
		logger.warn("Unable to find the dax file for workflow " + title)
	return None
	
	
def create_image_gallery(file_name, wf_uuid_list, wf_parent_uuid_list, uuid_image_map, wf_uuid_label_map, isDax):
	"""
	Creates the image gallery.
	@param file_mame the output file name
	@param wf_uuid_list the list of all workflow id's
	@param wf_parent_uuid_list the list of parent workflow id's corresponding to each workflow
	@param uuid_image_map uuid and image file mapping
	@param wf_uuid_label_map uuid and label mapping
	@isDax boolean indicating whether it is a dax file or not.
	"""
	wf_uuid_parent_ref = None
	try:
		fh = open(file_name, "w")
		content = """
<html>
<head>
<style>
.imgbox
{
float:left;
text-align:center;
width:450px;
height:450px;
margin:4px;
margin-bottom:8px;
padding:0px;
}
.thumbnail
{
width:300px;
height:300px;
margin:3px;
}
.box
{
width:450px;
padding:0px;
}
.workflow
{
clear:both;
}
</style>
</head>
<body>
""" + plot_utils.create_home_button() + """
<center>
"""
		if isDax:
			content += "<h3>DAX Graph </h3>" + NEW_LINE_STR
		else:
			content += "<h3>DAG Graph </h3>"+ NEW_LINE_STR
		for index in range(len(wf_uuid_list)):
			uuid = wf_uuid_list[index]
			image = uuid_image_map[index]
			label = wf_uuid_label_map[index]
			parent_uuid =wf_parent_uuid_list[index]
			if parent_uuid is None:
				content += "<h3 class= 'workflow'> Top level workflow (" + uuid + ")</h3>"
			else:
				if parent_uuid != wf_uuid_parent_ref:
					wf_uuid_parent_ref = parent_uuid 
					content += "<h3 class= 'workflow'> Sub  workflow's of workflow (" +  parent_uuid + ")</h3>"
			content += "<div  class ='imgbox' >"
			if image is None:
				content += "<a class= 'thumbnail' href ='#'>\n"
				content +="<img src ='images/not_available.jpg' height='300px' width='300px'>\n"
				content +="</img>\n</a>"
				content +="<div class ='box'>\n"
				content += "wf_uuid :" + uuid +"<br/>"
				if isDax:
					content+= "dax label :" + label
				else:
					if image is not None:
						content += "dag label :" + image
				content +="</div>"					
			else:
				content +="<a class= 'thumbnail'  href ='" + image + "'>"
				content +="<img src ='" + image + "' height='300px' width='300px'>"
				content +="</img>\n</a>\n"
				content +="<div class ='box'>\n"
				content += "wf_uuid :" + uuid +"<br/>"
				if isDax:
					content += "dax label :" + label
				else:
					if label is not None:
						content += "dag label :" + label
				content += "</div>"
			content += "</div>"
		content += """
</center>
</body>
</html>
"""
		fh.write( content)
	except IOError:
		logger.error("Unable to write to file " + data_file)
		sys.exit(1)
	else:
		fh.close()	
	
		



def createOuterhtml(wf_stats, wf_info, submit_dir, output_dir, log_level):
	"""
	Generates the outer html file which links to all the charts and graph
	@param wf_info WorkflowInfo object reference
	@param wf_stats StampedeStatistics reference
	@param submit_dir submit directory
	@param output_dir the output directory path
	@param log_level logging level
	"""
	wf_uuid = wf_info.wf_uuid
	title =  str(wf_info.wf_uuid) + " (" + str(wf_info.dax_label) + ")"
	data_file = os.path.join(output_dir, "index.html")
	gantt_chart_parent_file = os.path.join("gantt_chart/" + wf_uuid + ".html")
	dag_graph_parent_file = os.path.join("dag_graph/" + wf_uuid + ".html")
	dax_graph_parent_file = os.path.join("dax_graph/" + wf_uuid + ".html")
	host_chart_parent_file = os.path.join("host_chart/" + wf_uuid + ".html")
	breakdown_chart_parent_file = os.path.join("breakdown_chart/" + wf_uuid + ".html")
	
	time_summary_output_dir = None
	breakdown_summary_output_dir = None
	if calc_time_chart:
		time_summary_output_dir = output_dir
		pegasus_time.setup(submit_dir, time_summary_output_dir, pegasus_env_path, log_level)
	if  calc_breakdown_chart:
		breakdown_summary_output_dir = output_dir
		pegasus_breakdown.setup(submit_dir, breakdown_summary_output_dir, pegasus_env_path, log_level)
	try:
		fh = open(data_file, "w")
		content = """
<html>
<head>
<style type ='text/css'>
#breakdown_chart{
border:2px solid orange;
}
#breakdown_chart_footer_div{
border:2px solid  #C35617;
border-top-style:none;
}
#time_chart{
border:2px solid orange;
}
#time_chart_footer_div{
border:2px solid #C35617;
border-top-style:none;
}
#breakdown_chart_legend_div{
color:#0066CC;
}
.header_level1{
font-family:"Times New Roman", Times, serif; 
font-size:36px;
}
.header_level2{
font-family:"Times New Roman", Times, serif; 
font-size:30px;
padding-top:25px;
}
</style>

</head>
<body>
<center>
<div class ='header_level1'>Pegasus plots </div>
"""
		if calc_time_chart or calc_breakdown_chart:
			content += """
<script type='text/javascript' src='js/protovis-r3.2.js'>
</script>
			"""
		if calc_gantt_chart :
			content += "<a href ='" + gantt_chart_parent_file +  "'>Workflow Execution Gantt Chart (Per workflow)</a><br/>" + NEW_LINE_STR
		if calc_host_chart :
			content += "<a href ='" + host_chart_parent_file + "'>Host Over Time Chart (Per workflow)</a><br/>" + NEW_LINE_STR
		if calc_breakdown_chart:
			content += """<a href ='#breakdown_chart_div'>Invocation Breakdown Chart (Across workflows)</a><br/>""" + NEW_LINE_STR
			content += "<a href ='" + breakdown_chart_parent_file + "'>Invocation Breakdown Chart (Per workflow)</a><br/>" + NEW_LINE_STR
		if calc_time_chart :
			content += """<a href ="#time_chart_div">Time Chart (Across workflows)</a><br/>""" + NEW_LINE_STR
		if calc_dax_graph :
			content += "<a href ='" + dax_graph_parent_file + "'>DAX graph</a><br/>" + NEW_LINE_STR
		if calc_dag_graph :
			content += "<a href ='" + dag_graph_parent_file + "'>DAG graph</a><br/>" + NEW_LINE_STR
		content += """<a href ='#env_div'> Workflow environment</a><br/><br/><br/>""" + NEW_LINE_STR
		content += """<div id ='env_div' class ='header_level2'> Workflow environment </div>"""
		content += plot_utils.print_property_table(wf_info.wf_env, False, ":")
		if calc_breakdown_chart:
			content += """
<div id ='breakdown_chart_div' class ='header_level2'> Invocation breakdown chart (Across workflows) </div>
			"""
			populate.populate_transformation_details(wf_stats, wf_info)
			logger.debug("Generating breakdown chart for the workflow " + title + " ... ")
			content += pegasus_breakdown.create_breakdown_plot(wf_info, breakdown_summary_output_dir)
		if calc_time_chart :
			populate.populate_time_details(wf_stats, wf_info)
			logger.debug("Generating time chart for the workflow " + title + " ... ")
			content += """
<div id ='time_chart_div' class ='header_level2'> Time chart (Across workflows) </div>
			"""
			content += pegasus_time.create_time_plot(wf_info, time_summary_output_dir)
		
		content += """
</center>
</body>
</html>
"""
		fh.write(content)
		plots_output = NEW_LINE_STR + "SUMMARY".center(100, '*')
		plots_output +=  NEW_LINE_STR
		plots_output +=  NEW_LINE_STR
		plots_output += "Graphs and charts generated by pegasus-plots can be viewed by opening the generated html file in the web browser  : \n" + data_file
		plots_output += NEW_LINE_STR
		plots_output +=  NEW_LINE_STR
		plots_output += "".center(100, '*')
		print plots_output
	except IOError:
		logger.error("Unable to write to file " + data_file)
		sys.exit(1)
	else:
		fh.close()
		
def create_charts(submit_dir, output_dir, config_properties, log_level):
	"""
	Generates all the graphs and charts 
	@submit_dir submit directory pathd
	@output_dir the output directory path
	@config_properties path to the pegasus property file
	@log_level logging level
	"""
	wf_uuid_dax_image = []
	wf_uuid_dax_label = []
	wf_uuid_dag_image = []
	wf_uuid_dag_label = []
	wf_uuid_parent = []
	populate.setup(submit_dir, config_properties)
	wf_uuid_list = populate.get_workflows_uuid()
	if len(wf_uuid_list) == 0:
		logger.error("Unable to populate workflow information.")
	if calc_gantt_chart:
		gantt_chart_output_dir = os.path.join(output_dir, "gantt_chart")
		pegasus_gantt.setup(submit_dir, gantt_chart_output_dir, pegasus_env_path, log_level)
	if calc_host_chart:
		host_chart_output_dir = os.path.join(output_dir, "host_chart")
		pegasus_host_over_time.setup(submit_dir, host_chart_output_dir, pegasus_env_path, log_level)
	if calc_breakdown_chart:
		breakdown_chart_output_dir = os.path.join(output_dir, "breakdown_chart")
		pegasus_breakdown.setup(submit_dir, breakdown_chart_output_dir, pegasus_env_path, log_level)
		
	if calc_dag_graph:
		dag_graph_output_dir = os.path.join(output_dir, "dag_graph")
		setup(dag_graph_output_dir)
	if calc_dax_graph :
		dax_graph_output_dir = os.path.join(output_dir, "dax_graph")
		setup(dax_graph_output_dir)
	for wf_uuid in wf_uuid_list:
		wf_stats, wf_info = populate.populate_chart(wf_uuid)
		title =  str(wf_uuid) + " (" + str(wf_info.dax_label) +")"
		logger.info("Generating graphs/charts for the workflow " + title +" ... ")
		if calc_gantt_chart or calc_host_chart:
			populate.populate_job_instance_details(wf_stats, wf_info)
		wf_uuid_parent.append(wf_info.parent_wf_uuid)
		if calc_gantt_chart :
			logger.debug("Generating gantt chart for the workflow " + title +" ... ")
			pegasus_gantt.generate_chart(wf_info)
		if calc_host_chart :
			logger.debug("Generating host chart for the workflow " + title +" ... ")
			pegasus_host_over_time.generate_chart(wf_info)
		if calc_breakdown_chart:
			populate.populate_transformation_details(wf_stats, wf_info)
			logger.debug("Generating breakdown chart for the workflow " + title +" ... ")
			pegasus_breakdown.generate_chart(wf_info)
		
		if calc_dag_graph :
			populate.populate_job_details(wf_stats, wf_info)
			if wf_info.total_jobs <= max_graph_nodes:
				logger.debug("Generating dag graph for the workflow " + title +" ... ")
				if generate_dag_graph(wf_info, dag_graph_output_dir) is None:
					wf_uuid_dag_image.append(None)
				else:
					wf_uuid_dag_image.append(wf_info.wf_uuid+".png")
			else:
				wf_uuid_dag_image.append(None)
				logger.info("Dag graph for workflow '" + title + "'  was not created because number of jobs is more than the maximum graph node generation limit ")
			wf_uuid_dag_label.append( wf_info.dag_label)
		if calc_dax_graph :
			populate.populate_task_details(wf_stats, wf_info)
			if wf_info.total_tasks <= max_graph_nodes:
				logger.debug("Generating dax graph for the workflow " + title +" ... ")
				if generate_dax_graph(wf_info, dax_graph_output_dir) is None:
					wf_uuid_dax_image.append(None)
				else:
					wf_uuid_dax_image.append(wf_info.wf_uuid+".png")
			else:
				wf_uuid_dax_image.append(None)
				logger.info("Dax graph for workflow '" + title +"' was not created because number of tasks is more than the maximum  graph node generation limit ")
			wf_uuid_dax_label.append(wf_info.dax_label)
		wf_stats.close()
		wf_stats = None
		wf_info = None
	
	root_wf_stats, root_wf_info = populate.populate_chart(wf_uuid_list[0], True)
	if calc_dax_graph :
		data_file =  os.path.join(dax_graph_output_dir, wf_uuid_list[0] + ".html")
		create_image_gallery(data_file, wf_uuid_list, wf_uuid_parent, wf_uuid_dax_image, wf_uuid_dax_label, True )
	if calc_dag_graph :
		data_file =  os.path.join(dag_graph_output_dir, wf_uuid_list[0] + ".html")
		create_image_gallery(data_file, wf_uuid_list, wf_uuid_parent, wf_uuid_dag_image, wf_uuid_dag_label, False )
	
	createOuterhtml(root_wf_stats, root_wf_info, submit_dir, output_dir, log_level)
	root_wf_stats.close()
	return
	

def set_plotting_level(plots_level):
	"""
	Sets the plotting level 
	@param plot_level
	"""
	global calc_dax_graph
	global calc_dag_graph
	global calc_gantt_chart
	global calc_host_chart
	global calc_time_chart
	global calc_breakdown_chart
	if plots_level =='all':
		calc_dax_graph = True
		calc_dag_graph = True
		calc_gantt_chart = True
		calc_host_chart = True
		calc_time_chart = True
		calc_breakdown_chart = True
	elif plots_level =='all_charts':
		calc_gantt_chart = True
		calc_host_chart = True
		calc_time_chart = True
		calc_breakdown_chart = True
	elif plots_level =='all_graphs':
		calc_dax_graph = True
		calc_dag_graph = True
	elif plots_level == 'dax_graph':
		calc_dax_graph = True
	elif plots_level == 'dag_graph':
		calc_dag_graph = True
	elif plots_level == 'gantt_chart':
		calc_gantt_chart = True
	elif plots_level == 'host_chart':
		calc_host_chart = True
	elif plots_level == 'time_chart':
		calc_time_chart = True
	else:
		calc_breakdown_chart = True	
	
	

# ---------main----------------------------------------------------------------------------
def main():
	# Configure command line option parser
	prog_usage = prog_base +" [options] SUBMIT DIRECTORY" 
	parser = optparse.OptionParser(usage=prog_usage)
	parser.add_option("-o", "--output", action = "store", dest = "output_dir",
			help = "writes the output to given directory.")
	parser.add_option("-c", "--conf", action = "store", type = "string", dest = "config_properties",
			  help = "specifies the properties file to use. This option overrides all other property files.")
	parser.add_option("-m", "--max-graph-nodes", action = "store", type = "int", dest = "max_graph_nodes",
		  help = "Maximum limit on the number of tasks/jobs in the dax/dag upto which the graph should be generated; Default value is 100")
	parser.add_option("-p", "--plotting-level", action = "store", dest = "plotting_level",
			choices=['all', 'all_charts', 'all_graphs', 'dax_graph', 'dag_graph', 'gantt_chart', 'host_chart', 'time_chart', 'breakdown_chart'],
			help = "specifies the chart and graphs to generate. Valid levels are: all, all_charts, all_graphs, dax_graph,\
					dag_graph, gantt_chart, host_chart, time_chart, breakdown_chart; Default is all_charts.")
	parser.add_option("-i", "--ignore-db-inconsistency", action = "store_const", const = 0, dest = "ignore_db_inconsistency",
			help = "turn off the check for db consistency")
	parser.add_option("-v", "--verbose", action="count", default=0, dest="verbose", 
			help="Increase verbosity, repeatable")
	parser.add_option("-q", "--quiet", action="count", default=0, dest="quiet", 
			help="Decrease verbosity, repeatable")
	# Parse command line options
	(options, args) = parser.parse_args()
	
	if len(args) > 1:
		parser.error("Invalid argument")
		sys.exit(1) 
	
	if len(args) < 1:
                submit_dir = os.getcwd()	
	else:
		submit_dir = os.path.abspath(args[0])
	
        # Copy options from the command line parser
	log_level = 1
	log_level_str = "info" 
	log_level += (options.verbose - options.quiet)
	if log_level <= 0:
		log_level_str = "error"
	elif log_level == 1:
		log_level_str = "warning"
	elif log_level == 2:
		log_level_str = "info"
	elif log_level >= 3:
		log_level_str = "debug"
	setup_logger(log_level_str)
	logger.info(prog_base + " : initializing...")
	
	if options.ignore_db_inconsistency is None:
		if not utils.loading_completed(submit_dir):
			if utils.monitoring_running(submit_dir):
				logger.warning("pegasus-monitord still running. Please wait for it to complete. ")
			else:
				logger.warning("Please run pegasus monitord in replay mode. ")
			sys.exit(1)
	else:
		logger.warning("The tool is meant to be run after the completion of workflow run.")
		
	if options.plotting_level is not None:
		 plotting_level = options.plotting_level
	else:
		plotting_level =  'all_charts'
	set_plotting_level(plotting_level)
		
	if options.output_dir is not None:
		output_dir = options.output_dir
	else :
		output_dir = os.path.join(submit_dir, DEFAULT_OUTPUT_DIR)
	utils.create_directory(output_dir, True)
	if options.max_graph_nodes is not None:
		global max_graph_nodes
		max_graph_nodes = options.max_graph_nodes
	
	try:
		create_charts(submit_dir, output_dir, options.config_properties, log_level_str)
	except SystemExit:
		sys.exit(1)
	except:
		logger.warning(traceback.format_exc())
		sys.exit(1)

	sys.exit(0)

if __name__ == '__main__':
	main()
