#!/usr/bin/python3
import re
import subprocess
import os
import sys
import json

#supported CPU types
g_dmi_cpu_fields = {
	"Socket Designation":  ["CPU","CPU1","CPU2","P0","P1","CPU0","AM5","AM4"],
	"Version":  [
		"Intel(R) Xeon(R) Silver 4110 CPU @ 2.10GHz",
		"AMD EPYC 7282 16-Core Processor",
  		"AMD EPYC 7281 16-Core Processor",
		"AMD EPYC 7252 8-Core Processor",
		"AMD EPYC 7272 12-Core Processor",
		"AMD EPYC 7452 32-Core Processor",
		"AMD EPYC 8324P 32-Core Processor",
		"AMD EPYC 9454 48-Core Processor",
		"AMD EPYC 8224P 24-Core Processor",
		"AMD EPYC 8534P 64-Core Emb Processor",
	  	"AMD EPYC 8124P 16-Core Processor",
      	"AMD EPYC 8024P 8-Core Processor",
		"AMD Ryzen 9 7950X 16-Core Processor",
		"AMD Ryzen 5 5500GT with Radeon Graphics",
		"AMD Ryzen 7 5700G with Radeon Graphics",
		"Intel(R) Xeon(R) CPU E3-1220 v6 @ 3.00GHz",
		"Intel(R) Xeon(R) Silver 4210 CPU @ 2.10GHz",
		"Intel(R) Xeon(R) Silver 4310 CPU @ 2.10GHz",
		"Intel(R) Xeon(R) Silver 4410Y",
		"Intel(R) Xeon(R) Gold 5418Y",
		"Intel(R) Xeon(R) Gold 6230R CPU @ 2.10GHz",
		"Intel(R) Xeon(R) Gold 6430",
		"Intel(R) Xeon(R) E E-2434",
		"Intel(R) Xeon(R) E E-2478",
  		"Intel(R) Core(TM) i7-14700"
	],
	"Current Speed": None,
	"Max Speed": None,
}


g_ipmitool_sensor_fields = {
	"CPU0 Temp": "(C)",
	"CPU1 Temp": "(C)",
	"CPU2 Temp": "(C)",
	"TEMP_CPU": "(C)",
	"PW Consumption": "(W)",
	"FAN1": "(RPM)",
	"FAN2": "(RPM)",
	"FAN3": "(RPM)",
	"FAN4": "(RPM)",
	"FAN5": "(RPM)",
	"FAN6": "(RPM)",
	"FANA": "(RPM)",
	"FANB": "(RPM)",
	"P1-DIMMA1 Temp": " (C)",
	"P1-DIMMB1 Temp": " (C)",
	"P1-DIMMD1 Temp": " (C)",
	"P1-DIMME1 Temp": " (C)",
	"P2-DIMMA1 Temp": " (C)",
	"P2-DIMMB1 Temp": " (C)",
	"P2-DIMMD1 Temp": " (C)",
	"P2-DIMME1 Temp": " (C)",
	"CPU Temp": " (C)",
	"DIMMA1 Temp": " (C)",
	"DIMMA2 Temp": " (C)",
	"DIMMB1 Temp": " (C)",
	"DIMMB2 Temp": " (C)",
	"DIMMC1 Temp": " (C)",
	"DIMMC2 Temp": " (C)",
	"DIMMD1 Temp": " (C)",
	"DIMMD2 Temp": " (C)",
	"DIMME1 Temp": " (C)",
	"DIMME2 Temp": " (C)",
	"DIMMF1 Temp": " (C)",
	"DIMMF2 Temp": " (C)",
	"DIMMG1 Temp": " (C)",
	"DIMMH1 Temp": " (C)",
	"DDR4_A Temp": " (C)", #ASRockRack 
	"DDR4_B Temp": " (C)", #ASRockRack
	"DDR4_C Temp": " (C)", #ASRockRack
	"DDR4_D Temp": " (C)", #ASRockRack
	"DDR4_E Temp": " (C)", #ASRockRack
	"DDR4_F Temp": " (C)",  #ASRockRack
	"CPU0_TEMP": " (C)", #Giga Computing (GIGABYTE)
	"CPU1_TEMP": " (C)", #Giga Computing (GIGABYTE)
	"DIMMG0_TEMP": " (C)", #Giga Computing (GIGABYTE)
	"DIMMG1_TEMP": " (C)", #Giga Computing (GIGABYTE)
	"DIMMG2_TEMP": " (C)", #Giga Computing (GIGABYTE)
	"DIMMG3_TEMP": " (C)", #Giga Computing (GIGABYTE)
	"CPU0_FAN": " (RPM)", #Giga Computing (GIGABYTE)
	"CPU1_FAN": " (RPM)", #Giga Computing (GIGABYTE)
	"SYS_FAN1": " (RPM)", #Giga Computing (GIGABYTE)
	"SYS_FAN2": " (RPM)", #Giga Computing (GIGABYTE)
	"SYS_FAN3": " (RPM)", #Giga Computing (GIGABYTE)
	"SYS_FAN4": " (RPM)", #Giga Computing (GIGABYTE)
	"SYS_FAN5": " (RPM)", #Giga Computing (GIGABYTE)
 	"SYS_FAN6": " (RPM)", #Giga Computing (GIGABYTE)
	"SYS_POWER": " (W)", #Giga Computing (GIGABYTE)
	"TEMP_CPU": " (C)", 
	"TEMP_DDR5_A1 ": " (C)",
	"TEMP_DDR5_B1 ": " (C)",
	"CPU1 Power": " (W)",
}

