#!/usr/bin/python3
################################################################################
# pci:
#   used to return information about pci cards in a .json
#   format. This is a helper script for use with the
#   cockpit-hardware package (https://github.com/45Drives/cockpit-hardware)
#
# Copyright (C) 2020, Mark Hooper   <mhooper@45drives.com>
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
################################################################################

import subprocess
import re
import json
import os

g_dmi_fields = [
    "Designation",
    "Type",
    "Current Usage",
    "ID",
    "Bus Address",
]

g_storcli64_fields = [
    "SAS9305-16i",
    "SAS9305-24i",
    "AVAGOMegaRAIDSAS9361-16i",
]

g_network_card_models = [
    "X540-AT2",
    "X550",
    "XL710",
    "XXV710",
    "82599ES",
    "BCM57412",
    "MT27800",
]

g_sata_controllers = [
    "ASM1062",
    # "ASM1164"
]

BA_EPC612D8A = {
    "ff00:16:00.0": "0000:17:00.0",
    "ff00:16:02.0": "0000:1c:00.0",
    "ff00:64:00.0": "0000:65:00.0",
    "ff00:64:02.0": "0000:66:00.0",
    "ff00:b2:00.0": "0000:b3:00.0",
    "ff00:b2:02.0": "0000:b4:00.0",
}

PCIE_SLOT_TYPE_GIGABYTE_ME03 = {
    "PCIE_1": "PCI Express 4 x16",
    "PCIE_2": "PCI Express 4 x8",
    "PCIE_3": "PCI Express 4 x16",
    "PCIE_4": "PCI Express 5 x16",
    "PCIE_5": "PCI Express 5 x16",
    "PCIE_6": "PCI Express 5 x16",
}

BA_LUT_ME03 = {
    "PCIE_1": "0000:c2:00.0",
    "PCIE_2": "0000:c4:00.0",
    "PCIE_3": ["0000:01:00.0", "0000:02:00.0"],
    "PCIE_4": "0000:41:00.0",
    "PCIE_5": ["0000:01:00.0", "0000:02:00.0"],
    "PCIE_6": "0000:c1:00.0",
}

PCIE_SLOT_TYPE_GIGABYTE_MZ73 = {
    "PCIE_1": "PCI Express 5 x16",
    "PCIE_2": "PCI Express 5 x16",
    "PCIE_3": "PCI Express 5 x16",
    "PCIE_4": "PCI Express 5 x16",
}

BA_LUT_MZ73 = {
    "PCIE_1": ["0000:a2:00.0"],
    "PCIE_2": ["0000:81:00.0", "0000:84:00.0"],
    "PCIE_3": ["0000:01:00.0"],
    "PCIE_4": ["0000:41:00.0"],
    "M2_0": ["0000:61:00.0"],
}

PCIE_SLOT_TYPE_GIGABYTE_MH53 = {
    "PCIE_1": "PCI Express 5 x16",
    "PCIE_2": "PCI Express 5 x16",
    "PCIE_3": "PCI Express 5 x16",
    "PCIE_5": "PCI Express 5 x1",
    "PCIE_6": "PCI Express 5 x16",
    "PCIE_7": "PCI Express 5 x16",
}

MH53_AUX_SLOT_TYPE = {
    "NVME0": "NVMe Slot (PCI Express 5 x8)",
    "M2_0": "M.2 Slot (PCI Express 5 x4)",
    "M2_1": "M.2 Slot (PCI Express 5 x4)",
    "M2_2": "M.2 Slot (PCI Express 5 x4)",
    "M2_3": "M.2 Slot (PCI Express 5 x4)",
}

MH53_SLOT_ORDER = [
    "PCIE_1",
    "PCIE_2",
    "PCIE_3",
    "PCIE_5",
    "PCIE_6",
    "PCIE_7",
    "NVME0",
    "M2_0",
    "M2_1",
    "M2_2",
    "M2_3",
]

BA_LUT_MH53 = {
    "PCIE_1": ["0000:01:00.0", "0000:00:01.1", "0000:00:01.2", "0000:00:01.3", "0000:00:01.4"],
    "PCIE_2": ["0000:41:00.0", "0000:40:01.2", "0000:40:01.3", "0000:40:01.4"],
    "PCIE_3": ["0000:81:00.0", "0000:80:01.1", "0000:80:01.2", "0000:80:01.3", "0000:80:01.4"],
    "PCIE_5": ["0000:21:00.0"],
    "PCIE_6": ["0000:c1:00.0", "0000:c0:01.2", "0000:c0:01.3", "0000:c0:01.4"],
    "PCIE_7": ["0000:e1:00.0", "0000:e0:01.2", "0000:e0:01.3", "0000:e0:01.4"],
    "NVME0": ["0000:61:00.0", "0000:60:01.1", "0000:60:01.2"],
    "M2_0": ["0000:a0:01.1"],
    "M2_1": ["0000:a0:01.2"],
    "M2_2": ["0000:a0:01.3"],
    "M2_3": ["0000:a0:01.4"],
}

