#!/usr/bin/env python3
"""usb3check - watch USB/Thunderbolt enumeration and report cable tier.

Plug a device in while this runs and it reports the negotiated link speed and
which cable tier that qualifies for.

    Tier 1  Thunderbolt 4/5, 240W EPR   dock uplink, displays, external SSDs
    Tier 2  >=10Gbps, 240W EPR          general data
    Tier 3  charge-only, 240W EPR       bench power, laptop charging
    Tier 4  everything else             phone charging, or the bin

Data rate is measured. Wattage is NOT - nothing in sysfs exposes the cable's
e-marker, so the tier printed is the ceiling the data rate allows. Confirm the
power rating from the purchase or an inline analyzer before promoting a cable.

Reads /sys/bus/usb/devices and /sys/bus/thunderbolt/devices. No root, no deps.

    ./usb3check.py           watch continuously
    ./usb3check.py --once    print current devices grouped by speed
"""

import argparse
import os
import sys
import time

USB_SYSFS = "/sys/bus/usb/devices"
TB_SYSFS = "/sys/bus/thunderbolt/devices"

TIER_NOTES = {
    1: "Thunderbolt link established. Tier 1 if the cable is also 240W EPR.",
    2: "SuperSpeed link established. Tier 2 if the cable is also 240W EPR.",
    3: "No SuperSpeed link - data-incapable. Tier 3 if 240W EPR, else tier 4.",
    4: "No link worth keeping. Tier 4.",
}


def read_attr(path, name):
    try:
        with open(os.path.join(path, name)) as fh:
            return fh.read().strip()
    except OSError:
        return None


def usb_snapshot():
    """Map of sysfs name -> device info, excluding root hubs and interfaces."""
    devices = {}
    try:
        entries = os.listdir(USB_SYSFS)
    except OSError as exc:
        sys.exit(f"cannot read {USB_SYSFS}: {exc}")

    for name in entries:
        if name.startswith("usb") or ":" in name:
            continue  # root hub, or an interface rather than a device
        path = os.path.join(USB_SYSFS, name)
        raw_speed = read_attr(path, "speed")
        if raw_speed is None:
            continue
        try:
            speed = float(raw_speed)
        except ValueError:
            continue
        devices[name] = {
            "kind": "usb",
            "speed": speed,
            "vid": read_attr(path, "idVendor") or "????",
            "pid": read_attr(path, "idProduct") or "????",
            "product": read_attr(path, "product") or "",
            "vendor": read_attr(path, "manufacturer") or "",
            "class": read_attr(path, "bDeviceClass") or "",
        }
    return devices


def tb_snapshot():
    """Map of sysfs name -> Thunderbolt device info. Empty if no TB subsystem."""
    devices = {}
    if not os.path.isdir(TB_SYSFS):
        return devices

    for name in os.listdir(TB_SYSFS):
        if name.startswith("domain") or ":" in name:
            continue  # domain controller, or a retimer/port entry
        path = os.path.join(TB_SYSFS, name)
        if read_attr(path, "device_name") is None and read_attr(path, "device") is None:
            continue
        devices[name] = {
            "kind": "thunderbolt",
            "product": read_attr(path, "device_name") or "",
            "vendor": read_attr(path, "vendor_name") or "",
            "authorized": read_attr(path, "authorized") or "0",
            "rx": read_attr(path, "rx_speed") or "",
            "tx": read_attr(path, "tx_speed") or "",
            "generation": read_attr(path, "generation") or "",
        }
    return devices


def snapshot():
    combined = usb_snapshot()
    for name, info in tb_snapshot().items():
        combined[f"tb:{name}"] = info
    return combined


def fmt_speed(mbps):
    if mbps >= 1000:
        return f"{mbps / 1000:g} Gbps"
    return f"{mbps:g} Mbps"


def generation(mbps):
    if mbps >= 20000:
        return "USB 3.2 Gen 2x2"
    if mbps >= 10000:
        return "USB 3.1 Gen 2 (SuperSpeed+)"
    if mbps >= 5000:
        return "USB 3.0 (SuperSpeed)"
    if mbps >= 480:
        return "USB 2.0 (High Speed)"
    return "USB 1.x (Full/Low Speed)"


