# Dependancy: pandas
# Usage:
# curl -o Locality.cpp https://raw.githubusercontent.com/apple/foundationdb/main/fdbrpc/Locality.cpp
# # (You can replace /main/ by the tag version you want, for example /7.2.9/)
# python generate.py
# # output.csv is the output and can be imported in any spreadsheet editor (gdoc, excel...) for color-coding

from collections import defaultdict
import pandas as pd


PRIORITIES = {
	"ProcessClass::BestFit": 0,
	"ProcessClass::GoodFit": 1,
	"ProcessClass::UnsetFit": 2,
	"ProcessClass::OkayFit": 3,
	"ProcessClass::WorstFit": 4,
	"ProcessClass::NeverAssign": 5,
}

def string_to_role_class(s):
	if not "ProcessClass::" in s:
		raise IOError(s)
	parsed = s.split("ProcessClass::")[1].split(":")[0]
	if parsed.endswith("Class"):
		parsed = parsed[:-5]
	return parsed

def get_priority(l):
	parsed = "ProcessClass::" + l.split("ProcessClass::")[1].split(";")[0]
	return PRIORITIES[parsed]

class_to_process = defaultdict(dict) # dict[class][role] = priority
defaults = {} # role -> default priority

with open("Locality.cpp", "r") as f:
	lines = f.readlines()
	index = 0

	# Find where the switch statement starts
	while "switch(role)" not in lines[index].replace(" ", ""): 
		index += 1
	index += 1

	# Now, iterate for each role, stops when we arrive at default: (for role)
	while True:
		# Get to next statement
		while "case ProcessClass::" not in lines[index] and "default:" not in lines[index]:
			index += 1
		if "default:" in lines[index]:
			break

		role = string_to_role_class(lines[index])
		index += 1

		classes = [] # Since a switch statement can fallthrough, collect all classes first
		while "}" not in lines[index]:
			# Either we have a priority (this if) OR we have a class (next)
			if "return ProcessClass" in lines[index]:
				priority = get_priority(lines[index])
				for class_ in classes:
					c = string_to_role_class(class_)
					class_to_process[c][role] = priority
				classes = []
			elif "case ProcessClass::" in lines[index]:
				classes.append(lines[index])
			elif "default:" in lines[index]:
				defaults[role] = get_priority(lines[index+1])

			index += 1

# Add all defaults to the final dictionary
for c in class_to_process.keys():
	for r in defaults.keys():
		if r not in class_to_process[c]:
			class_to_process[c][r] = defaults[r]


# Generate a CSV via pandas that can be imported into any spreadsheet editor (gdoc, excel...)
print(class_to_process)
df = pd.DataFrame(class_to_process)
df.to_csv("output.csv")