PCIE_SLOT_LUT_B550I = {
    "J10": "PCIEX16",
    "J3700 M.2 Slot": "M2A_CPU",
    "J3708 PCIE x4 slot from Promontory": "M2B_SB",
}

BA_LUT_B550I = {
    "PCIEX16": ["0000:01:00.0", "0000:01:00.1"],
}

NETWORK_BA_LUT_B550I = {
    "00:01.0": "01:00.0",
    "00:01.1": "01:00.1",
}

# Gigabyte MW34-SP0-00 (Intel W680 chipset)
# dmidecode -t 9 bus addresses are completely wrong on this board.
# Root port mapping confirmed across 5 lspci-tv runs:
#   Run 1 (4 NVMe + HBA + NIC):         00:01.0→HBA, 00:1d.0→NIC, 00:1d.4→M2M_1
#   Run 2 (3 NVMe + HBA + NIC):         00:01.1→HBA, 00:01.0→empty, no 00:1d.4
#   Run 3 (only M2M_CPU):               no 00:01.x, no 00:1d.x, no 00:1a/1b
#   Run 4 (NIC + HBA + M2M_CPU+M2M_3): 00:01.1→HBA, 00:01.0→empty, 00:1a.0→Sandisk
#   Run 5 (NIC + 2×HBA + M2M_CPU):     00:01.0→MegaRAID, 00:01.1→SAS3416, 00:1d.0→NIC
# All root ports vanish entirely when their slot is empty.
# Only PCIE_3/PCIE_6 share the 00:01.x group (CPU lanes, dynamic numbering).
# When both are populated, BIOS assigns 00:01.0 and 00:01.1 — order is not guaranteed.
BA_LUT_MW34 = {
    "PCIE_1":  ["0000:00:1d.0"],                    # PCH — dedicated root port
    "M2M_1":   ["0000:00:1d.4"],                    # PCH — dedicated, absent when empty
    "PCIE_3":  ["0000:00:01.0", "0000:00:01.1"],   # CPU — shares 00:01.x with PCIE_6
    "PCIE_6":  ["0000:00:01.0", "0000:00:01.1"],   # CPU — shares 00:01.x with PCIE_3
    "M2M_CPU": ["0000:00:06.0"],                    # CPU — dedicated root port
    "M2M_2":   ["0000:00:1b.4"],                    # PCH — dedicated root port
    "M2M_3":   ["0000:00:1a.0"],                    # PCH — dedicated root port
}


def run_command(args):
    try:
        result = subprocess.run(
            args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True,
            check=False,
        )
        return result.stdout or ""
    except Exception:
        return ""


def get_motherboard_model_server_info():
    json_path = "/etc/45drives/server_info/server_info.json"
    default_info = {
        "Motherboard": [{
            "Manufacturer": "?",
            "Product Name": "?",
            "Serial Number": "?",
        }]
    }

    if not os.path.exists(json_path):
        return default_info

    try:
        with open(json_path, "r", encoding="utf-8") as f:
            si = json.load(f)
        return {
            "Motherboard": [{
                "Manufacturer": si.get("Motherboard", {}).get("Manufacturer", "?"),
                "Product Name": si.get("Motherboard", {}).get("Product Name", "?"),
                "Serial Number": si.get("Motherboard", {}).get("Serial Number", "?"),
            }]
        }
    except Exception:
        return default_info


def parse_lspci_output():
    lspci_output = run_command(["lspci", "-mm"])
    bus_address_pattern = re.compile(r"^([0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9])")
    parsed_data = {}

    motherboard_info = get_motherboard_model_server_info()
    motherboard_model = motherboard_info["Motherboard"][0].get("Product Name", "")
    exclusion_strings = ["BCM57416", "Non-Essential Instrumentation"]

    for line in lspci_output.splitlines():
        match = bus_address_pattern.match(line)
        if not match:
            continue

        bus_address = "0000:" + match.group(1)
        rest_of_line = re.findall(r'"[^"]*"|[\S]+', line)[1:]

        if motherboard_model.startswith("MZ73") and any(
            exclusion_string in " ".join(rest_of_line)
            for exclusion_string in exclusion_strings
        ):
            continue

        if rest_of_line:
            parsed_data[bus_address] = rest_of_line

    return parsed_data


