import os, sys
import profile

# Nice labels for algorithms of all kinds
ALG_LABEL = {
	"AES_128_GCM_SHA256": "AES128",
	"AES_256_GCM_SHA384": "AES256",
	"CHACHA20_POLY1305_SHA256": "CHACHA20",
	"ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,ECDHE_RSA_WITH_AES_128_GCM_SHA256": "AES128",
	"ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,ECDHE_RSA_WITH_AES_256_GCM_SHA384": "AES256",
	"ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256,ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256": "CHACHA20",
	"prime256v1": "prime256v1",
	"rsa2048": "rsa2048",
	"rsa3072": "rsa3072",
	"rsa4096": "rsa4096",
	"X25519": "X25519",
	"SECP256R1": "SECP256R1",
	"SECP384R1": "SECP384R1"
}

# Nice labels for TLS versions using ciphers
VER_LABEL = {
	"AES_128_GCM_SHA256": "1.3",
	"AES_256_GCM_SHA384": "1.3",
	"CHACHA20_POLY1305_SHA256": "1.3",
	"ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,ECDHE_RSA_WITH_AES_128_GCM_SHA256": "1.2",
	"ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,ECDHE_RSA_WITH_AES_256_GCM_SHA384": "1.2",
	"ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256,ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256": "1.2"
}

# Titles for measured quantities
OBJ_TITLE = {
	"cpu": "CPU time",
	"energy": "energy consumption",
	"profile": "time profile",
}
# Logfile column names
COL = {
	"cpu": "cpu",
	"energy": "Wh",
	"cipher": "cipher",
	"cert": "alg",
	"kex": "kex"
}
# Physical units by object
UNIT = {
	"cpu": "s",
	"energy": "W",
	"profile": "samples",
}
# Titles for criteria
CRITERION_TITLE = {
	"cipher": "cipher",
	"cert": "signature algorithm",
	"kex": "key exchange"
}

# Where gnuplot files, data files and images are output
PLOTS_DIR = "/dev/shm/plots"

def gnuplot_histogram(**kwargs):
	cluster = ""
	for i in range(kwargs["nb_impls"]-1):
		cluster += """, "" using {}:xticlabels(1) title col""".format(i+4)
	f = open("{plots_dir}/{object}_by_{criterion}_{side}_{record}.gnuplot".format(plots_dir=PLOTS_DIR, **kwargs), "w")
	f.write("""\
set terminal pngcairo enhanced font "CMU Sans Serif,11" fontscale 1.0 size 800, 600
set output "{plots_dir}/{object}_by_{criterion}_{side}_{record}.png"
set boxwidth 0.9 absolute
set style fill solid 1.0 border lt -1
set style histogram clustered gap 1 title textcolor lt -1
set style data histograms
set title font "CMU Sans Serif,12" "{object_title} by {criterion_title} ({record}, {side} side) ({unit})"
#set xtics border in scale 0,0 nomirror rotate by -45 autojustify
set xtics border in scale 0,0 nomirror autojustify
#set key fixed right top vertical Right noreverse noenhanced autotitle nobox
set key fixed left top vertical Left reverse noenhanced autotitle nobox
set colorbox vertical origin screen 0.9, 0.2 size screen 0.05, 0.6 front noinvert bdefault
set xrange [ * : * ] noreverse writeback
set yrange [ 0 : * ]
set grid y lt 1 lw .75 lc "gray"
plot \
	newhistogram "", "{plots_dir}/{object}_by_{criterion}_{side}_{record}.dat" using 2:xticlabels(1) notitle col, \
	newhistogram "", "{plots_dir}/{object}_by_{criterion}_{side}_{record}.dat" using 3:xticlabels(1) title col{cluster}
""".format(plots_dir=PLOTS_DIR, cluster=cluster, **kwargs))
	f.close()
	os.system("gnuplot {plots_dir}/{object}_by_{criterion}_{side}_{record}.gnuplot".format(plots_dir=PLOTS_DIR, **kwargs))

def gnuplot_stacked_histogram(**kwargs):
	cluster = ""
	#for i in range(kwargs["nb_impls"]-1):
	#	cluster += """, "" using {}:xticlabels(1) title col""".format(i+4)
	f = open("{plots_dir}/{object}_by_{criterion}_{side}_{record}.gnuplot".format(plots_dir=PLOTS_DIR, **kwargs), "w")
	f.write("""\
set terminal pngcairo enhanced font "CMU Sans Serif,11" fontscale 1.0 size 800, 600
set output "{plots_dir}/{object}_by_{criterion}_{side}_{record}.png"
set boxwidth 0.9 absolute
set style fill solid 1.0 border lt -1
set style histogram rowstacked
set style data histograms
set title font "CMU Sans Serif,12" "{object_title} by {criterion_title} ({record}, {side} side) ({unit})"
set xtics border in scale 0,0 nomirror noenhanced rotate by -15 autojustify
set key fixed left top vertical Left noenhanced autotitle nobox invert reverse opaque
set colorbox vertical origin screen 0.9, 0.2 size screen 0.05, 0.6 front noinvert bdefault
set xrange [ * : * ] noreverse writeback
set yrange [ 0 : * ]
set grid y lt 1 lw .75 lc "gray"
plot for [i=2:{nb_functions}] "{plots_dir}/{object}_by_{criterion}_{side}_{record}.dat" using i:xticlabels(1) title col
""".format(plots_dir=PLOTS_DIR, cluster=cluster, **kwargs))
	f.close()
	os.system("gnuplot {plots_dir}/{object}_by_{criterion}_{side}_{record}.gnuplot".format(plots_dir=PLOTS_DIR, **kwargs))

