#!/usr/bin/python3

from pathlib import Path
import re
from collections import namedtuple
import sys
import os
import argparse

VERSION = (1,0,0)

try:
    ifo = os.environ["ifo"]
except KeyError:
    ifo = None

try:
    IFO = os.environ["IFO"]
except KeyError:
    IFO = None

if IFO and not ifo:
    ifo = IFO.lower()

argparser = argparse.ArgumentParser(description = "Show info on installed cards for LIGO front end servers, CyMACs, and test stands.")

argparser.add_argument("-v","--version",action='store_true')
args = argparser.parse_args()

if args.version:
    print(".".join([str(i) for i in VERSION]))
    sys.exit(0)

pci_path = Path("/sys/bus/pci/devices")

dev_path_re = re.compile(r"^([0-9a-fA-F]{4}):([0-9a-fA-F]{2}):([0-9a-fA-F]{2})\.([0-9a-fA-F])$")


class DeviceID(namedtuple('DeviceID', ['vendor', 'device'])):
    def __repr__(self):
        return "DeviceID(vendor=0x%04x, device=0x%04x)" % self

    def __str__(self):
        return "(%04x %04x)" % self


vendor_database = {}
device_database = {}
subsystem_database = {}

device_database_files = ["/usr/share/misc/pci.ids", "/usr/share/showcards/pci.ids",
                        "/home/controls/pci/pci.ids",
                         "/opt/rtcds/userapps/release/cds/h1/scripts/pci.ids",
                         "/opt/rtcds/userapps/release/cds/l1/scripts/pci.ids"]
if ifo:
    device_database_files += [f"/opt/rtcds/userapps/release/cds/{ifo}/pci.ids",
                              f"/opt/rtcds/userapps/release/cds/{ifo}/scripts/pci.ids",]

db_vend_re = re.compile(r"^([0-9a-fA-F]{4})\s+(\S.*\S)\s*$")
db_device_re = re.compile(r"^\t([0-9a-fA-F]{4})\s+(\S.*\S)\s*$")
db_subsystem_re = re.compile(r"^\t\t([0-9a-fA-F]{4})\s+([0-9a-fA-F]{4})\s+(\S.*\S)\s*$")


def scan_device_database_file(fname):
    with open(fname, "rt") as f:
        current_vendor = None
        current_device = None
        block = f.read()
        lines_n = block.split("\n")
        lines = []
        for ln in lines_n:
            lines += ln.split("\r")

        for line_raw in lines:
            line = line_raw.split("#")[0]
            if current_device:
                m = db_subsystem_re.search(line)
                if m:
                    subsystem_id = DeviceID(int(m.group(1), 16), int(m.group(2), 16))
                    subsystem_database[current_device][subsystem_id] = m.group(3)
                    continue
            if current_vendor:
                m = db_device_re.search(line)
                if m:
                    dev_id = DeviceID(current_vendor, int(m.group(1), 16))
                    device_database[dev_id] = m.group(2)
                    subsystem_database[dev_id] = {}
                    current_device = dev_id
                    continue
            m = db_vend_re.search(line)
            if m:
                current_device = None
                vend_id = int(m.group(1), 16)
                vendor_database[vend_id] = m.group(2)
                current_vendor = vend_id
                continue


for filename in device_database_files:
    try:
        scan_device_database_file(filename)
    except IOError:
        pass

extender_ids = [DeviceID(0x10b5, 0x8606), DeviceID(0x10b5, 0x8608), DeviceID(0x10b5, 0x8604),
                DeviceID(0x111d, 0x8032)]
skip_2_ids = [DeviceID(0x10b5, 0x8606)]  # these devices have a phony slot 2
bridge_ids = [DeviceID(0x10b5, 0x8111), DeviceID(0x10b5, 0x9056), DeviceID(0x1a03, 0x1150)]

# not shown as cards
ignored_ids = [DeviceID(0x11c8, 0x1611), DeviceID(0x11c8, 0x2611)  # extra dolphins ids,
               ] + bridge_ids + extender_ids

card_names = {
    DeviceID(0x10b5, 0x3101): "General Standards 16AI64",
    DeviceID(0x10b5, 0x3120): "General Standards 16AO16",
    DeviceID(0x1221, 0x8682): "Contec Binary I/O 64",
    DeviceID(0x1221, 0x8632): "Contec Binary I/O 16",
    DeviceID(0x10ee, 0xd8c6): "LIGO Timing Board",
    DeviceID(0x10b5, 0x3574): "General Standards 20A08",
    DeviceID(0x10b5, 0x3357): "General Standards 18AO8",
}


def get_card_name(card_dev):
    if card_dev.id in device_database and card_dev.subsystem_id in subsystem_database[card_dev.id]:
        return subsystem_database[card_dev.id][card_dev.subsystem_id]
    if card_dev.subsystem_id in device_database:
        return device_database[card_dev.subsystem_id]
    if card_dev.subsystem_id in card_names:
        return card_names[card_dev.subsystem_id]
    return "Unknown device"


def get_vendor(vid):
    if vid in vendor_database:
        return vendor_database[vid]
    return "Unknown vendor"


def find_device_paths(p, parent=None):
    devices = [Device(x, parent) for x in p.iterdir() if dev_path_re.search(x.name)]
    return devices


def read_hex_word_from_text_file(p):
    with open(str(p), "rt") as f:
        text = f.read().strip()
        return int(text, 16)


def is_interesting_device(did):
    return did.vendor not in [0xaaaa, 0x0, 0xffff,0x15d9,0x8086] and did.device not in [0x0, 0xffff] and did not in ignored_ids