def _is_lspci_endpoint(device_info):
    if not device_info:
        return False
    cls = str(device_info[0]).strip('"').lower()
    return "bridge" not in cls


def _pcie_gen_from_speed(speed_gts):
    try:
        speed = float(speed_gts)
    except (TypeError, ValueError):
        return None
    if speed >= 64:
        return 6
    if speed >= 32:
        return 5
    if speed >= 16:
        return 4
    if speed >= 8:
        return 3
    if speed >= 5:
        return 2
    if speed >= 2.5:
        return 1
    return None


def infer_pcie_type_from_sysfs(bus_address):
    if not bus_address or bus_address == "0000:00:00.0":
        return None

    base = f"/sys/bus/pci/devices/{bus_address}"
    max_width_path = os.path.join(base, "max_link_width")
    max_speed_path = os.path.join(base, "max_link_speed")

    try:
        with open(max_width_path, "r", encoding="utf-8") as f:
            width = f.read().strip()
        with open(max_speed_path, "r", encoding="utf-8") as f:
            speed_raw = f.read().strip()
    except Exception:
        return None

    speed_match = re.search(r"([0-9]+(?:\.[0-9]+)?)", speed_raw)
    if not width or not speed_match:
        return None

    gen = _pcie_gen_from_speed(speed_match.group(1))
    if not gen:
        return None
    return f"PCI Express {gen} x{width}"


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


def is_pci_bridge(bus_address, lspci_devices):
    """Check if a PCI device is a bridge (root port, switch port, etc.)."""
    device_info = lspci_devices.get(bus_address)
    if not device_info:
        return False
    cls = str(device_info[0]).strip('"').lower()
    return "bridge" in cls


def has_downstream_endpoint(bus_address, lspci_devices):
    """Recursively check if a PCI bridge has any non-bridge endpoint downstream."""
    real_path = os.path.realpath(f"/sys/bus/pci/devices/{bus_address}")
    if not os.path.isdir(real_path):
        return False
    for entry in os.listdir(real_path):
        if not _PCI_ADDR_RE.match(entry):
            continue
        if _is_lspci_endpoint(lspci_devices.get(entry)):
            return True
        if has_downstream_endpoint(entry, lspci_devices):
            return True
    return False


def slot_has_pci_device(bus_address, lspci_devices):
    """Return True if a real endpoint device exists at or behind *bus_address*."""
    if not bus_address or bus_address == "0000:00:00.0":
        return False
    if _is_lspci_endpoint(lspci_devices.get(bus_address)):
        return True
    if is_pci_bridge(bus_address, lspci_devices):
        return has_downstream_endpoint(bus_address, lspci_devices)
    return False


def first_endpoint_addr(bus_address, lspci_devices):
    """Return the PCI address of the first non-bridge endpoint at or behind *bus_address*."""
    if not bus_address or bus_address == "0000:00:00.0":
        return None
    if _is_lspci_endpoint(lspci_devices.get(bus_address)):
        return bus_address
    real_path = os.path.realpath(f"/sys/bus/pci/devices/{bus_address}")
    if not os.path.isdir(real_path):
        return None
    for entry in os.listdir(real_path):
        if not _PCI_ADDR_RE.match(entry):
            continue
        result = first_endpoint_addr(entry, lspci_devices)
        if result:
            return result
    return None


def fix_slot_type(slot, motherboard_model):
    designation = slot.get("Designation", "")
    if motherboard_model == "MH53-G40-000":
        static_type = PCIE_SLOT_TYPE_GIGABYTE_MH53.get(designation) or MH53_AUX_SLOT_TYPE.get(designation)
        if static_type:
            slot["Type"] = static_type
            slot["PCI Type"] = static_type
            return

    slot_type = str(slot.get("Type", "")).strip()
    if slot_type not in ("", "-", "Unknown", "<OUT OF SPEC>"):
        return

    inferred = infer_pcie_type_from_sysfs(slot.get("Bus Address", ""))
    if inferred:
        slot["Type"] = inferred
        slot["PCI Type"] = inferred


