#!/usr/bin/env python3

import argparse
from typing import NamedTuple

import yaml


class RptSensor(NamedTuple):
    name: str
    entityId: int
    typeId: int
    evtType: int
    sensorId: int
    fru: int
    targetPath: str


sampleDimmTemp = {
    "bExp": 0,
    "entityID": 32,
    "entityInstance": 2,
    "interfaces": {
        "xyz.openbmc_project.Sensor.Value": {
            "Value": {"Offsets": {255: {"type": "int64_t"}}}
        }
    },
    "multiplierM": 1,
    "mutability": "Mutability::Write|Mutability::Read",
    "offsetB": -127,
    "path": "/xyz/openbmc_project/sensors/temperature/dimm0_temp",
    "rExp": 0,
    "readingType": "readingData",
    "scale": -3,
    "sensorNamePattern": "nameLeaf",
    "sensorReadingType": 1,
    "sensorType": 1,
    "serviceInterface": "org.freedesktop.DBus.Properties",
    "unit": "xyz.openbmc_project.Sensor.Value.Unit.DegreesC",
}
sampleCoreTemp = {
    "bExp": 0,
    "entityID": 208,
    "entityInstance": 2,
    "interfaces": {
        "xyz.openbmc_project.Sensor.Value": {
            "Value": {"Offsets": {255: {"type": "int64_t"}}}
        }
    },
    "multiplierM": 1,
    "mutability": "Mutability::Write|Mutability::Read",
    "offsetB": -127,
    "path": "/xyz/openbmc_project/sensors/temperature/p0_core0_temp",
    "rExp": 0,
    "readingType": "readingData",
    "scale": -3,
    "sensorNamePattern": "nameLeaf",
    "sensorReadingType": 1,
    "sensorType": 1,
    "serviceInterface": "org.freedesktop.DBus.Properties",
    "unit": "xyz.openbmc_project.Sensor.Value.Unit.DegreesC",
}
samplePower = {
    "bExp": 0,
    "entityID": 10,
    "entityInstance": 13,
    "interfaces": {
        "xyz.openbmc_project.Sensor.Value": {
            "Value": {"Offsets": {255: {"type": "int64_t"}}}
        }
    },
    "multiplierM": 2,
    "offsetB": 0,
    "path": "/xyz/openbmc_project/sensors/power/p0_power",
    "rExp": 0,
    "readingType": "readingData",
    "scale": -6,
    "sensorNamePattern": "nameLeaf",
    "sensorReadingType": 1,
    "sensorType": 8,
    "serviceInterface": "org.freedesktop.DBus.Properties",
    "unit": "xyz.openbmc_project.Sensor.Value.Unit.Watts",
}

sampleDcmiSensor = {
    "instance": 1,
    "dbus": "/xyz/openbmc_project/sensors/temperature/p0_core0_temp",
    "record_id": 91,
}


def openYaml(f):
    return yaml.load(open(f))


def saveYaml(y, f, safe=True):
    if safe:
        noaliasDumper = yaml.dumper.SafeDumper
        noaliasDumper.ignore_aliases = lambda self, data: True
        yaml.dump(
            y, open(f, "w"), default_flow_style=False, Dumper=noaliasDumper
        )
    else:
        yaml.dump(y, open(f, "w"))


def getEntityIdAndNamePattern(p, intfs, m):
    key = (p, intfs)
    match = m.get(key, None)
    if match is None:
        # Workaround for P8's occ sensors, where the path look like
        # /org/open_power/control/occ_3_0050
        if (
            "/org/open_power/control/occ" in p
            and "org.open_power.OCC.Status" in intfs
        ):
            return (210, "nameLeaf")
        raise Exception("Unable to find sensor", key, "from map")
    return (m[key]["entityID"], m[key]["sensorNamePattern"])


# Global entity instances
entityInstances = {}


def getEntityInstance(id):
    instanceId = entityInstances.get(id, 0)
    instanceId = instanceId + 1
    entityInstances[id] = instanceId
    print("EntityId:", id, "InstanceId:", instanceId)
    return instanceId


def loadRpt(rptFile):
    sensors = []
    with open(rptFile) as f:
        next(f)
        next(f)
        for line in f:
            fields = line.strip().split("|")
            fields = list(map(str.strip, fields))
            sensor = RptSensor(
                fields[0],
                int(fields[2], 16) if fields[2] else None,
                int(fields[3], 16) if fields[3] else None,
                int(fields[4], 16) if fields[4] else None,
                int(fields[5], 16) if fields[5] else None,
                int(fields[7], 16) if fields[7] else None,
                fields[9],
            )
            # print(sensor)
            sensors.append(sensor)
    return sensors


