#!/usr/bin/env python3

"""
# Copyright (C) 2010 - 2021 Xilinx, Inc.  All rights reserved.
# Copyright (C) 2022 - 2023 Advanced Micro Devices, Inc.  All Rights Reserved.

# SPDX-License-Identifier: MIT
"""

import subprocess
import argparse
import re
import glob
import sys
import distro
import os.path
from periphery import GPIO
from periphery import I2C
from periphery import MMIO

cmds = {'boardid': '/usr/sbin/ipmi-fru',
        'bootfw_status': '/usr/bin/image_update',
        'bootfw_update': '/usr/bin/image_update',
        'listapps': 'dfx-mgr-client',
        'loadapp': 'dfx-mgr-client',
        'unloadapp': 'dfx-mgr-client',
        'xlnx_platformstats': '/usr/bin/xlnx_platformstats',
        'ddrqos': '/usr/bin/ddr-qos',
        'axiqos': '/usr/bin/axi-qos',
        }


def pcie_format(match):
    pcie = match.group(1)
    if pcie is not None:
        pcie = pcie.replace('h', '').replace(' ', '')
        return f"""FRU Board PCIe Info:
    Vendor ID: {pcie[0:4]}
    Device ID: {pcie[4:8]}
    SubVendor ID: {pcie[8:12]}
    SubDevice ID: {pcie[12:16]}
"""


def uuid_format(match):
    uuid = match.group(1)
    if uuid is not None:
        return "FRU Board UUID: " + uuid.replace('h', '').replace(' ', '')


def board_custom_format(fru):
    # Common for all Xilinx boards
    fru = fru.replace('FRU Board Custom Info:',
                      'FRU Board Revision Number:', 1)
    fru = fru.replace('FRU Board Custom Info:', 'FRU Board PCIe Info:', 1)
    fru = re.sub('FRU Board PCIe Info:(.*)\n', pcie_format, fru)
    fru = fru.replace('FRU Board Custom Info:', 'FRU Board UUID:', 1)
    fru = re.sub('FRU Board UUID:(.*)\n', uuid_format, fru)
    return fru


def mac_format(fru):
    # Kria SOM specific interpratation of MAC record
    if re.search('FRU Board Product Name: SM.*', fru):
        fru = fru.replace('FRU OEM MAC ID 0:', 'FRU OEM PS MAC ID 0:', 1)
    # Kria Carrier card specific interpretation of MAC record
    prod = re.search('FRU Board Product Name: SCK-(.*)-G', fru)
    if prod and prod.group(1) == 'KR':
        fru = fru.replace('FRU OEM MAC ID 0:', 'FRU OEM PS MAC ID 1:', 1)
        fru = fru.replace('FRU OEM MAC ID 1:', 'FRU OEM PL MAC ID 0:', 1)
        fru = fru.replace('FRU OEM MAC ID 2:', 'FRU OEM PL MAC ID 1:', 1)
    elif prod and prod.group(1) == 'KD':
        fru = fru.replace('FRU OEM MAC ID 0:', 'FRU OEM PL MAC ID 0:', 1)
        fru = fru.replace('FRU OEM MAC ID 1:', 'FRU OEM PL MAC ID 1:', 1)
    return fru


def free_form_format(fru):
    # Kria SOM specific interpratation of MAC record
    if re.search('FRU Board Product Name: SM.*', fru):
        fru = fru.replace('FRU OEM Memory:', 'FRU OEM Primary Boot Device:', 1)
        fru = fru.replace('FRU OEM Memory:',
                          'FRU OEM Secondary Boot Device:', 1)
        fru = fru.replace('FRU OEM Memory:', 'FRU OEM PS DDR Memory:', 1)
        fru = fru.replace('FRU OEM Memory:', 'FRU OEM PL DDR Memory:', 1)
    return fru


def boardid_format(fru):
    # Xilinx specific interpretation of custom board info
    fru = board_custom_format(fru)
    # Xilinx specific interpretation of MAC record
    fru = mac_format(fru)
    # Xilinx specific interpretation of free form record
    fru = free_form_format(fru)
    return fru


def get_eeprom(args):
    eeprom = []
    for arg in args:
        if '--fru-file' in arg:
            eeprom.append(arg.split('=')[1])

    if not eeprom:
        eeprom = glob.glob('/sys/devices/platform/axi/*.i2c/*/*/eeprom')

    if not eeprom:
        sys.exit('No EEPROM nodes found')

    return eeprom