def normalize_mh53_pci_slots(
    cards,
    lspci_devices,
):
    allowed = set(PCIE_SLOT_TYPE_GIGABYTE_MH53) | set(MH53_AUX_SLOT_TYPE)
    grouped = {}

    for slot in cards:
        designation = slot.get("Designation", "")
        bus = slot.get("Bus Address", "")
        if designation not in allowed:
            continue
        if not bus or bus == "0000:00:00.0":
            continue
        grouped.setdefault(designation, []).append(slot)

    endpoint_buses = {
        bus for bus, devinfo in lspci_devices.items()
        if _is_lspci_endpoint(devinfo)
    }

    normalized = []
    for designation in MH53_SLOT_ORDER:
        candidates = grouped.get(designation, [])
        if not candidates:
            continue

        def rank(slot):
            usage = str(slot.get("Current Usage", "")).strip().lower()
            bus = slot.get("Bus Address", "")
            return (
                usage == "in use",
                bus in endpoint_buses,
                bool(bus and bus != "0000:00:00.0"),
            )

        chosen = max(candidates, key=rank).copy()
        preferred_bus = ""

        for ba in BA_LUT_MH53.get(designation, []):
            if ba in endpoint_buses:
                preferred_bus = ba
                break

        if not preferred_bus:
            for ba in BA_LUT_MH53.get(designation, []):
                if ba in lspci_devices:
                    preferred_bus = ba
                    break

        if preferred_bus:
            chosen["Bus Address"] = preferred_bus

        chosen["Designation"] = designation
        chosen["Type"] = PCIE_SLOT_TYPE_GIGABYTE_MH53.get(
            designation,
            MH53_AUX_SLOT_TYPE.get(designation, chosen.get("Type", "Unknown")),
        )

        is_aux_slot = designation == "NVME0" or designation.startswith("M2_")
        if is_aux_slot:
            in_use = any(c.get("Bus Address", "") in endpoint_buses for c in candidates)
            if not in_use and preferred_bus:
                in_use = preferred_bus in endpoint_buses
        else:
            in_use = any(rank(c)[0] or rank(c)[1] for c in candidates)
            if not in_use and preferred_bus:
                in_use = preferred_bus in endpoint_buses or preferred_bus in lspci_devices

        chosen["Current Usage"] = "In Use" if in_use else "Available"
        fix_slot_type(chosen, "MH53-G40-000")
        normalized.append(chosen)

    return normalized if normalized else cards


def reassign_mh53_slot_ids(slots):
    order_rank = {name: i for i, name in enumerate(MH53_SLOT_ORDER)}
    pcie_slots = [s for s in slots if s.get("Designation", "").startswith("PCIE_")]
    aux_slots = [s for s in slots if not s.get("Designation", "").startswith("PCIE_")]

    pcie_slots.sort(key=lambda s: order_rank.get(s.get("Designation", ""), 999))
    aux_slots.sort(key=lambda s: order_rank.get(s.get("Designation", ""), 999))

    ordered = pcie_slots + aux_slots
    for idx, slot in enumerate(ordered, start=1):
        slot["ID"] = str(idx)
    return ordered


def finalize_slot(
    current_slot,
    cards,
    slot_seen,
    motherboard_model,
):
    if not current_slot:
        return

    designation = current_slot.get("Designation", "")
    bus_addr = current_slot.get("Bus Address", "")

    if not designation or designation == "N/A":
        return

    if not bus_addr or bus_addr == "0000:00:00.0":
        return

    if motherboard_model == "ME03-CE0-000" and designation in PCIE_SLOT_TYPE_GIGABYTE_ME03:
        current_slot["Type"] = PCIE_SLOT_TYPE_GIGABYTE_ME03[designation]
        candidate_addresses = BA_LUT_ME03.get(designation, [])
        if not isinstance(candidate_addresses, list):
            candidate_addresses = [candidate_addresses]
        current_slot["Current Usage"] = (
            "In Use" if current_slot.get("Bus Address") in candidate_addresses else "Available"
        )

        slot_seen.setdefault(designation, [])
        if current_slot not in slot_seen[designation]:
            slot_seen[designation].append(current_slot.copy())
            cards.append(current_slot.copy())
        return

    if motherboard_model == "B550I AORUS PRO AX":
        mapped = PCIE_SLOT_LUT_B550I.get(designation)
        if mapped:
            current_slot["Designation"] = mapped
            designation = mapped

    if motherboard_model == "MC13-LE1-000":
        normalized_designation = designation.replace("_", "")
        if normalized_designation in slot_seen:
            return
        slot_seen[normalized_designation] = True
        cards.append(current_slot.copy())
        return

    if motherboard_model == "MZ73-LM0-000":
        base_designation = re.sub(r"_[A-Z]$", "", designation)
        if base_designation in slot_seen:
            return
        slot_seen[base_designation] = True
        cards.append(current_slot.copy())
        return

    if designation not in slot_seen:
        slot_seen[designation] = True
        cards.append(current_slot.copy())