def get_motherboard_model_server_info():
	json_path = "/etc/45drives/server_info/server_info.json"
	if os.path.exists(json_path):
		f = open(json_path, "r")
		si = json.load(f)
		f.close()
		mobo_obj = {
			"Motherboard": [{
				"Manufacturer": si["Motherboard"]["Manufacturer"],
				"Product Name": si["Motherboard"]["Product Name"],
				"Serial Number": si["Motherboard"]["Serial Number"]
			}
			]
		}
		return json.dumps(mobo_obj)
	else:
		return "{\"Motherboard\":[{\"Manufacturer\": \"?\", \"Product Name\": \"?\", \"Serial Number\": \"?\"}]}"


def get_cpu_info():
	cpu = []
	try:
		dmi_result = subprocess.Popen(["dmidecode","-t","4"],stdout=subprocess.PIPE,universal_newlines=True).stdout
	except:
		#print("ERROR: dmidecode is not installed")
		return False
	if dmi_result is None:
		return False
	for line in dmi_result:
		for field in g_dmi_cpu_fields.keys():
			regex = re.search(r"^\s({fld}):\s+(.*)".format(fld=field),line)
			if regex != None:
				regex_group1_str = str(regex.group(1)).rstrip()
				regex_group2_str = str(regex.group(2)).rstrip()
				if g_dmi_cpu_fields[regex_group1_str] != None:
					if regex_group2_str in g_dmi_cpu_fields[regex_group1_str]:
						cpu.append((regex_group1_str,regex_group2_str))
					# elif field == "Socket Designation" and regex_group2_str == "CPU":
					# 	cpu.append((regex_group1_str,"CPU1"))
					else:
						cpu.append((regex_group1_str,regex_group2_str))
				else:
					cpu.append((regex_group1_str,regex_group2_str))
	cpu_dict = {
		"CPU":[]
	}

	cpu_count = len(cpu)//len(g_dmi_cpu_fields)
	for i in range(0,cpu_count):
		cpu_dict["CPU"].append(dict(cpu[i*len(g_dmi_cpu_fields):i*len(g_dmi_cpu_fields)+(len(g_dmi_cpu_fields))]))
	return cpu_dict["CPU"]

