#!/usr/bin/python3
# print useful diagnostics info about any attached LIGO PCIe Timing Cards (LPTC).
# Uses sysfs information generated by gpstime kernel module, which must be installed.

import argparse
import os
import sys
from collections import namedtuple

LPTC_SYSFS = "/sys/kernel/gpstime/lptc"
VOLTAGE_TOLERANCE = 0.10  # +/- 10% before flagging out of range

# ---------- color ----------
_USE_COLOR = False
_COLORS = {
    'red': '31', 'green': '32', 'yellow': '33',
    'cyan': '36', 'dim': '2', 'bold': '1',
}


def _c(text, color):
    if not _USE_COLOR or color is None:
        return text
    return f"\033[{_COLORS[color]}m{text}\033[0m"


def _setup_color(mode):
    global _USE_COLOR
    if mode == 'always':
        _USE_COLOR = True
    elif mode == 'never':
        _USE_COLOR = False
    else:
        _USE_COLOR = sys.stdout.isatty() and not os.environ.get('NO_COLOR')


# ---------- sysfs helpers ----------
def read_lptc_val(fname):
    with open(os.path.join(LPTC_SYSFS, fname), "rt") as f:
        return f.readline().strip()


def read_lptc_int(fname):
    return int(read_lptc_val(fname), 16)


def read_lptc_array(fname):
    with open(os.path.join(LPTC_SYSFS, fname), "rt") as f:
        line = f.readline().strip()
    return [int(x, 16) for x in line.split(",")]


def check_sysfs_available():
    if not os.path.isdir(LPTC_SYSFS):
        sys.stderr.write(
            f"Error: {LPTC_SYSFS} not found.\n"
            f"The 'gpstime' kernel module does not appear to be loaded, "
            f"or no LIGO PCIe Timing Card is attached.\n"
        )
        sys.exit(1)


# ---------- bitfield rendering ----------
# kind: 'ok' (set = healthy, unset = concerning),
#       'alarm' (set = problem),
#       'info' (neutral; only displayed when set)
Flag = namedtuple("Flag", ['name', 'inverted', 'kind'],
                  defaults=[False, 'info'])


def _print_section(title):
    print(_c(f"\n=== {title} ===", 'bold'))


def _print_bitfield(label, value, flags, verbose):
    """Print decoded bitfield, grouping flags by severity."""
    set_ok, missing_ok, set_alarm, set_info, unset_info = [], [], [], [], []
    for mask, _flag in flags.items():
        if isinstance(_flag, str):
            flag = Flag(_flag)
        else:
            flag = _flag
        active = bool(value & mask) != bool(flag.inverted)
        if flag.kind == 'alarm':
            (set_alarm if active else unset_info).append(flag.name)
        elif flag.kind == 'ok':
            (set_ok if active else missing_ok).append(flag.name)
        else:
            (set_info if active else unset_info).append(flag.name)

    print(f"  {label}:")
    if verbose:
        print(f"    raw = 0x{value:08x}")
    if set_alarm:
        print(f"    {_c('ALARM:', 'red')} {', '.join(set_alarm)}")
    if missing_ok:
        print(f"    {_c('not active:', 'yellow')} {', '.join(missing_ok)}")
    if set_ok:
        print(f"    {_c('ok:', 'green')} {', '.join(set_ok)}")
    if set_info:
        print(f"    set: {', '.join(set_info)}")
    if verbose and unset_info:
        print(f"    {_c('unset: ' + ', '.join(unset_info), 'dim')}")