def dmidecode():
    dmi_output = run_command(["dmidecode", "-t", "9"])
    if not dmi_output:
        return []

    cards = []
    id_counter = 1
    current_slot = None
    slot_seen = {}
    disk_map = get_disk_by_path()

    motherboard_info = get_motherboard_model_server_info()
    motherboard_model = motherboard_info["Motherboard"][0].get("Product Name", "")

    pcie_pattern = re.compile(r"(PCIE_\d+)")

    for line in dmi_output.splitlines():
        if line.startswith("Handle "):
            finalize_slot(current_slot, cards, slot_seen, motherboard_model)
            current_slot = {}
            continue

        if current_slot is None:
            continue

        for field in g_dmi_fields:
            regex = re.search(r"^\s+({fld}):\s+(.*)".format(fld=field), line)
            if regex:
                current_slot[regex.group(1)] = regex.group(2)

        if "ID" not in current_slot:
            current_slot["ID"] = str(id_counter)
            id_counter += 1

        if "Designation" not in current_slot:
            continue

        designation = current_slot["Designation"]
        stripped_designation = re.sub(r"_[A-Z]$", "", designation)
        current_slot["Designation"] = stripped_designation

        if designation == "N/A":
            current_slot = None
            continue

        match = pcie_pattern.search(stripped_designation)
        normalized_designation = match.group(1) if match else stripped_designation
        actual_type = current_slot.get("Type", "Unknown")

        if motherboard_model.startswith("ME03") and normalized_designation in PCIE_SLOT_TYPE_GIGABYTE_ME03:
            current_slot["Type"] = PCIE_SLOT_TYPE_GIGABYTE_ME03[normalized_designation]
        elif motherboard_model.startswith("MZ73") and normalized_designation in PCIE_SLOT_TYPE_GIGABYTE_MZ73:
            current_slot["Type"] = PCIE_SLOT_TYPE_GIGABYTE_MZ73[normalized_designation]
        elif normalized_designation.startswith("U2_"):
            current_slot["Type"] = "MCIO Port"
        elif normalized_designation.startswith("OCU"):
            if actual_type and not str(actual_type).startswith("OCulink"):
                current_slot["Type"] = f"OCulink ({actual_type})"
        elif normalized_designation.startswith("M2_"):
            if actual_type and not str(actual_type).startswith("M.2 Slot"):
                current_slot["Type"] = f"M.2 Slot ({actual_type})"

            if motherboard_model.startswith("MZ73"):
                candidate_buses = BA_LUT_MZ73.get(current_slot["Designation"], [])
                if any(bus in disk_map for bus in candidate_buses):
                    current_slot["Current Usage"] = "In Use"
        else:
            current_slot["Type"] = actual_type

        if current_slot.get("Current Usage") != "In Use":
            current_slot["Current Usage"] = "Available"

        if motherboard_model.startswith("EPC621D8A"):
            bus_addr = current_slot.get("Bus Address", "")
            if bus_addr in BA_EPC612D8A:
                current_slot["Bus Address"] = BA_EPC612D8A[bus_addr]

    finalize_slot(current_slot, cards, slot_seen, motherboard_model)

    if motherboard_model == "MH53-G40-000":
        return normalize_mh53_pci_slots(cards, parse_lspci_output())
    return cards


