#!/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

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
	"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,
 	"DIMMAB Temp": " (C)", #X13SAE-F
 	"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_DDR5_C1": " (C)", #ASRockRack
	"TEMP_DDR5_E1": " (C)", #ASRockRack
	"TEMP_DDR5_G1": " (C)", #ASRockRack
	"TEMP_DDR5_H1": " (C)", #ASRockRack
	"TEMP_DDR5_I1": " (C)", #ASRockRack
	"TEMP_DDR5_K1": " (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, X12SPL-F, 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","DIMMAB Temp"],
	"DIMMB1":["DIMMABCD","P1_DIMMA~D","DIMMABCD Temp","DIMMAB Temp"],
 	"DIMMA2":["DIMMAB Temp"],
	"DIMMB2":["DIMMAB 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",
}

GENOA_locator_tags = {
    "P0 CHANNEL A": "TEMP_DDR5_A1",
    "P0 CHANNEL B": "TEMP_DDR5_B1",
    "P0 CHANNEL C": "TEMP_DDR5_C1",
    "P0 CHANNEL E": "TEMP_DDR5_E1",
    "P0 CHANNEL G": "TEMP_DDR5_G1",
    "P0 CHANNEL H": "TEMP_DDR5_H1",
    "P0 CHANNEL I": "TEMP_DDR5_I1",
    "P0 CHANNEL K": "TEMP_DDR5_K1",
}

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",
}

# Replace your current PROART_X870E_locator_tags with:
PROART_X870E_locator_tags = {
    "P0 CHANNEL A": {
        "DIMM 0": "DIMM 0",
        "DIMM 1": "DIMM 1",
    },
    "P0 CHANNEL B": {
        "DIMM 0": "DIMM 2",
        "DIMM 1": "DIMM 3",
    },
}

locator_tags.update({
    "DIMMA2": ["DIMMA2", "DIMMA1", "DIMMA1 Temp"],
    "DIMMB2": ["DIMMB2", "DIMMB1", "DIMMB1 Temp"],
})
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:
		return False
	if dmi_result is None:
		return False

	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]
				is_proart_x870e = ("PROART" in motherboard_model.upper()) and ("X870E" in motherboard_model.upper())
				if is_proart_x870e and bank_locator:
					bl = bank_locator.strip()  # "P0 CHANNEL A/B"
					loc = current_slot.get("Locator", "").strip()  # "DIMM 0/1"
					if bl in PROART_X870E_locator_tags and loc in PROART_X870E_locator_tags[bl]:
						current_slot["Locator"] = PROART_X870E_locator_tags[bl][loc]			
				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]]
				elif motherboard_model.startswith("GENOAD8X") and bank_locator:
					chan = bank_locator.split()[-1]          # "A" … "K"
					current_slot["Locator"] = f"DIMM {chan}"  # e.g. "DIMM A"
					if bank_locator in GENOA_locator_tags:
						current_slot["Temperature Keys"] = [GENOA_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'			
		is_proart_x870e = ("PROART" in motherboard_model.upper()) and ("X870E" in motherboard_model.upper())
		if is_proart_x870e and bank_locator:
			bl = bank_locator.strip()
			loc = current_slot.get("Locator", "").strip()
			if bl in PROART_X870E_locator_tags and loc in PROART_X870E_locator_tags[bl]:
				current_slot["Locator"] = PROART_X870E_locator_tags[bl][loc]			
		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]]
		if motherboard_model.startswith("GENOAD8X") and bank_locator:
			chan = bank_locator.split()[-1]
			current_slot["Locator"] = f"DIMM {chan}"
			if bank_locator in GENOA_locator_tags:
				current_slot["Temperature Keys"] = [GENOA_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"]]

	return slot_entries

def main():
    ram = dmidecode()
    temps = ipmitool_sensor()
    output_str = "{\"Ram Info\":[]}"

    if ram:
        for module in ram:
            # Ensure required DMI fields exist
            for key in g_dmi_fields:
                if key not in module:
                    module[key] = "-"

            # Resolve temperature
            temp_found = False
            if temps:
                name = module.get("Locator", "")
                # candidates: exact locator, "<locator> Temp"
                candidates = [name, f"{name} Temp"] if name else []
                # plus any Temperature Keys (both exact and "... Temp")
                for lk in module.get("Temperature Keys", []):
                    candidates.extend([lk, f"{lk} Temp"])

                for c in candidates:
                    if c in temps:
                        module["Temperature"] = temps[c]
                        temp_found = True
                        break

            # Normalize special/absent values
            if module.get("Temperature") == "na (C)":
                module["Temperature"] = "-"

            # If a DIMM is installed but we still have no temperature, mark as N/A
            if (module.get("Size") and module["Size"] != "No Module Installed" and
                (not module.get("Temperature") or module["Temperature"] == "-")):
                module["Temperature"] = "N/A"
                module.setdefault("Temperature Keys", [])

            # Clean up vendor "NO DIMM" placeholders
            if module.get("Manufacturer") == "NO DIMM":
                module["Manufacturer"] = "-"
            if module.get("Serial Number") == "NO DIMM":
                module["Serial Number"] = "-"

        # Build output
        output_str = "{\"Ram Info\":["
        for module in ram:
            output_str += json.dumps(module) + ","
        output_str = output_str[:-1] + "]}"
        print(json.dumps(json.loads(output_str), indent=4))
    else:
        print(json.dumps(json.loads(output_str), indent=4))

if __name__ == "__main__":
    main()