#!/usr/bin/python3
################################################################################
# ram:
# 	used to return information about the ram configuration in a .json
#   format. This is a helper script for use with the
#   cockpit-hardware package (https://github.com/45Drives/cockpit-hardware)
#
# Copyright (C) 2020, Mark Hooper   <mhooper@45drives.com>
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#   
################################################################################
import subprocess
import re
import os
import json
import sys

g_dmi_fields = [
	"Size",
	"Form Factor",
	"Locator",
	"Type",
	"Manufacturer",
	"Serial Number"
]

g_ipmitool_sensor_fields = {
	"P1-DIMMA1 Temp": " (C)",
	"P1-DIMMB1 Temp": " (C)",
	"P1-DIMMC1 Temp": " (C)",
	"P1-DIMMD1 Temp": " (C)",
	"P1-DIMME1 Temp": " (C)",
	"P1-DIMMF1 Temp": " (C)",
	"P1-DIMMG1 Temp": " (C)",
	"P1-DIMMH1 Temp": " (C)",
	"P2-DIMMA1 Temp": " (C)",
	"P2-DIMMB1 Temp": " (C)",
	"P2-DIMMC1 Temp": " (C)",
	"P2-DIMMD1 Temp": " (C)",
	"P2-DIMME1 Temp": " (C)",
	"P2-DIMMF1 Temp": " (C)",
	"P2-DIMMG1 Temp": " (C)",
	"P2-DIMMH1 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)",
	"DIMMG2 Temp": " (C)",
	"DIMMH1 Temp": " (C)",
	"DIMMH2 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
 	"TEMP_DDR5_A1": " (C)", #ASRockRack
	"TEMP_DDR5_B1": " (C)", #ASRockRack
	"DIMMABCD Temp": " (C)", #X12SPL-F
	"DIMMEFGH Temp": " (C)", #X12SPL-F
	"P1_DIMMA~D Temp": " (C)", #X12DPi-N6
	"P1_DIMME~H Temp": " (C)", #X12DPi-N6
	"P2_DIMMA~D Temp": " (C)", #X12DPi-N6
	"P2_DIMME~H Temp": " (C)", #X12DPi-N6
  	"DIMMG0_TEMP": " (C)", #Giga Computing (GIGABYTE)
  	"DIMMG1_TEMP": " (C)", #Giga Computing (GIGABYTE)
   	"DIMMG2_TEMP": " (C)", #Giga Computing (GIGABYTE)
  	"DIMMG3_TEMP": " (C)", #Giga Computing (GIGABYTE)
   	"TEMP_DDR5_A1": " (C)", #ASRockRack
	"TEMP_DDR5_B1": " (C)", #ASRockRack
	"TEMP_CPU1_DDR4A": " (C)", #ASRockRack
 	"TEMP_CPU1_DDR4B": " (C)", #ASRockRack
  	"TEMP_CPU1_DDR4C": " (C)", #ASRockRack
   	"TEMP_CPU1_DDR4D": " (C)", #ASRockRack
	"TEMP_CPU1_DDR4E": " (C)", #ASRockRack
	"TEMP_CPU1_DDR4F": " (C)", #ASRockRack
	"TEMP_CPU1_DDR4G": " (C)", #ASRockRack
	"TEMP_CPU1_DDR4H": " (C)", #ASRockRack
}

