# Benchmark: LRU eviction vs add-then-remove-after-use on ESP32.
# Run on device: mpremote run esp32/benchmark_peers.py
# (add/del_peer are timed; send() may fail if no peer is listening - timing still valid)
import espnow
import network
import time

BROADCAST = b"\xff\xff\xff\xff\xff\xff"
MAX_PEERS = 20
ITERATIONS = 50
PAYLOAD = b"x" * 32  # small payload

network.WLAN(network.STA_IF).active(True)
esp = espnow.ESPNow()
esp.active(True)
esp.add_peer(BROADCAST)

# Build 19 dummy MACs so we have 20 peers total (broadcast + 19).
def mac(i):
    return bytes([0, 0, 0, 0, 0, i])
peers_list = [mac(i) for i in range(1, 20)]
for p in peers_list:
    esp.add_peer(p)

# One "new" MAC we'll add/remove.
new_mac = bytes([0, 0, 0, 0, 0, 99])

def bench_lru():
    """LRU: ensure_peer (evict oldest + add new), send, update last_used."""
    last_used = {BROADCAST: time.ticks_ms()}
    for p in peers_list:
        last_used[p] = time.ticks_ms()
    # Pre-remove one so we have 19; ensure_peer(new) will add 20th.
    esp.del_peer(peers_list[-1])
    last_used.pop(peers_list[-1], None)
    # Now 19 peers. Each iteration: ensure_peer(new) -> add_peer(new), send, update.
    # Next iter: ensure_peer(new) -> already there, just send. So we need to force
    # eviction each time: use a different "new" each time so we always evict+add.
    t0 = time.ticks_us()
    for i in range(ITERATIONS):
        addr = bytes([0, 0, 0, 0, 0, 50 + (i % 30)])  # 30 different "new" MACs
        peers = esp.get_peers()
        peer_macs = [p[0] for p in peers]
        if addr not in peer_macs:
            if len(peer_macs) >= MAX_PEERS:
                oldest_mac = None
                oldest_ts = time.ticks_ms()
                for m in peer_macs:
                    if m == BROADCAST:
                        continue
                    ts = last_used.get(m, 0)
                    if ts <= oldest_ts:
                        oldest_ts = ts
                        oldest_mac = m
                if oldest_mac is not None:
                    esp.del_peer(oldest_mac)
                    last_used.pop(oldest_mac, None)
            esp.add_peer(addr)
        esp.send(addr, PAYLOAD)
        last_used[addr] = time.ticks_ms()
    t1 = time.ticks_us()
    return time.ticks_diff(t1, t0)

def bench_add_then_remove():
    """Add peer, send, del_peer (remove after use). At 20 we must del one first."""
    # Start full: 20 peers. To add new we del any one, add new, send, del new.
    victim = peers_list[0]
    t0 = time.ticks_us()
    for i in range(ITERATIONS):
        esp.del_peer(victim)  # make room
        esp.add_peer(new_mac)
        esp.send(new_mac, PAYLOAD)
        esp.del_peer(new_mac)
        esp.add_peer(victim)  # put victim back so we're at 20 again
    t1 = time.ticks_us()
    return time.ticks_diff(t1, t0)

def bench_send_existing():
    """Baseline: send to existing peer only (no add/del)."""
    t0 = time.ticks_us()
    for _ in range(ITERATIONS):
        esp.send(peers_list[0], PAYLOAD)
    t1 = time.ticks_us()
    return time.ticks_diff(t1, t0)

print("ESP-NOW peer benchmark ({} iterations)".format(ITERATIONS))
print()

# Baseline: send to existing peer
try:
    us = bench_send_existing()
    print("Send to existing peer only: {:>8} us total  {:>7.1f} us/iter".format(us, us / ITERATIONS))
except Exception as e:
    print("Send existing failed:", e)
print()

# LRU: evict oldest then add new, send
try:
    us = bench_lru()
    print("LRU (evict oldest + add + send): {:>8} us total  {:>7.1f} us/iter".format(us, us / ITERATIONS))
except Exception as e:
    print("LRU failed:", e)
print()

# Add then remove after use
try:
    us = bench_add_then_remove()
    print("Add then remove after use:       {:>8} us total  {:>7.1f} us/iter".format(us, us / ITERATIONS))
except Exception as e:
    print("Add-then-remove failed:", e)
print()
print("Done.")