from util import (
    MyArgParseHelpFormatter,
    open_device,
    get_hwnetid,
    dict_align_str,
)

from util_eth import (
    load_macsec_from_yaml,
)

from util_sfp import (
    read_macsec_rx_secy_stats,
    read_macsec_tx_secy_stats,
    read_macsec_rx_sc_stats,
    read_macsec_tx_sc_stats,
    sfp_query_verify_macsec_support,
)

import ics
import argparse

###############################################################################
# This loads a single macsec configuration from selected yaml into supported
#  devices and exits
###############################################################################


DESCRIPTION = "MACSec utility"


def main():
    parser = get_parser()
    args = parser.parse_args()
    run_test(args, None)


def run_test(args, power_supply):
    print(f"Attempting to load {args.yaml}...")

    fail = 0

    # connect to test device1 and set network settings for the macsec test
    with open_device(args.serial_number) as device:
        # Just load default every time in case we get some crazy stuff
        ics.load_default_settings(device)

        # if sfp module, query for mdio i2c address
        if args.sfp:
            sfp_i2c_mdio_address = sfp_query_verify_macsec_support(
                device, args.config_netid, args.reset
            )
        # if we only want reset, exit here
        if args.reset:
            exit(0)

        # if we only want the phy macsec counters, read then exit
        if args.counters:
            # Read secy, sc stats
            print("\nReading PHY MACsec rx/tx stats...\n")
            rx_secy_stats = read_macsec_rx_secy_stats(
                device, get_hwnetid(args.config_netid), sfp_i2c_mdio_address
            )
            tx_secy_stats = read_macsec_tx_secy_stats(
                device, get_hwnetid(args.config_netid), sfp_i2c_mdio_address
            )
            rx_sc_stats = read_macsec_rx_sc_stats(
                device, get_hwnetid(args.config_netid), sfp_i2c_mdio_address
            )
            tx_sc_stats = read_macsec_tx_sc_stats(
                device, get_hwnetid(args.config_netid), sfp_i2c_mdio_address
            )
            # print non-zero macsec counts
            print("rx_secy_stats:\n")
            print(dict_align_str(rx_secy_stats, False, True))
            print("tx_secy_stats:\n")
            print(dict_align_str(tx_secy_stats, False, True))
            print("rx_sc_stats:\n")
            print(dict_align_str(rx_sc_stats, False, True))
            print("tx_sc_stats:\n")
            print(dict_align_str(tx_sc_stats, False, True))
            exit(0)

        # load yaml macsec config
        load_macsec_from_yaml(device, args.config_netid, args.yaml, args.sfp)

        if fail:  # bail out if something fails
            print(f"{fail} tests failed.")
            exit(1)


def get_parser():
    parser = argparse.ArgumentParser(
        description="MACsec utility script",
        formatter_class=MyArgParseHelpFormatter,
    )
    parser.add_argument(
        "serial_number",
        help="The serial number of the device to connect to",
    )
    parser.add_argument(
        "--yaml",
        help="yaml file name to be loaded/configured.",
        default="test_sfp_macsec_Rule_wildcard_SecY_128_strip_strict_sci_SA_sak0.yml",
    )
    parser.add_argument(
        "--sfp",
        help="sfp module.",
        default=False,
        action="store_true",
    )
    parser.add_argument(
        "--config_netid",
        help="netid to config the sfp module.",
        default="NETID_I2C2",
    )
    parser.add_argument(
        "--reset",
        help="Only reset and clear and disable macsec on the device, then exit.",
        default=False,
        action="store_true",
    )
    parser.add_argument(
        "--counters",
        help="Only read the phy macsec counters, then exit.",
        default=False,
        action="store_true",
    )
    return parser


if __name__ == "__main__":
    main()