# check slot_entries for translatable locator tags (ASRockRack Motherboards, X12 and H12 boards, and GigaByte boards)
locator_tags = {
	"CPU1_DIMM_A1":["DDR4_A"],
	"CPU1_DIMM_A2":["DDR4_A"],
	"CPU1_DIMM_B1":["DDR4_B"],
	"CPU1_DIMM_C1":["DDR4_C"],
	"CPU1_DIMM_D1":["DDR4_D"],
	"CPU1_DIMM_D2":["DDR4_D"],
	"CPU1_DIMM_E1":["DDR4_E"],
	"CPU1_DIMM_F1":["DDR4_F"],
	"DIMMA1":["DIMMABCD","P1_DIMMA~D","DIMMABCD Temp"],
	"DIMMB1":["DIMMABCD","P1_DIMMA~D","DIMMABCD Temp"],
	"DIMMC1":["DIMMABCD","P1_DIMMA~D","DIMMABCD Temp"],
	"DIMMD1":["DIMMABCD","P1_DIMMA~D","DIMMABCD Temp"],
	"DIMME1":["DIMMEFGH","P1_DIMME~H","DIMMEFGH Temp"],
	"DIMMF1":["DIMMEFGH","P1_DIMME~H","DIMMEFGH Temp"],
	"DIMMG1":["DIMMEFGH","P1_DIMME~H","DIMMEFGH Temp"],
	"DIMMH1":["DIMMEFGH","P1_DIMME~H","DIMMEFGH Temp"],
	"P1-DIMMA1": ["P1_DIMMA~D"],
	"P1-DIMMB1": ["P1_DIMMA~D"],
	"P1-DIMMB2": ["P1_DIMMA~D"],
	"P1-DIMMC1": ["P1_DIMMA~D"],
	"P1-DIMMD1": ["P1_DIMMA~D"],
	"P2-DIMMA1": ["P2_DIMMA~D"],
	"P2-DIMMB1": ["P2_DIMMA~D"],
	"P2-DIMMB2": ["P2_DIMMA~D"],
	"P2-DIMMC1": ["P2_DIMMA~D"],
	"P2-DIMMD1": ["P2_DIMMA~D"],
	"P1-DIMME1": ["P1_DIMME~H"],
	"P1-DIMMF1": ["P1_DIMME~H"],
	"P1-DIMMG1": ["P1_DIMME~H"],
	"P1-DIMMH1": ["P1_DIMME~H"],
	"P2-DIMME1": ["P1_DIMME~H"],
	"P2-DIMMF1": ["P1_DIMME~H"],
	"P2-DIMMG1": ["P1_DIMME~H"],
	"P2-DIMMH1": ["P1_DIMME~H"],
	"DIMM_P0_A0": ["DIMMG0", "DIMMG0_TEMP"],
	"DIMM_P0_A1": ["DIMMG0", "DIMMG0_TEMP"],
 	"DIMM_P1_A0": ["DIMMG0", "DIMMG0_TEMP"],
	"DIMM_P1_A1": ["DIMMG0", "DIMMG0_TEMP"],
	"DIMM_P0_B0": ["DIMMG0", "DIMMG0_TEMP"],
	"DIMM_P0_B1": ["DIMMG0", "DIMMG0_TEMP"],
	"DIMM_P0_C0": ["DIMMG0", "DIMMG0_TEMP"],
	"DIMM_P0_D0": ["DIMMG0", "DIMMG0_TEMP"],
	"DIMM_P0_D1": ["DIMMG0", "DIMMG0_TEMP"],
	"DIMM_P0_E0": ["DIMMG1", "DIMMG1_TEMP"],
	"DIMM_P0_E1": ["DIMMG1", "DIMMG1_TEMP"],
	"DIMM_P0_F0": ["DIMMG1", "DIMMG1_TEMP"],
	"DIMM_P0_F1": ["DIMMG1", "DIMMG1_TEMP"],
	"DIMM_P0_H0": ["DIMMG1", "DIMMG1_TEMP"],
	"DIMM_P0_H1": ["DIMMG1", "DIMMG1_TEMP"],
	"DIMM_P0_G0": ["DIMMG1", "DIMMG1_TEMP"],
	"DIMM_P1_I0": ["DIMMG2"],
	"DIMM_P1_J0": ["DIMMG2"],
	"DIMM_P1_K0": ["DIMMG2"],
	"DIMM_P1_L0": ["DIMMG2"],
	"DIMM_P1_M0": ["DIMMG3"],
	"DIMM_P1_N0": ["DIMMG3"],
	"DIMM_P1_O0": ["DIMMG3"],
	"DIMM_P1_P0": ["DIMMG3"],
}

MZ73_locator_tags = {
	"DIMM_P0_A0": ["DIMMG0"],
	"DIMM_P0_B0": ["DIMMG0"],
	"DIMM_P0_C0": ["DIMMG0"],
	"DIMM_P0_D0": ["DIMMG0"],
	"DIMM_P0_E0": ["DIMMG0"],
	"DIMM_P0_F0": ["DIMMG0"],

	"DIMM_P0_H0": ["DIMMG1"],
	"DIMM_P0_G0": ["DIMMG1"],
	"DIMM_P0_I0": ["DIMMG1"],
	"DIMM_P0_J0": ["DIMMG1"],
	"DIMM_P0_K0": ["DIMMG1"],
	"DIMM_P0_L0": ["DIMMG1"],

	"DIMM_P1_M0": ["DIMMG2"],
	"DIMM_P1_N0": ["DIMMG2"],
	"DIMM_P1_O0": ["DIMMG2"],
	"DIMM_P1_P0": ["DIMMG2"],
	"DIMM_P1_Q0": ["DIMMG2"],
	"DIMM_P1_R0": ["DIMMG2"],

	"DIMM_P1_S0": ["DIMMG3"],
	"DIMM_P1_T0": ["DIMMG3"],
	"DIMM_P1_U0": ["DIMMG3"],
	"DIMM_P1_V0": ["DIMMG3"],
	"DIMM_P1_W0": ["DIMMG3"],
	"DIMM_P1_X0": ["DIMMG3"],
}