def make_log_plot(logs, exp, criterion, side, obj, record):
	f = open(f"/dev/shm/plots/{obj}_by_{criterion}_{side}_{record}.dat", "w")
	ciphers = {}
	impls = []
	plain_line = None
	idle_val = None
	
	for log in logs:
		if log["exp"] == "idle":
			idle_val = float(log[COL[obj]]) / float(log["time"])
		if log["exp"] != exp or log["record"] != record:
			continue
		if log["setup"] == "none":
			plain_line = "plain {}".format(float(log[COL[obj]]) - idle_val * float(log["time"]))
	
	if plain_line == None:
		return
	
	for log in logs:
		if log["exp"] == exp and log["record"] == record and log["setup"] == side:
			if log[COL[criterion]] not in ciphers:
				ciphers[log[COL[criterion]]] = {}
			ciphers[log[COL[criterion]]][log["impl"]] = float(log[COL[obj]]) - idle_val * float(log["time"])
			if log["impl"] not in impls:
				impls.append(log["impl"])
	impls.sort()
	f.write("{} none {}\n".format(criterion, " ".join(impls)))
	f.write(plain_line+" -"*len(impls)+"\n")
	for cipher in ciphers:
		f.write("{}({}) - {}\n".format(
			ALG_LABEL[cipher],
			VER_LABEL[log["cipher"]],
			" ".join([
				str(ciphers[cipher][impl])
				for impl in impls
			]),
		))
	f.close()
	gnuplot_histogram(object=obj, criterion=criterion, side=side, object_title=OBJ_TITLE[obj], criterion_title=CRITERION_TITLE[criterion], unit=UNIT[obj], nb_impls=len(impls), record=record)

def make_profile_plot(logs, exp, criterion, side, record, no_flamegraph=False):
	f = open(f"/dev/shm/plots/profile_by_{criterion}_{side}_{record}.dat", "w")
	runs = []
	functions = []
	
	for log in logs:
		if log["exp"] == exp and log["record"] == record and log["setup"] == side:
			svg_filename = log["prof"] + ".svg"
			if not no_flamegraph:
				os.system("flamegraph --perfdata {} -o {}".format(log["prof"], svg_filename))
			profile_results = profile.extract_from_file(svg_filename)
			print(profile_results)
			for function in profile_results:
				if function not in functions:
					functions.append(function)
			runs.append({
				criterion: log[COL[criterion]],
				"impl": log["impl"],
				"functions": profile_results,
			})
	f.write("{} {}\n".format(criterion, " ".join(functions)))
	for run in runs:
		f.write("\"{} {}({})\" {}\n".format(
			run["impl"],
			ALG_LABEL[run[criterion]],
			VER_LABEL[log["cipher"]],
			" ".join([
				str(run["functions"][function][0])
				for function in functions
			]),
		))
	f.close()
	gnuplot_stacked_histogram(object="profile", criterion=criterion, side=side, object_title=OBJ_TITLE["profile"], criterion_title=CRITERION_TITLE[criterion], unit=UNIT["profile"], record=record, nb_functions=len(functions)+1)

if __name__ == "__main__":
	cmd = sys.argv[1]
	logfile_name = sys.argv[2]
	logfile = open(logfile_name, "r")
	lines = logfile.readlines()
	logfile.close()

	colnames = lines[0].removesuffix("\n").split(" ")

	logs = []
	records = {}
	for line in lines[1:]:
		cols = line.removesuffix("\n").split(" ")
		log = {}
		for col in range(len(cols)):
			log[colnames[col]] = cols[col]
		if log["record"] != "-":
			records[log["record"]] = ()
		logs.append(log)
	
	os.makedirs("/dev/shm/plots", exist_ok=True)

	no_flamegraph = "-f" in sys.argv

	if cmd == "log":
		for side in ["client", "server"]:
			for record in records:
				make_log_plot(logs, "impl-cipher-ver", "cipher", side, "cpu", record)
				make_log_plot(logs, "impl-cipher-ver", "cipher", side, "energy", record)
				make_log_plot(logs, "impl-cert-ver", "cert", side, "cpu", record)
				make_log_plot(logs, "impl-cert-ver", "cert", side, "energy", record)
				make_log_plot(logs, "impl-kex-ver", "kex", side, "cpu", record)
				make_log_plot(logs, "impl-kex-ver", "kex", side, "energy", record)
	elif cmd == "prof":
		for side in ["client-local", "server-local"]:
			for record in records:
				make_profile_plot(logs, "impl-cipher-ver", "cipher", side, record, no_flamegraph=no_flamegraph)
				make_profile_plot(logs, "impl-cert-ver", "cert", side, record, no_flamegraph=no_flamegraph)
				make_profile_plot(logs, "impl-kex-ver", "kex", side, record, no_flamegraph=no_flamegraph)