class Device(object):
    subsystem_id = "None"

    def __init__(self, path, parent=None):
        self.path = Path(path)
        self.parent = parent
        self.pci_address = self.path.name

        m = dev_path_re.search(path.name)
        if m:
            self.bus_num = tuple([int(x, 16) for x in m.groups()])

        # read some info from files in path
        vendor_id = read_hex_word_from_text_file(path / "vendor")
        device_id = read_hex_word_from_text_file(path / "device")
        self.id = DeviceID(vendor_id, device_id)

        subsystem_vendor = read_hex_word_from_text_file(path / "subsystem_vendor")
        subsystem_device = read_hex_word_from_text_file(path / "subsystem_device")
        self.subsystem_id = DeviceID(subsystem_vendor, subsystem_device)

        self.children = find_device_paths(self.path, self)
        self.children.sort()

        self.skip_slot_2 = self.id in skip_2_ids  # (slot 2 is a fantom slot)

    @property
    def vendor(self):
        return self.id.vendor

    @property
    def device(self):
        return self.id.device

    @property
    def bus_number(self):
        return "%04x:%02x:%02x.%1x" % self.bus_num

    @property
    def pci_num(self):
        return self.bus_num[1]

    def __repr__(self):
        return f"Device({self.path})"

    def has_one_as_child(self, bus_numbers):
        """return true if at least one of bus_nums is a child (recursively) of self.

        used to rule out extenders that are parents of other extenders
        we only want to be looking at each extender once."""

        for child in self.children:
            if child.bus_num in bus_numbers:
                return True
            if child.has_one_as_child(bus_numbers):
                return True

        return False

    def get_terminal_devices(self):
        """
        :return: a list in address order of devices with interesting subsystem DeviceIDs
        """
        if is_interesting_device(self.subsystem_id):
            ret = [self]
        else:
            ret = []
        for child in self.children:
            ret += child.get_terminal_devices()
        return ret

    def get_slots(self):
        """
        Return the terminal devices from each slot of an extender
        """

        slot_list = []
        for i in range(len(self.children)):
            if self.skip_slot_2 and i == 1:
                continue
            child = self.children[i]
            slot_list.append(child.get_terminal_devices())
        return slot_list

    def __lt__(self, other):
        return self.bus_num < other.bus_num

    def __eq__(self, other):
        return self.bus_num == other.bus_num

    def __hash__(self):
        return self.bus_num[3] + self.bus_num[2] * (16 << 1) + self.bus_num[1] * (16 << 3) + self.bus_num[0] * (16 << 5)


def treeify_device_list(dev_list):
    dev_list.sort()
    ret = []
    found = {}

    def fillchildren(dev):
        found[dev] = True
        for c in dev.children:
            fillchildren(c)

    for dev in dev_list:
        if dev not in found:
            ret.append(dev)
            fillchildren(dev)

    return ret


def top_level_extenders(dev_tree):
    """
    Return an ordered list of top level extenders.

    :param dev_tree:
    :return:
    """
    ret = []
    for dev in dev_tree:
        if dev.id in extender_ids:
            ret.append(dev)
        else:
            ret += top_level_extenders(dev.children)
    return ret


root_devices = find_device_paths(pci_path)
dev_tree = treeify_device_list(root_devices)
# dev_tree = root_devices

# these are the cards in the server
top_extenders = top_level_extenders(dev_tree)
extender_cards = []
for ext in top_extenders:
    card = ext
    while card.parent is not None:
        card = card.parent
    extender_cards.append(card)


def find_large_extender(devs):
    ret = []
    for dev in devs:
        if dev.id in extender_ids and len(dev.children) > 3:
            ret.append(dev)
        else:
            ret += find_large_extender(dev.children)
    return ret


print(f"Found {len(top_extenders)} bus extenders")

if len(top_extenders) > 0:
    fb = top_extenders[0]

    sys.stdout.write(f"Extender type {fb.id} {get_card_name(fb)}.  ")
    if fb.id == DeviceID(0x10b5, 0x8608):
        print("Possible V4 or V5 chassis")
    elif fb.id == DeviceID(0x111d, 0x8032):
        print("Possible V3 or earlier chassis")
    else:
        print("Unrecognized chassis")

board_num = 1
ext_num = 1

for board in dev_tree:
    if board in extender_cards:
        print(f"Card {board_num}")
        extenders = find_large_extender([board])
        if len(extenders) == 0:
            print(f"\tno expansion board found.  Disconnected?")
        else:
            for ext in extenders:
                print(f"    Extender {ext_num}")
                slots = ext.get_slots()
                for slot_num in range(len(slots)):
                    sys.stdout.write(f"        slot {slot_num + 1}: ")
                    cards = slots[slot_num]
                    if len(cards) == 0:
                        sys.stdout.write(f"empty")
                    elif len(cards) == 1:
                        name = get_card_name(cards[0])
                        vendor = get_vendor(cards[0].subsystem_id.vendor)
                        sys.stdout.write(f"{cards[0].subsystem_id} {vendor},{name}")
                    else:
                        sys.stdout.write(f"ERROR, found more than one card!")
                        for card in cards:
                            name = get_card_name(card)
                            vendor = get_vendor(card.subsystem_id.vendor)
                            sys.stdout.write(f" [{card.subsystem_id} {vendor}, {name}]")
                    sys.stdout.write("\n")

                ext_num += 1
        board_num += 1
    else:  # not an extender
        term_devs = board.get_terminal_devices()
        if len(term_devs) > 0:
            print(f"Card {board_num}")
            for tdev in term_devs:
                name = get_card_name(tdev)
                vendor = get_vendor(tdev.subsystem_id.vendor)
                sys.stdout.write(f"    {tdev.subsystem_id} {vendor},{name}\n")
            board_num += 1

