# Test configuring transmit rates for ESP-NOW on ESP32

import sys
import time

try:
    import network
    import espnow
except ImportError:
    print("SKIP")
    raise SystemExit

# ESP8266 doesn't support multiple ESP-NOW data rates
if sys.platform == "esp8266":
    print("SKIP")
    raise SystemExit

# Currently the config(rate=...) implementation is not compatible with ESP32-C6
# (this test passes when C6 is receiver, but not if C6 is sender.)
if "ESP32-C6" in sys.implementation._machine:
    print("SKIP")
    raise SystemExit

# ESP32-C2 doesn't support Long Range mode. This test is currently written assuming
# LR mode can be enabled.
if "ESP32-C2" in sys.implementation._machine:
    print("SKIP")
    raise SystemExit


timeout_ms = 1000
default_pmk = b"MicroPyth0nRules"
CHANNEL = 9


def init_sta():
    sta = network.WLAN(network.WLAN.IF_STA)
    e = espnow.ESPNow()
    e.active(True)
    sta.active(True)
    sta.disconnect()  # Force AP disconnect for any saved config, important so the channel doesn't change
    sta.config(channel=CHANNEL)
    e.set_pmk(default_pmk)
    # Enable both default 802.11 modes and Long Range modes
    sta.config(protocol=network.WLAN.PROTOCOL_LR | network.WLAN.PROTOCOL_DEFAULT)
    return sta, e


# Receiver
def instance0():
    sta, e = init_sta()
    try:
        multitest.globals(PEER=sta.config("mac"))
        multitest.next()
        while True:
            peer, msg = e.recv(timeout_ms)
            if peer is None:
                print("Timeout")
                break
            # Note that we don't have any way in Python to tell what data rate this message
            # was received with, so we're assuming the rate was correct.
            print(msg)
    finally:
        # Important to stop both even if the test is skipped due to the other instance
        sta.active(False)
        e.active(False)


# Sender
def instance1():
    sta, e = init_sta()
    multitest.next()
    peer = PEER
    try:
        e.add_peer(peer)
        # Test normal, non-LR rates
        for msg, rate in (
            (b"default rate", None),
            (b"5Mbit", espnow.RATE_5M),
            (b"11Mbit", espnow.RATE_11M),
            (b"24Mbit", espnow.RATE_24M),
            (b"54Mbit", espnow.RATE_54M),
            (b"250K LR", espnow.RATE_LORA_250K),
            (b"500K LR", espnow.RATE_LORA_500K),
            # switch back to non-LR rates to check it's all OK
            (b"1Mbit again", espnow.RATE_1M),
            (b"11Mbit again", espnow.RATE_11M),
        ):
            if rate is not None:
                e.config(rate=rate)
            for _ in range(3):
                e.send(peer, msg)
            time.sleep_ms(50)  # give messages some time to be received before continuing
        e.del_peer(peer)
    finally:
        sta.active(False)
        e.active(False)
