"""Device-side radio diagnostic test (MicroPython).

Checks:
1) STA/AP bring-up on channel 5
2) ESP-NOW init and broadcast peer add
3) Broadcast TX test packet send
4) RX wait window to see any incoming ESP-NOW frames
"""

import espnow
import machine
import network
import time
import ubinascii


CHANNEL = 5
RX_WINDOW_MS = 3000
WDT_TIMEOUT_MS = 10000
BROADCAST_MAC = b"\xff\xff\xff\xff\xff\xff"
TEST_PAYLOAD = b"\x4c\x05\x01\x00\x00\x00"


def _mac_hex(mac_bytes):
    try:
        return ubinascii.hexlify(mac_bytes).decode()
    except Exception:
        return "?"


def run_diag(channel=CHANNEL, rx_window_ms=RX_WINDOW_MS):
    wdt = machine.WDT(timeout=WDT_TIMEOUT_MS)
    wdt.feed()

    print("diag start")
    print("cpu freq", machine.freq())

    sta = network.WLAN(network.STA_IF)
    ap = network.WLAN(network.AP_IF)

    # Clean start
    try:
        sta.active(False)
        ap.active(False)
        time.sleep_ms(100)
    except Exception as e:
        print("wifi reset failed", repr(e))

    # STA setup
    try:
        sta.active(True)
        sta.config(pm=network.WLAN.PM_NONE)
        sta.config(channel=channel)
        print("sta ok ch", sta.config("channel"), "mac", _mac_hex(sta.config("mac")))
    except Exception as e:
        print("sta setup failed", repr(e))

    # AP setup
    try:
        ap.active(True)
        try:
            ap.config(essid="diag-ap", channel=channel, hidden=True)
        except TypeError:
            ap.config(essid="diag-ap", channel=channel)
        print("ap ok ch", ap.config("channel"), "mac", _mac_hex(ap.config("mac")))
    except Exception as e:
        print("ap setup failed", repr(e))

    wdt.feed()

    # ESP-NOW setup
    try:
        e = espnow.ESPNow()
        e.active(True)
        print("espnow active ok")
    except Exception as e_err:
        print("espnow init failed", repr(e_err))
        return

    # Add broadcast peer
    try:
        e.add_peer(BROADCAST_MAC, channel=channel)
        print("add bcast peer ok")
    except TypeError:
        try:
            e.add_peer(BROADCAST_MAC)
            print("add bcast peer ok (no channel arg)")
        except Exception as e_err:
            print("add bcast peer failed", repr(e_err))
    except Exception as e_err:
        print("add bcast peer failed", repr(e_err))

    # TX test
    try:
        ok = e.send(BROADCAST_MAC, TEST_PAYLOAD, True)
        print("tx bcast", ok, "len", len(TEST_PAYLOAD))
    except Exception as e_err:
        print("tx bcast failed", repr(e_err))

    # RX window
    print("rx window ms", rx_window_ms)
    t_end = time.ticks_add(time.ticks_ms(), rx_window_ms)
    rx_count = 0
    while time.ticks_diff(t_end, time.ticks_ms()) > 0:
        wdt.feed()
        host, msg = e.recv(100)
        if host:
            rx_count += 1
            print("rx", rx_count, _mac_hex(host), "len", len(msg))
        time.sleep_ms(5)

    print("diag done rx_count", rx_count)


if __name__ == "__main__":
    run_diag()