def process_pci_slots_and_devices(
    pci_slots,
    devices,
    lspci_devices,
):
    matched_devices = {}
    motherboard_info = get_motherboard_model_server_info()
    motherboard_model = motherboard_info["Motherboard"][0].get("Product Name", "")
    mw34_claimed = set()

    for slot in pci_slots:
        designation = slot.get("Designation", "")
        slot_type = slot.get("Type", "")
        current_usage = slot.get("Current Usage", "")
        slot_key = re.sub(r"_[A-Z]$", "", designation)

        if motherboard_model.startswith("MZ73"):
            slot["Current Usage"] = "Available"

            if slot_key in BA_LUT_MZ73:
                bus_address_candidates = BA_LUT_MZ73[slot_key]
                for candidate_address in bus_address_candidates:
                    device_info = lspci_devices.get(candidate_address)
                    if not device_info:
                        continue

                    slot["Bus Address"] = candidate_address
                    slot["Current Usage"] = "In Use"

                    if len(device_info) >= 3:
                        card_type = device_info[0].strip('"')
                        if card_type == "Ethernet controller":
                            slot["Card Type"] = "Network Card"
                            manufacturer = device_info[1].strip('"')
                            if manufacturer == "Intel Corporation":
                                slot["Card Model"] = device_info[2].replace('"', "")
                            elif manufacturer == "Mellanox Technologies" and len(device_info) >= 5:
                                slot["Card Model"] = device_info[4].replace('"', "")
                            else:
                                slot["Card Model"] = " ".join(device_info[1:3]).replace('"', "")
                    break

        # ---------- MW34 special handling (claim-and-remove) ----------
        if motherboard_model.startswith("MW34"):
            slot["Current Usage"] = "Available"
            candidates = BA_LUT_MW34.get(slot_key, [])
            claimed = False
            for cand in candidates:
                if cand in mw34_claimed:
                    continue
                if slot_has_pci_device(cand, lspci_devices):
                    endpoint = first_endpoint_addr(cand, lspci_devices)
                    slot["Bus Address"] = endpoint or cand
                    slot["Current Usage"] = "In Use"
                    mw34_claimed.add(cand)
                    claimed = True
                    break
            if not claimed:
                # Pick first unclaimed candidate that exists, or clear
                slot["Bus Address"] = ""
                for cand in candidates:
                    if cand not in mw34_claimed and cand in lspci_devices:
                        slot["Bus Address"] = cand
                        break

        if designation not in matched_devices:
            matched_devices[designation] = slot
        else:
            existing_slot = matched_devices[designation]
            if current_usage == "In Use" and existing_slot.get("Current Usage", "") != "In Use":
                matched_devices[designation] = slot
            elif current_usage == "Available" and existing_slot.get("Current Usage", "") == "Unknown":
                matched_devices[designation] = slot
            elif current_usage == "Unknown" or slot_type == "<OUT OF SPEC>":
                if existing_slot.get("Current Usage", "") != "In Use":
                    matched_devices[designation] = slot

    for device in devices:
        device_bus = device.get("Bus Address")
        if not device_bus:
            continue

        for slot in matched_devices.values():
            if slot.get("Bus Address") == device_bus:
                for k, v in device.items():
                    if k in ("Type", "PCI Type"):
                        continue
                    slot[k] = v
                break

    return list(matched_devices.values())


def lspci_hba():
    lspci_hba_result = run_command(
        ["lspci", "-d", "1000:*", "-vv", "-i", "/opt/45drives/tools/pci.ids"]
    )
    if not lspci_hba_result:
        return []

    hba_cards = []
    rx_pci = re.compile(
        r'^(\w\w:\w\w\.\w).*\n.*(?:(?:(?:^\t).*\n)+^.*)?'
        r'(9600-16i|9600-24i|SAS9305-16i|SAS9305-24i|HBA 9405W-16i|9361-16i|'
        r'HBA 9400-16i|LSI HBA 9400-16i|SAS3416|9361-24i|9660-16i).*\n',
        re.MULTILINE,
    )

    for match in rx_pci.finditer(lspci_hba_result):
        hba_cards.append({
            "Model": match.group(2),
            "Bus Address": "0000:" + match.group(1),
        })

    return hba_cards


def network():
    network_result = run_command(
        ["/usr/share/cockpit/45drives-motherboard/helper_scripts/network"]
    )
    if not network_result:
        return []

    try:
        network_output = json.loads(network_result)
    except json.JSONDecodeError:
        return []

    network_cards = []
    for connection in network_output.get("Network Info", []):
        if "PCI Slot" in connection:
            network_cards.append(connection)

    return network_cards


def find_primary_bus_for_port(port_bus_address):
    return NETWORK_BA_LUT_B550I.get(port_bus_address, port_bus_address)


def getNetworkCardModel(busAddress):
    trimmedBusAddress = busAddress[5:] if busAddress.startswith("0000:") else busAddress
    primary_bus_address = find_primary_bus_for_port(trimmedBusAddress) or ""
    lspci_output = run_command(["lspci"])
    if not lspci_output:
        return "unknown"

    for line in lspci_output.splitlines():
        regex_addr = re.search(rf"^{re.escape(primary_bus_address)}\s(.*)$", line)
        if regex_addr is None:
            continue

        for model in g_network_card_models:
            if re.search(re.escape(model), regex_addr.group(1)):
                return model

    return "unknown"