def get_sensor_readings():
	try:
		ipmitool_sensor_result = subprocess.Popen(
			["ipmitool", "sensor"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True).stdout
	except:
		print("Error running 'ipmitool sensor'")
		sys.exit(1)
	if ipmitool_sensor_result is None:
		sys.exit(1)

	sensor_readings = []
	for line in ipmitool_sensor_result:
		for field in g_ipmitool_sensor_fields.keys():
			regex = re.search(r"^({fld})\s+\|\s+(\S+).*".format(fld=field), line)
			if regex != None:
				# if regex.group(1) == "CPU Temp" or regex.group(1) == "CPU0_TEMP":
				#     # ipmitool does not use the socket name
				#     # we need to manually intervene and store it as
				#     # CPU1 Temp
				#     sensor_readings.append(
				#         ("CPU1 Temp", regex.group(2) +
				#          g_ipmitool_sensor_fields[regex.group(1)])
				#     )
				# else:
				if regex.group(1) == "TEMP_CPU":
					sensor_readings.append(
						("CPU1 Temp", regex.group(2) + g_ipmitool_sensor_fields[regex.group(1)])
					)
				else:
					reading = (regex.group(1), regex.group(2) + g_ipmitool_sensor_fields[regex.group(1)])
					# print(reading)
	
					# All other readings are fine, use the field as captured.
					sensor_readings.append(reading)
	return dict(sensor_readings)


def main():
	cpus = get_cpu_info()
	sensor_readings = get_sensor_readings()
	server_info_json = get_motherboard_model_server_info()
	server_info = json.loads(server_info_json)  # Parse the JSON string into a Python dictionary
	mobo_model = server_info.get("Motherboard", [{}])[0].get("Product Name", "?")
 
	if cpus:
		cpu_arr = []
		# for cpu in cpus:
		#     if (cpu["Socket Designation"]+" Temp") in sensor_readings.keys():
		#         cpu_temp = sensor_readings[cpu["Socket Designation"]+" Temp"]
		#     elif (cpu["Socket Designation"] == 'P0'):
		#         cpu_temp = sensor_readings['CPU1 Temp']
		#     else:
		#         cpu_temp = '-'
					
		#     cpu_arr.append(
		#         { 
		#             'socket': cpu["Socket Designation"], 
		#             'model': cpu["Version"],
		#             'maxSpeed': cpu["Max Speed"],
		#             'currentSpeed': cpu["Current Speed"],
		#             'temperature': cpu_temp
		#         },
		#     )
		for cpu in cpus:
			sock = cpu["Socket Designation"]

			# ----- map socket designation to the right IPMI sensor name -----
			if sock in ("CPU0", "P0", "AM5", "SP6"):
				cand_keys = ("CPU0_TEMP", "CPU Temp")          # try these in order
			elif sock in ("CPU1", "P1"):
				cand_keys = ("CPU1_TEMP", "CPU Temp", "CPU1 Temp",)
			else:                                              # fallback - anything that
				cand_keys = ("CPU Temp",)                      # might exist
			# ---------------------------------------------------------------

			# pick the first key that actually exists in ipmitool output
			for k in cand_keys:
				if k in sensor_readings:
					cpu_temp = sensor_readings[k]
					break
			else:
				cpu_temp = "-"            # none of the candidates were present

			cpu_arr.append(
				{
					"socket":        sock,
					"model":         cpu["Version"],
					"maxSpeed":      cpu["Max Speed"],
					"currentSpeed":  cpu["Current Speed"],
					"temperature":   cpu_temp,
				}
			)


		result = {
			"cpus": cpu_arr
		}
	else:
		result = {
			"cpus": [
				{ 
					'socket': "UNKNOWN", 
					'model': "UNKNOWN",
					'maxSpeed': "-",
					'currentSpeed': "-",
					'temperature': "-"
				}
			]
		}
	print(json.dumps(result, indent=4))

if __name__ == "__main__":
	main()