def getDimmTempPath(p):
    # Convert path like: /sys-0/node-0/motherboard-0/dimmconn-0/dimm-0
    # to: /xyz/openbmc_project/sensors/temperature/dimm0_temp
    import re

    dimmconn = re.search(r"dimmconn-\d+", p).group()
    dimmId = re.search(r"\d+", dimmconn).group()
    return "/xyz/openbmc_project/sensors/temperature/dimm{}_temp".format(
        dimmId
    )


def getMembufTempPath(name):
    # Convert names like MEMBUF0_Temp or CENTAUR0_Temp
    # to: /xyz/openbmc_project/sensors/temperature/membuf0_temp
    # to: /xyz/openbmc_project/sensors/temperature/centaur0_temp
    return "/xyz/openbmc_project/sensors/temperature/{}".format(name.lower())


def getCoreTempPath(name, p):
    # For different rpts:
    # Convert path like:
    #   /sys-0/node-0/motherboard-0/proc_socket-0/module-0/p9_proc_s/eq0/ex0/core0 (for P9)  # noqa: E501
    # to: /xyz/openbmc_project/sensors/temperature/p0_core0_temp
    # or name like: CORE0_Temp (for P8)
    # to: /xyz/openbmc_project/sensors/temperature/core0_temp (for P8)
    import re

    if "p9_proc" in p:
        splitted = p.split("/")
        socket = re.search(r"\d+", splitted[4]).group()
        core = re.search(r"\d+", splitted[9]).group()
        return (
            "/xyz/openbmc_project/sensors/temperature/p{}_core{}_temp".format(
                socket, core
            )
        )
    else:
        core = re.search(r"\d+", name).group()
        return "/xyz/openbmc_project/sensors/temperature/core{}_temp".format(
            core
        )


def getPowerPath(name):
    # Convert name like Proc0_Power
    # to: /xyz/openbmc_project/sensors/power/p0_power
    import re

    r = re.search(r"\d+", name)
    if r:
        index = r.group()
    else:
        # Handle cases like IO_A_Power, Storage_Power_A
        r = re.search(r"_[A|B|C|D]", name).group()[-1]
        index = str(ord(r) - ord("A"))
    prefix = "p"
    m = None
    if "memory_proc" in name.lower():
        prefix = None
        m = "centaur"
    elif "pcie_proc" in name.lower():
        m = "pcie"
    elif "io" in name.lower():
        m = "io"
    elif "fan" in name.lower():
        m = "fan"
    elif "storage" in name.lower():
        m = "disk"
    elif "total" in name.lower():
        prefix = None
        m = "total"
    elif "proc" in name.lower():
        # Default
        pass

    ret = "/xyz/openbmc_project/sensors/power/"
    if prefix:
        ret = ret + prefix + index
    if m:
        if prefix:
            ret = ret + "_" + m
        else:
            ret = ret + m
    if prefix is None:
        ret = ret + index
    ret = ret + "_power"
    return ret


def getDimmTempConfig(s):
    r = sampleDimmTemp.copy()
    r["entityInstance"] = getEntityInstance(r["entityID"])
    r["path"] = getDimmTempPath(s.targetPath)
    return r


def getMembufTempConfig(s):
    r = sampleDimmTemp.copy()
    r["entityID"] = 0xD1
    r["entityInstance"] = getEntityInstance(r["entityID"])
    r["path"] = getMembufTempPath(s.name)
    return r


def getCoreTempConfig(s):
    r = sampleCoreTemp.copy()
    r["entityInstance"] = getEntityInstance(r["entityID"])
    r["path"] = getCoreTempPath(s.name, s.targetPath)
    return r


def getPowerConfig(s):
    r = samplePower.copy()
    r["entityInstance"] = getEntityInstance(r["entityID"])
    r["path"] = getPowerPath(s.name)
    return r


def isCoreTemp(p):
    import re

    m = re.search(r"p\d+_core\d+_temp", p)
    return m is not None


def getDcmiSensor(i, sensor):
    import re

    path = sensor["path"]
    name = path.split("/")[-1]
    m = re.findall(r"\d+", name)
    socket, core = int(m[0]), int(m[1])
    instance = socket * 24 + core + 1
    r = sampleDcmiSensor.copy()
    r["instance"] = instance
    r["dbus"] = path
    r["record_id"] = i
    return r


