#!/usr/bin/env python3
"""Talk to an ET-SoC-1 card's service processor over the management node, without libDM.

The message is the device-management wire format: an 8-byte header (size of the whole message,
tag, command code, flags) followed by the command's payload. It is pushed with
ETSOC1_IOCTL_PUSH_SQ and the reply popped with ETSOC1_IOCTL_POP_CQ, the two ioctls libDM's
device layer uses (devicelayer/src/DevicePcie.cpp:754-819). The kernel driver checks only the
message size, never the command code (et-driver/et-soc1-pcie.c:857-910, et_vqueue.c:867-899),
so a command libDM does not know, like DM_CMD_GET_ASIC_PVT_SAMPLES (74), goes through as is.

Hold the card's lock while running it (the node has one opener):
  flock -n /run/lock/etsoc-shire0.lock python3 pvt_ctl.py temp
"""
import argparse
import ctypes
import errno
import fcntl
import json
import os
import select
import struct
import sys
import time

MAGIC = 0xE7  # ESPERANTO_PCIE_IOCTL_MAGIC, et-driver/et_ioctl.h:15


def _ioc(direction, nr, size):
    return (direction << 30) | (size << 16) | (MAGIC << 8) | nr


class CmdDesc(ctypes.Structure):  # struct cmd_desc, et_ioctl.h:107
    _fields_ = [("cmd", ctypes.c_void_p), ("size", ctypes.c_uint16),
                ("sq_index", ctypes.c_uint16), ("flags", ctypes.c_uint8)]


class RspDesc(ctypes.Structure):  # struct rsp_desc, et_ioctl.h:120
    _fields_ = [("rsp", ctypes.c_void_p), ("size", ctypes.c_uint16), ("cq_index", ctypes.c_uint16)]


PUSH_SQ = _ioc(1, 6, ctypes.sizeof(CmdDesc))  # _IOW(MAGIC, 6, struct cmd_desc)
POP_CQ = _ioc(3, 7, ctypes.sizeof(RspDesc))   # _IOWR(MAGIC, 7, struct rsp_desc)
MAX_MSG = 1760                                # SP VQ maximum message size the device layer reports

DM_CMD_GET_MODULE_CURRENT_TEMPERATURE = 27
DM_CMD_GET_MODULE_RESIDENCY_THROTTLE_STATES = 29
DM_CMD_GET_ASIC_PVT_SAMPLES = 74              # Roman's patch; only on firmware that has it

PVT_GROUPS = {"minshire_ts": 0, "minshire_pd": 1, "minshire_vm": 2, "memshire_vm": 3,
              "ioshire": 4, "pshire_vm": 5, "ts_raw": 7}  # 6 (HILO_RESET) deliberately absent


def open_node(card, wait_s=5.0):
    path = "/dev/et%d_mgmt" % card
    end = time.time() + wait_s
    while True:
        try:
            return os.open(path, os.O_RDWR | os.O_NONBLOCK)
        except OSError as e:
            if e.errno in (errno.EPERM, errno.EBUSY) and time.time() < end:
                time.sleep(0.05)
                continue
            raise


def request(fd, cmd, payload=b"", timeout_s=2.0):
    """Send one command and return (reply bytes, host round-trip seconds)."""
    tag = (os.getpid() ^ int(time.time() * 1000)) & 0xFFFF or 1
    msg = struct.pack("<HHHH", 8 + len(payload), tag, cmd, 0) + payload
    sbuf = ctypes.create_string_buffer(msg, len(msg))
    t0 = time.perf_counter()
    fcntl.ioctl(fd, PUSH_SQ, CmdDesc(ctypes.cast(sbuf, ctypes.c_void_p), len(msg), 0, 0))
    rbuf = ctypes.create_string_buffer(MAX_MSG)
    desc = RspDesc(ctypes.cast(rbuf, ctypes.c_void_p), MAX_MSG, 0)
    poller = select.poll()
    poller.register(fd, select.POLLIN)
    end = time.time() + timeout_s
    while time.time() < end:
        poller.poll(20)
        try:
            n = fcntl.ioctl(fd, POP_CQ, desc)
        except OSError as e:
            if e.errno == errno.EAGAIN:
                continue
            raise
        data = rbuf.raw[:n]
        _size, rtag, rcmd, _flags = struct.unpack_from("<HHHH", data, 0)
        if rtag == tag and rcmd == cmd:
            return data, time.perf_counter() - t0
        print("skipped a message: cmd %d tag %d (%d bytes)" % (rcmd, rtag, n), file=sys.stderr)
    raise TimeoutError("no reply to command %d within %.1f s" % (cmd, timeout_s))