B550I_AORUS_locator_tags = {
	'P0 CHANNEL A': 'DIMM_A1',
	'P0 CHANNEL B': 'DIMM_B1'
}

EC266D2I_locator_tags = {
   	"Controller0-ChannelA-DIMM0": "TEMP_DDR5_A1",
 	"Controller1-ChannelA-DIMM0": "TEMP_DDR5_B1",
}

ROMED8_locator_tags = {
    "P0 CHANNEL A": "TEMP_CPU1_DDR4A",
    "P0 CHANNEL B": "TEMP_CPU1_DDR4B",
    "P0 CHANNEL C": "TEMP_CPU1_DDR4C",
    "P0 CHANNEL D": "TEMP_CPU1_DDR4D",
    "P0 CHANNEL E": "TEMP_CPU1_DDR4E",
    "P0 CHANNEL F": "TEMP_CPU1_DDR4F",
    "P0 CHANNEL G": "TEMP_CPU1_DDR4G",
    "P0 CHANNEL H": "TEMP_CPU1_DDR4H",
}

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 ipmitool_sensor():
    try:
        ipmitool_sensor_result = subprocess.Popen(
            ["ipmitool", "sensor"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True
        ).stdout
    except Exception as e:
        print(f"ERROR: Unable to run ipmitool: {e}")
        return {}
    if ipmitool_sensor_result is None:
        return {}

    sensor_readings = {}
    for line in ipmitool_sensor_result:
        # Match the sensor name and value using regex
        regex = re.search(r"^(\S+(?:\s\S+)*?)\s+\|\s+(\S+)", line)
        if regex:
            sensor_name = regex.group(1).strip()
            sensor_value = regex.group(2).strip()

            # Directly use the sensor name if it's in the fields list
            if sensor_name in g_ipmitool_sensor_fields:
                sensor_readings[sensor_name] = sensor_value + g_ipmitool_sensor_fields[sensor_name]
            else:
                # Fallback to truncation logic for older cases
                truncated_name = sensor_name[:-5] if len(sensor_name) > 5 else sensor_name
                if truncated_name in g_ipmitool_sensor_fields:
                    sensor_readings[truncated_name] = sensor_value + g_ipmitool_sensor_fields[truncated_name]
    
    return sensor_readings


def dmidecode():
	try:
		dmi_result = subprocess.Popen(
			["dmidecode","-t","17"],stdout=subprocess.PIPE, universal_newlines=True).stdout
	except:
		print("Error running 'dmidecode -t 17'")
		sys.exit(1)
	if dmi_result is None:
		sys.exit(1)

	slot_entries = []
	slot_template = {
		"Size": "No Module Installed",
		"Form Factor": "-",
		"Locator" : "-",
		"Type": "-",
		"Manufacturer": "-",
		"Serial Number" : "-",
		"Temperature" : "-",
		"Temperature Keys": [],
	}

	current_slot = {}
	motherboard_info = json.loads(get_motherboard_model_server_info())
	motherboard_model = motherboard_info['Motherboard'][0].get('Product Name', '')
	# print(f"mobo:{motherboard_model}")
	bank_locator = ""  # Variable to store the bank locator for each slot

	for line in dmi_result:
		if line.startswith("Handle "):
			if current_slot:
				locator = current_slot.get("Locator", "")
				# Filter: Handle if motherboard model is B550I AORUS PRO AX and the Locator is DIMM 1 with valid data
				if motherboard_model == "B550I AORUS PRO AX":
					if current_slot["Locator"] == "DIMM 1" and (
						current_slot["Size"] != "No Module Installed" or current_slot["Size"] == "No Module Installed"
					):
						# Apply B550I_AORUS locator tag based on Bank Locator
						if bank_locator in B550I_AORUS_locator_tags:
							current_slot["Locator"] = B550I_AORUS_locator_tags[bank_locator]
				if motherboard_model == "EC266D2I-2T/AQC" and locator in EC266D2I_locator_tags:
					current_slot["Temperature Keys"] = [EC266D2I_locator_tags[locator]]
				elif motherboard_model.startswith("ROMED8") and bank_locator:
                    # bank_locator looks like “P0 CHANNEL A”
					chan = bank_locator.split()[-1]          # "A" … "H"
					current_slot["Locator"] = f"DIMM {chan}"  # e.g. "DIMM A"
					if bank_locator in ROMED8_locator_tags:
						current_slot["Temperature Keys"] = [ROMED8_locator_tags[bank_locator]]
				slot_entries.append(current_slot.copy())
			current_slot = slot_template.copy()
			bank_locator = ""  # Reset bank locator for new entry
		elif current_slot:
			for field in g_dmi_fields:
				regex = re.search(r"^\s+({fld}):\s+(.*)".format(fld=field), line)
				if regex:
					current_slot[regex.group(1)] = regex.group(2)
			# Capture the Bank Locator
			if line.strip().startswith("Bank Locator:"):
				bank_locator = line.split(":", 1)[1].strip()

	# Append the last entry if it meets criteria
	if current_slot:
		locator = current_slot.get("Locator", "")
		if motherboard_model == "B550I AORUS PRO AX":
			if current_slot["Locator"] == "DIMM 1" and (
				current_slot["Size"] != "No Module Installed" or current_slot["Size"] == "No Module Installed"
			):
				if bank_locator in B550I_AORUS_locator_tags:
					current_slot["Locator"] = B550I_AORUS_locator_tags[bank_locator]
					current_slot["Temperature"] = 'N/A'
		if motherboard_model == "EC266D2I-2T/AQC" and locator in EC266D2I_locator_tags:
			current_slot["Temperature Keys"] = [EC266D2I_locator_tags[locator]]
		if motherboard_model.startswith("ROMED8") and bank_locator:
			chan = bank_locator.split()[-1]
			current_slot["Locator"] = f"DIMM {chan}"
			if bank_locator in ROMED8_locator_tags:
				current_slot["Temperature Keys"] = [ROMED8_locator_tags[bank_locator]]
			slot_entries.append(current_slot.copy())
	
		for slot in slot_entries:
			if motherboard_model.startswith("MZ73"):
				if "Locator" in slot and slot["Locator"] in MZ73_locator_tags:
					slot["Temperature Keys"] = MZ73_locator_tags[slot["Locator"]]
			else:
				if "Locator" in slot and slot["Locator"] in locator_tags:
					slot["Temperature Keys"] = locator_tags[slot["Locator"]]
		slot_entries.append(current_slot.copy())
  
	for slot in slot_entries:
		if motherboard_model.startswith("MZ73"):
			if "Locator" in slot and slot["Locator"] in MZ73_locator_tags:
				slot["Temperature Keys"] = MZ73_locator_tags[slot["Locator"]]
		else:
			if "Locator" in slot and slot["Locator"] in locator_tags:
				slot["Temperature Keys"] = locator_tags[slot["Locator"]]

	return slot_entries

def main():
	ram = dmidecode()
	temps = ipmitool_sensor()

	if ram and temps:
		for module in ram:
			for key in g_dmi_fields:
				if key not in module.keys():
					module[key] = "-"
			if module["Locator"] in temps.keys():
				module["Temperature"] = temps[module["Locator"]]
			elif "Temperature Keys" in module.keys():
				for lookup_key in module["Temperature Keys"]:
					if lookup_key in temps.keys():
						module["Temperature"] = temps[lookup_key]
			if module["Temperature"] == "na (C)":
				module["Temperature"] = "-"
			if module["Manufacturer"] == "NO DIMM":
				module["Manufacturer"] = "-"
			if module["Serial Number"] == "NO DIMM":
				module["Serial Number"] = "-"

	result = []

	for slot in ram:
		result.append({ 
			"locator": slot["Locator"], 
			"type": slot["Type"], 
			"size": slot["Size"], 
			"manufacturer": slot["Manufacturer"], 
			"serialNumber": slot["Serial Number"], 
			"temperature": slot["Temperature"] 
			}
		)
	
	print(json.dumps(result,indent=4))


if __name__ == "__main__":
    main()