# ---------- LPTC status ----------
def process_lptc_status(status, verbose):
    """Decodes Status and interrupt control register defined in
    DCC doc 2000406"""
    flags = {
        0x80000000: Flag('OK_LOCKED', kind='ok'),
        0x40000000: Flag('ROOT_NODE'),
        0x20000000: Flag('FANOUT_PORTS'),
        0x10000000: Flag('UPLINK_UP', kind='ok'),
        0x8000000:  Flag('LOSS_OF_SIGNAL', kind='alarm'),
        0x4000000:  Flag('OCXO_LOCKED', kind='ok'),
        0x2000000:  Flag('GPS_LOCKED', kind='ok'),
        0x1000000:  Flag('VCXO_VOLT_OUT_OF_RANGE', kind='alarm'),
        0x800000:   Flag('UTC_MODE'),
        0x400000:   Flag('LEAP_SEC_DECODE'),
        0x200000:   Flag('LEAP_DEC_PENDING'),
        0x100000:   Flag('LEAP_INC_PENDING'),
    }
    _print_bitfield("LPTC status", status, flags, verbose)

    leap_decode = bool(status & 0x400000)
    leap_seconds = (status & 0xff00) >> 8
    if leap_decode or verbose:
        print(f"  TAI-UTC offset (leap seconds): {leap_seconds}")

    active_msi = [i for i in range(4) if status & (1 << i)]
    if active_msi:
        print(f"  active MSI interrupts: {', '.join(str(i) for i in active_msi)}")
    elif verbose:
        print("  active MSI interrupts: none")


def process_backplane_status(status, verbose):
    flags = {
        0x200:  Flag("BP_PRESENT", kind='ok'),
        0x100:  Flag("X5_INPUT"),
        0x80:   Flag("X3_INPUT"),
        0x40:   Flag("X1_INPUT"),
        0x20:   Flag("TEMP_ALARM", kind='alarm'),
        0x2:    Flag("CLOCKS_RUNNING", kind='ok'),
        0x1:    Flag("CLOCKS_ACTIVE", kind='ok'),
    }
    _print_bitfield("backplane status", status, flags, verbose)
    bp_rev = (status & 0x18) >> 3
    print(f"  backplane revision: {bp_rev}")


def process_backplane_config(config, verbose):
    flags = {
        0x10:   "START_NEXT_IDLE",
        0x8:    "START_NEXT_SECOND",
        0x4:    Flag("GLOBAL_ENABLE", kind='ok'),
        0x2:    "WD_RESET_ON_READ",
        0x1:    "DISABLE_DUOTONE_OUTPUT",
    }
    _print_bitfield("backplane config", config, flags, verbose)


def process_duotone_shift(shift):
    print(f"  duotone shift: 0x{shift:x}")


def process_duotone_config(can_configure, amplitude, frequency_code):
    freqs = [
        "960, 961",
        "1920, 1921",
        "3840, 3841",
        "15424, 15423",
    ]
    if can_configure:
        print("  duotone: configurable")
        print(f"    amplitude: 0x{amplitude:x}")
        if 0 <= frequency_code < len(freqs):
            print(f"    frequencies: {freqs[frequency_code]} Hz")
        else:
            print(f"    {_c(f'unknown frequency code {frequency_code}', 'yellow')}")
    else:
        print("  duotone: not configurable")


# ---------- onboard ----------
def process_board_conf(conf):
    exp = conf & 0xf  # exponent for exponential method
    div = (conf & 0xff0) >> 4
    if exp > 0:
        in_clock_hz = 2**(exp + 10)
    else:
        in_clock_hz = 2**26 / (div + 1)
    print(f"  external sync frequency: {in_clock_hz} Hz")


def process_board_and_powersupply_status(status, verbose):
    flags = {
        0x100: Flag('SWTCH_REG_INT', kind='alarm'),
        0x4:   Flag('SWTCH_SPLY_TEMP_INT', kind='alarm'),
        0x8:   Flag('SWTCH_SPLY_IN_V_INT', kind='alarm'),
        0x10:  Flag('SWTCH_SPLY_1_INT', kind='alarm'),
        0x20:  Flag('SWTCH_SPLY_2_INT', kind='alarm'),
        0x40:  Flag('SWTCH_SPLY_3_INT', kind='alarm'),
        0x80:  Flag('SWTCH_SPLY_4_INT', kind='alarm'),
        0x2:   Flag('GBIT_XCVR_PWR_GOOD', kind='ok'),
        0x1:   Flag('SWTCH_SPLY_PWR_GOOD', kind='ok'),
    }
    for i in range(16):
        flags[0x10000 << i] = Flag(f"DIP_{i+1}", inverted=True)
    _print_bitfield("board/PSU status", status, flags, verbose)


def process_xadc_status(status, verbose):
    flags = {
        0x20: Flag('XADC_ENBLD', kind='ok'),
        0x10: Flag('VCCAUX_ALARM', kind='alarm'),
        0x8:  Flag('VCCINT_ALARM', kind='alarm'),
        0x4:  Flag('USER_TEMP_ALARM', kind='alarm'),
        0x2:  Flag('OVER_TEMP_ALARM', kind='alarm'),
        0x1:  Flag('ALARM', kind='alarm'),
    }
    _print_bitfield("XADC", status, flags, verbose)