def describe(info):
    label = " ".join(x for x in (info.get("vendor", ""), info.get("product", ""))).strip()
    if not label:
        label = f"{info.get('vid', '????')}:{info.get('pid', '????')}"
    if info["kind"] == "thunderbolt":
        gen = info.get("generation")
        label += f"  [thunderbolt{' gen ' + gen if gen else ''}]"
    elif info.get("class") == "09":
        label += "  [hub]"
    return label


def speed_column(info):
    if info["kind"] == "thunderbolt":
        return info.get("rx") or "TB link"
    return fmt_speed(info["speed"])


def classify(added):
    """Best tier the observed links qualify for."""
    if any(i["kind"] == "thunderbolt" and i.get("authorized") == "1" for i in added.values()):
        return 1
    usb = [i["speed"] for i in added.values() if i["kind"] == "usb"]
    if not usb:
        return 4
    best = max(usb)
    if best >= 10000:
        return 2
    if best >= 5000:
        return 2
    if best >= 480:
        return 3
    return 4


def report(added, removed):
    print(f"\n-- {time.strftime('%H:%M:%S')} " + "-" * 44)
    for name, info in sorted(removed.items()):
        print(f"  - {name:<14} {'':>10}  {describe(info)}")
    for name, info in sorted(added.items()):
        print(f"  + {name:<14} {speed_column(info):>10}  {describe(info)}")

    if not added:
        return

    tier = classify(added)
    usb = [i["speed"] for i in added.values() if i["kind"] == "usb"]
    print()
    if usb:
        print(f"  fastest USB link: {fmt_speed(max(usb))}  ({generation(max(usb))})")
    print(f"  TIER {tier}: {TIER_NOTES[tier]}")
    if tier >= 3:
        print("           a USB 3 hub always enumerates twice - once at 480 Mbps and")
        print("           again on the SuperSpeed bus. Only one showed up here.")
    if tier == 2:
        print("           not a Thunderbolt link. Tier 1 needs `boltctl list` with the dock.")


def print_tree(devices):
    if not devices:
        print("no devices found")
        return
    tb = {k: v for k, v in devices.items() if v["kind"] == "thunderbolt"}
    usb = {k: v for k, v in devices.items() if v["kind"] == "usb"}

    if tb:
        print("\nThunderbolt")
        for name, info in sorted(tb.items()):
            state = "authorized" if info.get("authorized") == "1" else "not authorized"
            print(f"  {name:<14} {info.get('rx', ''):>10}  {describe(info)}  [{state}]")

    by_speed = {}
    for name, info in usb.items():
        by_speed.setdefault(info["speed"], []).append((name, info))
    for speed in sorted(by_speed, reverse=True):
        print(f"\n{fmt_speed(speed)}  ({generation(speed)})")
        for name, info in sorted(by_speed[speed]):
            print(f"  {name:<14} {describe(info)}")


def watch(interval, settle):
    base = snapshot()
    fast = sum(1 for i in base.values() if i["kind"] == "usb" and i["speed"] >= 5000)
    tb = sum(1 for i in base.values() if i["kind"] == "thunderbolt")
    print(f"watching {USB_SYSFS}" + (f" and {TB_SYSFS}" if os.path.isdir(TB_SYSFS) else ""))
    print(f"baseline: {len(base)} devices, {fast} at SuperSpeed, {tb} Thunderbolt")
    print("plug something in - Ctrl-C to stop")

    while True:
        time.sleep(interval)
        current = snapshot()
        if current == base:
            continue

        # Let enumeration settle before reporting; a hub plus its downstream
        # devices arrive over several hundred milliseconds.
        while True:
            time.sleep(settle)
            following = snapshot()
            if following == current:
                break
            current = following

        added = {k: v for k, v in current.items() if k not in base}
        removed = {k: v for k, v in base.items() if k not in current}
        report(added, removed)
        base = current


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--once", action="store_true",
                        help="print current devices grouped by speed, then exit")
    parser.add_argument("--interval", type=float, default=0.4,
                        help="poll interval in seconds (default: 0.4)")
    parser.add_argument("--settle", type=float, default=1.0,
                        help="settle time after a change before reporting (default: 1.0)")
    args = parser.parse_args()

    if args.once:
        print_tree(snapshot())
        return

    try:
        watch(args.interval, args.settle)
    except KeyboardInterrupt:
        print()


if __name__ == "__main__":
    main()
