import getpass, time, requests, sys

def fetch(site, node, start, stop, user, psw):
	start_str = time.strftime("%Y-%m-%dT%H:%M", time.localtime(int(start)))
	stop_str = time.strftime("%Y-%m-%dT%H:%M", time.localtime(int(stop+61)))
	auth = requests.auth.HTTPBasicAuth(user, psw)
	url = f"https://api.grid5000.fr/stable/sites/{site}/metrics?nodes={node}&metrics=bmc_node_power_watt&start_time={start_str}&end_time={stop_str}"
	print(url)
	resp = requests.get(url, auth=auth)
	if resp.status_code != 200:
		print("Status: ", resp.status_code)
		print(resp.content)
		exit(1)
	data = resp.json()

	energy = 0
	last_power = None
	last_time = 0
	for sample in data:
		t = time.mktime(time.strptime(sample["timestamp"], "%Y-%m-%dT%H:%M:%S.%f%z"))
		if last_power != None:
			energy += last_power * (t - last_time)
			last_power = None
		if t >= start and t <= stop:
			last_power = sample["value"]
			last_time = t
	ws = energy
	wh = ws / 3600
	return {
		"Ws": ws,
		"Wh": wh,
		"w_avg": ws / (stop - start)
	}

def insert_wh_into_logfile(path, site, node, user, psw):
	# Read
	logfile = open(path, "r")
	lines = logfile.readlines()
	logfile.close()

	# Parse
	colnames = lines[0].removesuffix("\n").split(" ")
	col_start = colnames.index("start")
	col_start = colnames.index("stop")
	col_start = colnames.index("Wh")

	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)

	# Modify & Write
	abort = False
	outfile = open(path+"-wh", "w")
	outfile.write(lines[0])
	for log in logs:
		line = ""
		for col in colnames:
			if line != "":
				line += " "
			
			# Modify
			if col == "Wh" and (log[col] == "-" or log[col] == "0") and not abort:
				try:
					measure = fetch(site, node, float(log["start"]), float(log["stop"]), user, psw)
					log[col] = str(measure["Wh"])
				except Exception as e:
					print("Abort:", e)
					abort = True

			line += log[col]
		line += "\n"
		outfile.write(line)
	outfile.close()

def get_psw():
	psw = None
	try:
		f = open("/dev/shm/g5k_psw", "r")
		psw = f.read()
		f.close()
	except:
		pass
	if psw:
		return psw
	psw = getpass.getpass("Psw: ")
	f = open("/dev/shm/g5k_psw", "w")
	f.write(psw)
	f.close()
	return psw

def main():
	if len(sys.argv) < 5:
		print("Usage:")
		print("python g5watt.py <user> <site> <node> <logfile_path>")
		exit(0)

	user = sys.argv[1]
	site = sys.argv[2]
	node = sys.argv[3]
	path = sys.argv[4]

	psw = get_psw()
	insert_wh_into_logfile(path, site, node, user, psw)

if __name__ == "__main__":
	main()