def lsblk(device):
    lsblk_output = run_command(["lsblk", "-l", f"/dev/{device}"])
    if not lsblk_output:
        return []

    partitions = []
    for line in lsblk_output.splitlines():
        regex = re.search(r"^(\S+)\s+\S+\s+\S+\s+(\S+)\s+\S+\s+(\S+)(.*)$", line)
        if regex is not None and regex.group(1) != "NAME":
            partitions.append({
                "Name": regex.group(1),
                "Size": regex.group(2),
                "Type": regex.group(3),
                "Mount Point": regex.group(4).lstrip(),
            })
    return partitions


def sata():
    lspci_output = run_command(["lspci"])
    if not lspci_output:
        return []

    sata_cards = []
    for line in lspci_output.splitlines():
        for field in g_sata_controllers:
            regex = re.search(r"^(\S+).*({fld}).*$".format(fld=field), line)
            if regex is not None:
                sata_cards.append({
                    "Card Type": "Serial ATA Controller",
                    "Card Model": regex.group(2),
                    "Bus Address": "0000:" + regex.group(1),
                })

    ls_output = run_command(["ls", "-l", "/dev/disk/by-path"])
    if not ls_output:
        for card in sata_cards:
            card["Connections"] = []
        return sata_cards

    for card in sata_cards:
        drives = []
        short_bus = card["Bus Address"].replace("0000:", "")
        for line in ls_output.splitlines():
            drive_dict = {}
            regex = re.search(rf"pci-({re.escape(short_bus)})-ata-(\d)\s->\s\W+(.*)", line)
            if regex is not None:
                drive_dict["Device"] = regex.group(3)
                drive_dict["Path"] = "pci-" + regex.group(1) + "-ata-" + regex.group(2)
                drive_dict["Partitions"] = lsblk(drive_dict["Device"])
                drives.append(drive_dict.copy())
        card["Connections"] = drives

    return sata_cards


def get_disk_by_path():
    output = run_command(["ls", "-l", "/dev/disk/by-path"])
    if not output:
        return {}

    disk_map = {}
    pci_nvme_regex = re.compile(r"pci-([0-9a-fA-F:.-]+)-nvme-\d+")

    for line in output.splitlines():
        match = pci_nvme_regex.search(line)
        if match:
            pci_address = match.group(1)
            device = line.split(" -> ")[-1].strip().replace("../../", "")
            if not device.endswith("part"):
                disk_map[pci_address] = device

    return disk_map


def get_ocu_usage_from_by_path():
    usage = {"OCU1": [], "OCU2": []}
    bypath = "/dev/disk/by-path"

    try:
        entries = os.listdir(bypath)
    except Exception:
        return usage

    rx = re.compile(r"^pci-[0-9a-fA-F:.-]+-ata-(\d+)$")

    for name in entries:
        if "-part" in name or "-nvme-" in name or "-sas-" in name:
            continue

        m = rx.match(name)
        if not m:
            continue

        port = int(m.group(1))
        target = os.path.realpath(os.path.join(bypath, name))
        dev = os.path.basename(target)

        if 1 <= port <= 4:
            usage["OCU1"].append(dev)
        elif 5 <= port <= 8:
            usage["OCU2"].append(dev)

    return usage