def saveJson(data, f):
    import json

    with open(f, "w") as outfile:
        json.dump(data, outfile, indent=4)


def main():
    parser = argparse.ArgumentParser(
        description="Yaml tool for updating ipmi sensor yaml config"
    )
    parser.add_argument(
        "-i",
        "--input",
        required=True,
        dest="input",
        help="The ipmi sensor yaml config",
    )
    parser.add_argument(
        "-o",
        "--output",
        required=True,
        dest="output",
        help="The output yaml file",
    )
    parser.add_argument(
        "-m",
        "--map",
        dest="map",
        default="sensor_map.yaml",
        help="The sample map yaml file",
    )
    parser.add_argument(
        "-r", "--rpt", dest="rpt", help="The .rpt file generated by op-build"
    )
    parser.add_argument(
        "-f",
        "--fix",
        action="store_true",
        help="Fix entities and sensorNamePattern",
    )

    # -g expects output as yaml for mapping of entityID/sensorNamePattern
    # -d expects output as json config for dcmi sensors
    # Do not mess the output by enforcing only one argument is passed
    # TODO: -f and -r could be used together, and they are conflicted with
    #       -g or -d
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "-g",
        "--generate",
        action="store_true",
        help="Generate maps for entityID and sensorNamePattern",
    )
    group.add_argument(
        "-d",
        "--dcmi",
        action="store_true",
        help="Generate dcmi sensors json config",
    )

    args = parser.parse_args()
    args = vars(args)

    if args["input"] is None or args["output"] is None:
        parser.print_help()
        exit(1)

    y = openYaml(args["input"])

    if args["fix"]:
        # Fix entities and sensorNamePattern
        m = openYaml(args["map"])

        for i in y:
            path = y[i]["path"]
            intfs = tuple(sorted(list(y[i]["interfaces"].keys())))
            entityId, namePattern = getEntityIdAndNamePattern(path, intfs, m)
            y[i]["entityID"] = entityId
            y[i]["entityInstance"] = getEntityInstance(entityId)
            y[i]["sensorNamePattern"] = namePattern
            print(
                y[i]["path"],
                "id:",
                entityId,
                "instance:",
                y[i]["entityInstance"],
            )

    sensorIds = list(y.keys())
    if args["rpt"]:
        unhandledSensors = []
        rptSensors = loadRpt(args["rpt"])
        for s in rptSensors:
            if s.sensorId is not None and s.sensorId not in sensorIds:
                print(
                    "Sensor ID",
                    s.sensorId,
                    "not in yaml:",
                    s.name,
                    ", path:",
                    s.targetPath,
                )
                isAdded = False
                if "temp" in s.name.lower():
                    if "dimm" in s.targetPath.lower():
                        y[s.sensorId] = getDimmTempConfig(s)
                        isAdded = True
                    elif "core" in s.targetPath.lower():
                        y[s.sensorId] = getCoreTempConfig(s)
                        isAdded = True
                    elif (
                        "centaur" in s.name.lower()
                        or "membuf" in s.name.lower()
                    ):
                        y[s.sensorId] = getMembufTempConfig(s)
                        isAdded = True
                elif s.name.lower().endswith("_power"):
                    y[s.sensorId] = getPowerConfig(s)
                    isAdded = True

                if isAdded:
                    print(
                        "Added sensor id:",
                        s.sensorId,
                        ", path:",
                        y[s.sensorId]["path"],
                    )
                else:
                    unhandledSensors.append(s)

        print("Unhandled sensors:")
        for s in unhandledSensors:
            print(s)

    if args["generate"]:
        m = {}
        for i in y:
            path = y[i]["path"]
            intfs = tuple(sorted(list(y[i]["interfaces"].keys())))
            entityId = y[i]["entityID"]
            sensorNamePattern = y[i]["sensorNamePattern"]
            m[(path, intfs)] = {
                "entityID": entityId,
                "sensorNamePattern": sensorNamePattern,
            }
        y = m

    if args["dcmi"]:
        d = []
        for i in y:
            if isCoreTemp(y[i]["path"]):
                s = getDcmiSensor(i, y[i])
                d.append(s)
                print(s)
        saveJson(d, args["output"])
        return

    safe = False if args["generate"] else True
    saveYaml(y, args["output"], safe)


if __name__ == "__main__":
    main()