def process_temperature(temp):
    temp_c = (temp/65536.0) * 503.975 - 273.15
    color = None
    flag = ""
    if temp_c >= 85:
        color, flag = 'red', "  [HOT]"
    elif temp_c >= 70:
        color, flag = 'yellow', "  [warm]"
    print(f"  temperature: {_c(f'{temp_c:.1f} °C', color)}{flag}")


# ---------- voltage / current ----------
def calc_volts(raw, word_offset, gain, offset):
    counts = (raw >> (word_offset * 16)) & 0xffff
    return (counts / (2**16) / gain) + offset


# represents a single 16 bit voltage reading.
# name = display name for voltage
# register_offset = added to base address for the reading to get register.
#  A voltage from 'external power supply 8' has a register_offset of 7
# word_offset: number of 16 bit words to offset reading, with lowest being 1.
# gain: analog gain.
# nominal: approximately what voltage should be
VoltageReading = namedtuple("VoltageReading",
                            ['name', 'register_offset', 'word_offset',
                             'gain', 'offset', 'nominal'],
                            defaults=[0, 0],
                            )
CurrentReading = namedtuple("CurrentReading",
                            ['name', 'register_offset', 'word_offset',
                             'gain', 'offset'],
                            defaults=[-0.05, ])


def _volt_color_and_flag(volts, nominal):
    if nominal == 0:
        return None, ""
    err = abs(volts - nominal) / abs(nominal)
    if err > VOLTAGE_TOLERANCE:
        return 'red', "  [OUT OF RANGE]"
    if err > VOLTAGE_TOLERANCE / 2:
        return 'yellow', ""
    return 'green', ""


def print_volts(name, volts, nominal):
    color, flag = _volt_color_and_flag(volts, nominal)
    line = (f"    {name:<8} {volts:+7.2f} V  "
            f"(nominal {float(nominal):+.1f} V){flag}")
    print(_c(line, color))


def print_amps(name, amps):
    print(f"    {name:<8} {amps:+7.3f} A")


def calc_and_print_volts(vr, vals):
    # same calc used for volts and amps
    volts = calc_volts(vals[vr.register_offset], vr.word_offset, vr.gain, vr.offset)
    if isinstance(vr, VoltageReading):
        print_volts(vr.name, volts, vr.nominal)
    elif isinstance(vr, CurrentReading):
        print_amps(vr.name, volts)
    else:
        raise Exception("Unknown A2D reading type")


def process_voltage_readings(vrs, vals):
    for vr in vrs:
        calc_and_print_volts(vr, vals)


def read_and_process_internal_power():
    print("  Internal Power Supply:")
    vrs = [
        VoltageReading("VCCINT", register_offset=0, word_offset=1, gain=1/3, nominal=1),
        VoltageReading("VCCINT", register_offset=1, word_offset=0, gain=1/3, nominal=1.8),
        VoltageReading("VCCINT", register_offset=1, word_offset=1, gain=1/3, nominal=1),
    ]
    vals = read_lptc_array('internal_pwr')
    process_voltage_readings(vrs, vals)


def read_and_process_external_power():
    print("  External Power Supply:")
    vrs = [
        CurrentReading("3.3V",   register_offset=0, word_offset=0, gain=1),
        CurrentReading("VCCINT", register_offset=0, word_offset=1, gain=1),
        CurrentReading("VCCAUX", register_offset=1, word_offset=0, gain=1),
        CurrentReading("2.5V",   register_offset=1, word_offset=1, gain=1/3),
        VoltageReading("VREG",   register_offset=2, word_offset=0, gain=1/6,  nominal=5.1),
        VoltageReading("VDD",    register_offset=2, word_offset=1, gain=1/4,  nominal=2.5),
        VoltageReading("AVCC",   register_offset=3, word_offset=0, gain=2/3,  nominal=1.0),
        VoltageReading("AVTT",   register_offset=3, word_offset=1, gain=2/3,  nominal=1.2),
        VoltageReading("P5",     register_offset=4, word_offset=0, gain=1/6,  nominal=5),
        VoltageReading("N5",     register_offset=4, word_offset=1, gain=1/6,  offset=-6.25, nominal=-5),
        VoltageReading("VCC",    register_offset=5, word_offset=0, gain=1/5,  nominal=3.3),
        VoltageReading("N12",    register_offset=5, word_offset=1, gain=1/15, offset=-17.5, nominal=-12),
        VoltageReading("VADC",   register_offset=6, word_offset=0, gain=1/2,  nominal=1.8),
        VoltageReading("P10",    register_offset=6, word_offset=1, gain=1/11, nominal=10),
        VoltageReading("V12",    register_offset=7, word_offset=0, gain=1/15, nominal=12),
        CurrentReading("V12",    register_offset=7, word_offset=1, gain=1),
    ]
    vals = read_lptc_array('external_pwr')
    process_voltage_readings(vrs, vals)