def boardid(args):
    ignore = False
    if any('--ignore-errors' in arg for arg in args):
        ignore = True

    for file in get_eeprom(args):
        ret, fru = subprocess.getstatusoutput(cmds['boardid'] +
                                              ' --fru-file=' + file + ' --interpret-oem-data')
        if not ignore and ret != 0:
            print(fru)
            sys.exit('\nipmi-fru returned error code ' + str(ret))
        elif fru:
            print('\n' + boardid_format(fru))

    print('')


def product_name(args):
    ignore = False
    if any('--ignore-errors' in arg for arg in args):
        ignore = True

    prod = ''
    som = ''
    carrier = ''
    for file in get_eeprom(args):
        ret, fru = subprocess.getstatusoutput(cmds['boardid'] +
                                              ' --fru-file=' + file + ' --interpret-oem-data')
        if not ignore and ret != 0:
            raise RuntimeError('ipmi-fru returned error code ' + str(ret))
        elif fru:
            # Kria SOM
            match = re.search('FRU Board Product Name: SM.*-K(.*)-.*', fru)
            if match:
                som = match.group(1)
                continue
            # Kria Carrier Card
            match = re.search('FRU Board Product Name: SCK-(.*)-G', fru)
            if match:
                carrier = match.group(1).lower()
                continue
            # Other Xilinx Board
            match = re.search('FRU Board Product Name: (.*)', fru)
            if match:
                prod = match.group(1).lower()

    if som and carrier:
        prod = carrier + som + '0'

    return prod


def bootfw_status(args):
    subprocess.run([cmds['bootfw_status'], '-p'])


def bootfw_update(args):
    subprocess.run([cmds['bootfw_update']]+args)


def getpkgs(args):
    if ('-h' in args) or ('--help' in args):
        print("\ngetpkgs will use ipmi-fru to determine the target platform\nand search the package feed for compatible packagegroups\n")
        exit()
    try:
        prod = product_name(args + [' --ignore-errors'])
    except:
        sys.exit("No valid product name found, can't search package feed")
    print("\nSearching package feed for packages compatible with: " + prod + '\n')
    if distro.id() == 'ubuntu':
        subprocess.run('apt search xlnx-app-' + prod +
                       '- 2>/dev/null', shell=True)
    elif distro.id() == 'petalinux':
        subprocess.run('dnf list --available | grep packagegroup-' + prod +
                       '- | awk "!/-dev\./&&!/-lic\./&&!/-dbg\./&&!/-ptest\./"', shell=True)


def listapps(args):
    subprocess.run([cmds['listapps'], '-listPackage'])


def loadapp(args):
    subprocess.run([cmds['loadapp'], '-load']+args)


def unloadapp(args):
    subprocess.run([cmds['unloadapp'], '-remove']+args)


def xlnx_platformstats(args):
    subprocess.run([cmds['xlnx_platformstats']]+args)


def ddrqos(args):
    subprocess.run([cmds['ddrqos']]+args)


def axiqos(args):
    subprocess.run([cmds['axiqos']]+args)


def pl_node_state_update(state):
    # PMU call for PL node request/release
    pm_dev = open("/sys/kernel/debug/zynqmp-firmware/pm", "w")
    pm_dev.write("time echo " + state + " 69")


def pl_sequencer_on():
    gpio_out = GPIO("/dev/gpiochip1", 33, "out")
    gpio_out.write(True)
    gpio_out.close()
    pl_node_state_update("pm_request_node")


def pl_sequencer_off():
    pl_node_state_update("pm_release_node")
    gpio_out = GPIO("/dev/gpiochip1", 33, "out")
    gpio_out.write(False)
    gpio_out.close()


def pl_sequencer_status():
    #To get status of the PL sequencer we need to read the value of GPIO pin (MIO33),
    #but that GPIO configured as output, the only way to get the status of the GPIO is by
    #reading the data register(0xFF0A0044), here we are doing that using MMIO.
    gpio_mmio = MMIO(0xFF0A0044, 4)
    reg_value = gpio_mmio.read32(0x00)
    gpio_mmio.close()
    if ((reg_value >> 7) & 1) == 0:
        print("PL sequencer is off \n")
    elif ((reg_value >> 7) & 1) == 1:
        print("PL sequencer is on \n")
    else:
        print("Unknown PL sequencer state, reg_value: ", reg_value, '\n')


def pl_sequencer_on_via_i2c():
    i2c = I2C("/dev/i2c-1")
    msgs = [I2C.Message([0x08, 0x3E], read=False)]
    # Write '0x3E' to PL sequencer power register for PL switch on
    i2c.transfer(0x68, msgs)
    i2c.close()
    pl_node_state_update("pm_request_node")