def main():
    pci_slots = dmidecode()
    hba_cards = lspci_hba()
    network_cards = network()
    sata_cards = sata()
    lspci_devices = parse_lspci_output()

    mobo_info = get_motherboard_model_server_info()
    motherboard_model = mobo_info["Motherboard"][0]["Product Name"]

    for s in pci_slots:
        s.setdefault("Bus Address", "")

    if pci_slots and hba_cards:
        for hba in hba_cards:
            for slot in pci_slots:
                if motherboard_model.startswith("MZ73") and slot.get("Designation") in BA_LUT_MZ73:
                    if hba["Bus Address"] in BA_LUT_MZ73[slot["Designation"]]:
                        slot["Card Type"] = "HBA"
                        slot["Card Model"] = hba["Model"]
                elif hba["Bus Address"] == slot.get("Bus Address"):
                    slot["Card Type"] = "HBA"
                    slot["Card Model"] = hba["Model"]

    if pci_slots and network_cards:
        for slot in pci_slots:
            for card in network_cards:
                if motherboard_model.startswith("B550I") and slot.get("Designation") in BA_LUT_B550I:
                    if card.get("Bus Address") in BA_LUT_B550I[slot["Designation"]]:
                        slot["Card Type"] = "Network Card"
                        slot["Card Model"] = getNetworkCardModel(str(slot.get("Bus Address", "")))
                        slot.setdefault("Connections", [])
                        slot["Connections"].append(card)
                else:
                    normalized_card_id = str(card.get("PCI Slot", "")).replace("SLOT", "")
                    if normalized_card_id in (slot.get("Designation", ""), slot.get("ID", "")):
                        slot["Card Type"] = "Network Card"
                        slot["Card Model"] = getNetworkCardModel(str(slot.get("Bus Address", "")))
                        slot.setdefault("Connections", [])
                        slot["Connections"].append(card)

    if pci_slots and sata_cards:
        for slot in pci_slots:
            for card in sata_cards:
                if card.get("Bus Address") == slot.get("Bus Address"):
                    slot["Card Type"] = card["Card Type"]
                    slot["Card Model"] = card["Card Model"]
                    slot["Connections"] = card.get("Connections", [])

    all_devices = hba_cards + network_cards + sata_cards
    matched_devices = process_pci_slots_and_devices(pci_slots, all_devices, lspci_devices)

    # MW34: bus addresses are only correct after process_pci_slots_and_devices(),
    # so re-run HBA / network / SATA matching against the fixed endpoint addresses.
    if motherboard_model.startswith("MW34"):
        for slot in matched_devices:
            bus = slot.get("Bus Address", "")
            if not bus or bus == "0000:00:00.0":
                continue
            for hba in hba_cards:
                if hba["Bus Address"] == bus:
                    slot["Card Type"] = "HBA"
                    slot["Card Model"] = hba["Model"]
                    break
            # Try matching via network helper output first
            net_matched = False
            for card in network_cards:
                if card.get("Bus Address") == bus:
                    slot["Card Type"] = "Network Card"
                    slot["Card Model"] = getNetworkCardModel(bus)
                    slot.setdefault("Connections", []).append(card)
                    net_matched = True
                    break
            # Fallback: detect network card from lspci_devices (like MZ73)
            if not net_matched and slot.get("Card Type", "-") == "-":
                devinfo = lspci_devices.get(bus)
                if devinfo and "ethernet" in str(devinfo[0]).strip('"').lower():
                    slot["Card Type"] = "Network Card"
                    model = getNetworkCardModel(bus)
                    if model not in ("-", "unknown"):
                        slot["Card Model"] = model
                    elif len(devinfo) > 2:
                        slot["Card Model"] = devinfo[2].strip('"')
            for card in sata_cards:
                if card.get("Bus Address") == bus:
                    slot["Card Type"] = card["Card Type"]
                    slot["Card Model"] = card["Card Model"]
                    slot["Connections"] = card.get("Connections", [])
                    break

    for slot in matched_devices:
        fix_slot_type(slot, motherboard_model)

    if motherboard_model.startswith("EC266D2I"):
        ocu_usage = get_ocu_usage_from_by_path()
        for slot in matched_devices:
            desig = slot.get("Designation", "")
            if desig in ("OCU1", "OCU2") and ocu_usage.get(desig):
                slot["Current Usage"] = "In Use"

    seen_bus_addresses = {}
    for slot in matched_devices:
        bus = slot.get("Bus Address", "")
        desig = slot.get("Designation", "")
        if not bus:
            continue

        if bus in seen_bus_addresses and desig.startswith("M2"):
            slot["Card Type"] = "-"
            slot["Card Model"] = "-"
            slot["Firmware Version"] = "-"
        else:
            seen_bus_addresses[bus] = desig

    for slot in matched_devices:
        card_type = slot.get("Card Type", "-")
        bus = slot.get("Bus Address", "")
        if card_type not in ("", "-", None) and bus not in ("", "0000:00:00.0"):
            slot["Current Usage"] = "In Use"

    if motherboard_model.startswith("ME03"):
        other_devices = [d for d in matched_devices if d.get("Designation", "").startswith(("M2", "U2"))]
        pcie_devices = [d for d in matched_devices if not d.get("Designation", "").startswith(("M2", "U2"))]
        pcie_devices.sort(key=lambda x: x.get("Designation", ""))
        matched_devices = pcie_devices + other_devices

        for idx, device in enumerate(matched_devices, start=0):
            device["ID"] = str(idx)

    if motherboard_model == "MH53-G40-000":
        matched_devices = reassign_mh53_slot_ids(matched_devices)

    print(json.dumps({"PCI Info": matched_devices}, indent=4))


if __name__ == "__main__":
    main()