# ---------- slots ----------
def slot_config_freq(config):
    freq = config & 0xff
    if freq > 127:
        freq = 256 - freq
    return 2**freq


def process_slot_summary(slot, status, config, phase):
    """Compact one-line per-slot summary."""
    enabled = bool(config & 0x100)
    running = bool(status & 0x2)
    active = bool(status & 0x1)
    freq = slot_config_freq(config)
    phase_deg = phase * 360.0 / 2**32

    state = []
    state.append(_c("ENABLED", 'green') if enabled else _c("disabled", 'dim'))
    state.append(_c("running", 'green') if running else _c("not-running", 'yellow'))
    state.append(_c("active", 'green') if active else _c("inactive", 'dim'))
    print(f"  slot {slot:2d}: {' '.join(state)}  {freq} Hz  phase {phase_deg:.3f}°")


def process_slot_status(status, verbose):
    flags = {
        0x400000:   'BIT_2',
        0x100000:   'BIT_1',
        0x80000:    'DUOTONE_LAST_DAC',
        0x40000:    'DUOTONE_PENULT_ADC',
        0x20000:    'DUOTONE_LAST_ADC',
        0x2:        Flag('CLOCK_RUNNING', kind='ok'),
        0x1:        Flag('CLOCK_ACTIVE', kind='ok'),
    }
    _print_bitfield("slot status", status, flags, verbose)


def process_slot_config(config, verbose):
    flags = {
        0x400000:   'BIT_2_SET',
        0x200000:   'BIT_2_BIN',
        0x100000:   'BIT_1_SET',
        0x80000:    'BIT_1_BIN',
        0x40000:    'DUOTONE_PENULT_ADC',
        0x20000:    'DUOTONE_LAST_ADC',
        0x10000:    'LVDS_CLOCK_LINES',
        0x1000:     'IDLE_CLOCK_PULLUP',
        0x800:      'CLOCK_NEXT_IDLE',
        0x400:      'CLOCK_NEXT_SEC',
        0x200:      'CLOCK_INV',
        0x100:      Flag('CLOCK_ENABLE', kind='ok'),
    }
    _print_bitfield("slot config", config, flags, verbose)
    print(f"  frequency: {slot_config_freq(config)} Hz")


def process_slot_phase(phase):
    phase_deg = phase * 360.0 / 2**32
    print(f"  phase: {phase_deg:.3f}°")


# ---------- printers ----------
def _board_id_str(rawid):
    if (rawid >> 4) == 0x2000329:
        return f"board id {rawid & 0xf}"
    return f"raw board id 0x{rawid:x} (unrecognized)"


def _software_id_str(sid, rev):
    releases = {
        (1, 0x1f0): "Aug 2021",
        (2, 0x2f4): "2023 Feb 9",
    }
    if (sid >> 4) != 0x2000337:
        return f"raw software id 0x{sid:x} rev 0x{rev:x} (unrecognized)"
    sid_short = sid & 0xf
    release = releases.get((sid_short, rev), "version not recognized")
    return f"sw {sid_short} rev 0x{rev:x} ({release})"


def print_card_header():
    """Always-shown identification line."""
    try:
        rawid = read_lptc_int("board_id")
        sid = read_lptc_int("software_id")
        rev = read_lptc_int("software_rev")
    except IOError as e:
        print(_c(f"Could not identify card: {e}", 'red'))
        return
    print(_c(
        f"LIGO PCIe Timing Card  [{_board_id_str(rawid)}, {_software_id_str(sid, rev)}]",
        'bold'))