def pl_sequencer_off_via_i2c():
    pl_node_state_update("pm_release_node")
    i2c = I2C("/dev/i2c-1")
    # Write '0x00' to PL sequencer power register for PL switch off
    msgs = [I2C.Message([0x08, 0x00], read=False)]
    i2c.transfer(0x68, msgs)
    i2c.close()


def pl_sequencer_status_via_i2c():
    i2c = I2C("/dev/i2c-1")
    msgs = [I2C.Message([0xF5]), I2C.Message([0x00], read=True)]
    i2c.transfer(0x68, msgs)
    i2c.close()
    if msgs[1].data[0] == 0x9a:
        print("PL sequencer is on \n")
    elif msgs[1].data[0] == 0x04:
        print("PL sequencer is off \n")
    else:
        print("Unknown PL sequencer state: ", hex(msgs[1].data[0]), "\n")


def pwrctl(args):
    try:
        product = product_name(args + [' --ignore-errors'])
    except:
        sys.exit("Error during product name fetching, can't perform pwrctl operation")

    if product.find('K26') != -1 or product.find('kr260') != -1:
        if sys.argv[2] == "--off":
            pl_sequencer_off_via_i2c()
        elif sys.argv[2] == "--on":
            pl_sequencer_on_via_i2c()
        elif sys.argv[2] == "--status":
            pl_sequencer_status_via_i2c()
    elif product.find('K24') != -1 or product.find('kd240') != -1:
        if sys.argv[2] == "--off":
            pl_sequencer_off()
        elif sys.argv[2] == "--on":
            pl_sequencer_on()
        elif sys.argv[2] == "--status":
            pl_sequencer_status()
    else:
        print('Unkown device: xmutil pwrctl supports K24/K26 only\n')


def top(cmd, args):
    globals()[cmd[0]](args)


def desktop_disable(args):
    subprocess.run('systemctl isolate multi-user.target', shell=True)


def desktop_enable(args):
    subprocess.run('systemctl isolate graphical.target', shell=True)


def dp_unbind(args):
    subprocess.run(
        'echo -n fd4a0000.display > /sys/bus/platform/drivers/zynqmp-display/unbind', shell=True)


def dp_bind(args):
    subprocess.run(
        'echo -n fd4a0000.display > /sys/bus/platform/drivers/zynqmp-display/bind', shell=True)


if __name__ == "__main__":

    help_text = """
boardid: Reads FRU data based on EEPROM or file and prints information in human readable format.
    --fru-file=<file>    Pass a <file> containing fru data
    --ignore-errors      Ignore errors returned by fru-ipmi and print returned output

bootfw_status: Prints Qspi MFG version and date info along with persistent state values.

bootfw_update: Updates the primary boot device with a new BOOT.BIN in the inactive partition (either A or B).

getpkgs: Queries Xilinx package feeds and provides summary of relevant packages for active platform based on board ID information.
    --fru-file=<file>    Pass a <file> containing fru data
    --ignore-errors      Ignore errors returned by fru-ipmi and print returned output

listapps: Queries on target FW resource manager daemon of pre-built app bitstreams available on the platform and provides summary to CLI.

loadapp: Loads requested application configuration bitstream to programmable logic if the device is available.

unloadapp: Removes application bitstream. (Takes slot number, default 0)

xlnx_platformstats: Reads and prints a summary of the following performance related information:
    CPU Utilization for each configured CPU
    RAM utilization
    Swap memory Utilization
    SOM overall current, power, voltage utilization
    SysMon Temperatures(s)
    SOM power supply data summary reported by PMICs & ZU+ SysMon sources

ddrqos: Set QOS value for DDR slots on zynqmp platform.

axiqos: Set QOS value for AXI ports on zynqmp platform.

pwrctl: PL power control utility.
    --off    PL shutdown
    --on     PL powerup
    --status PL power status

desktop_disable: Disable the desktop.

desktop_enable: Enable the desktop.

dp_unbind: Unbind the display driver.

dp_bind: Bind the display driver.
"""

    parser = argparse.ArgumentParser(
        description=help_text, formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument('cmd', choices=['boardid', 'bootfw_status', 'bootfw_update', 'getpkgs', 'listapps', 'loadapp', 'unloadapp', 'xlnx_platformstats',
                        'ddrqos', 'axiqos', 'pwrctl', 'desktop_disable', 'desktop_enable', 'dp_unbind', 'dp_bind'], type=str, nargs=1, help='Enter a function')
    parser.add_argument('args', nargs=argparse.REMAINDER)
    args = parser.parse_args()
    top(args.cmd, args.args)