def reply_header(data):
    """dev_mgmt_rsp_header_t: cmn_header_t (8 B), then device_latency_usec (u64), status (i32), pad."""
    size, tag, cmd, flags = struct.unpack_from("<HHHH", data, 0)
    latency_us, status = struct.unpack_from("<Qi", data, 8)
    return {"cmd": cmd, "size": size, "device_us": latency_us, "status": status}


def temperature(fd):
    data, rtt = request(fd, DM_CMD_GET_MODULE_CURRENT_TEMPERATURE)
    h = reply_header(data)
    f = struct.unpack_from("<hhhhhhB", data, 24)  # current_temperature_t, rpc_types.h:81
    keys = ["ioshire_current", "ioshire_low", "ioshire_high", "minshire_avg", "minshire_low",
            "minshire_high", "pmic_sys"]
    return dict(h, host_ms=round(rtt * 1e3, 3), **dict(zip(keys, f)))


def throttle_residency(fd, state):
    payload = struct.pack("<B7x", state)  # device_mgmt_throttle_residency_cmd_t after the header
    data, rtt = request(fd, DM_CMD_GET_MODULE_RESIDENCY_THROTTLE_STATES, payload)
    h = reply_header(data)
    cum, avg, mx, mn = struct.unpack_from("<QQQQ", data, 24)  # residency_t, rpc_types.h:97
    return dict(h, state=state, host_ms=round(rtt * 1e3, 3), cumulative_us=cum, average_us=avg,
                maximum_us=mx, minimum_us=mn)


def ts_code_to_c(code):
    """RUN_1 nominal conversion, in floating point: 0.061 C per code."""
    return (57400 + 249400 * code / 4096 - 124700) / 1000


def pvt_samples(fd, group, first=0, count=255):
    """DM_CMD_GET_ASIC_PVT_SAMPLES (needs firmware with the patch). Returns entries of `stride` words."""
    payload = struct.pack("<BBB5x", PVT_GROUPS[group], first, count)
    data, rtt = request(fd, DM_CMD_GET_ASIC_PVT_SAMPLES, payload)
    h = reply_header(data)
    grp, first, count, stride, total = struct.unpack_from("<BBBBB", data, 24)
    fault_mask, = struct.unpack_from("<Q", data, 32)
    words = struct.unpack_from("<%dh" % (count * stride), data, 40)  # temperatures are int16
    entries = [list(words[i * stride:(i + 1) * stride]) for i in range(count)]
    out = dict(h, host_ms=round(rtt * 1e3, 3), group=grp, first=first, count=count, stride=stride,
               total=total, fault_mask=fault_mask, entries=entries)
    if group == "ts_raw":
        out["celsius"] = [[round(ts_code_to_c(c), 3) for c in e] for e in entries]
    return out


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("what", choices=["temp", "residency", "pvt"])
    ap.add_argument("--card", type=int, default=0)
    ap.add_argument("--state", type=int, default=0, help="throttle state for 'residency'")
    ap.add_argument("--group", default="ts_raw", choices=sorted(PVT_GROUPS), help="group for 'pvt'")
    ap.add_argument("--repeat", type=int, default=1)
    a = ap.parse_args()
    fd = open_node(a.card)
    try:
        for _ in range(a.repeat):
            if a.what == "temp":
                r = temperature(fd)
            elif a.what == "residency":
                r = throttle_residency(fd, a.state)
            else:
                r = pvt_samples(fd, a.group)
            print(json.dumps(r))
    finally:
        os.close(fd)


if __name__ == "__main__":
    main()