def print_health_summary():
    """One-line top-of-output verdict."""
    try:
        status = read_lptc_int("status")
    except IOError:
        return
    locked = bool(status & 0x80000000)
    gps = bool(status & 0x2000000)
    ocxo = bool(status & 0x4000000)
    los = bool(status & 0x8000000)
    vcxo_oor = bool(status & 0x1000000)

    items = []
    items.append(_c("LOCKED", 'green') if locked else _c("UNLOCKED", 'red'))
    items.append(_c("GPS", 'green') if gps else _c("no GPS", 'yellow'))
    items.append(_c("OCXO", 'green') if ocxo else _c("no OCXO", 'yellow'))
    if los:
        items.append(_c("LOSS_OF_SIGNAL", 'red'))
    if vcxo_oor:
        items.append(_c("VCXO_OUT_OF_RANGE", 'red'))

    healthy = locked and gps and ocxo and not los and not vcxo_oor
    verdict = _c("[OK]", 'green') if healthy else _c("[CHECK]", 'yellow')
    print(f"Status: {' · '.join(items)}    {verdict}")


def print_lptc_status(verbose):
    _print_section("LPTC Status")
    try:
        process_lptc_status(read_lptc_int("status"), verbose)
        process_backplane_status(read_lptc_int("backplane_status"), verbose)
        process_backplane_config(read_lptc_int("backplane_config"), verbose)
        process_duotone_config(
            read_lptc_int("duotone_can_config"),
            read_lptc_int("duotone_amplitude"),
            read_lptc_int("duotone_frequency"),
        )
        process_duotone_shift(read_lptc_int("duotone_shift"))
    except IOError as e:
        print(_c(f"  Could not read LPTC sysfs: {e}", 'red'))


def print_onboard_status(verbose):
    _print_section("Onboard")
    try:
        process_board_conf(read_lptc_int("brd_synch_factors"))
        process_board_and_powersupply_status(
            read_lptc_int("board_and_powersupply_status"), verbose)
        process_xadc_status(read_lptc_int("xadc_status"), verbose)
        process_temperature(read_lptc_int("temp"))
        read_and_process_internal_power()
        read_and_process_external_power()
    except IOError as e:
        print(_c(f"  Could not read LPTC sysfs: {e}", 'red'))


def print_slots(verbose):
    _print_section("Slots")
    for slot in range(1, 11):
        try:
            status = read_lptc_int(f"slot_{slot}/status")
            config = read_lptc_int(f"slot_{slot}/config")
            phase = read_lptc_int(f"slot_{slot}/phase")
        except IOError as e:
            print(_c(f"  slot {slot}: cannot read sysfs ({e})", 'red'))
            continue
        if verbose:
            print(_c(f"\n  --- slot {slot} ---", 'bold'))
            process_slot_status(status, verbose)
            process_slot_config(config, verbose)
            process_slot_phase(phase)
        else:
            process_slot_summary(slot, status, config, phase)


def main():
    parser = argparse.ArgumentParser(
        description="Show info about an attached LIGO PCIe Timing Card")
    parser.add_argument('--info', '-i', action='store_true',
                        help="show basic info, also shown when no args given")
    parser.add_argument('--onboard', '-b', action='store_true',
                        help="show on-board info")
    parser.add_argument('--slots', '-s', action='store_true',
                        help="show status and configuration for each slot")
    parser.add_argument('--all', '-a', action='store_true',
                        help="show all info")
    parser.add_argument('--verbose', '-v', action='store_true',
                        help="show raw register hex and all flags including unset")
    parser.add_argument('--color', choices=['auto', 'always', 'never'], default='auto',
                        help="colorize output (default: auto)")
    args = parser.parse_args()

    _setup_color(args.color)
    check_sysfs_available()

    print_card_header()
    print_health_summary()

    show_info = (not args.slots and not args.onboard) or args.all or args.info
    if show_info:
        print_lptc_status(args.verbose)
    if args.onboard or args.all:
        print_onboard_status(args.verbose)
    if args.slots or args.all:
        print_slots(args.verbose)


if __name__ == "__main__":
    main()
