atmos/kernel/

wifi.rs

// wifi.rs - CYW43455 WiFi driver for RPi3B+ via SDIO (Broadcom FMAC protocol)
// Firmware must be placed at SylFS /boot/brcmfmac43455-sdio.bin
#![allow(dead_code)]

extern crate alloc;
use crate::kernel::sdio;
use alloc::vec::Vec;
use spin::Mutex;

// ---- SDPCM transport layer constants ----
const SDPCM_HDR_LEN: usize = 12; // SDPCM header size
const SDPCM_CHAN_CONTROL: u8 = 0; // ioctl channel
const SDPCM_CHAN_EVENT: u8 = 1; // async event channel
const SDPCM_CHAN_DATA: u8 = 2; // Ethernet data channel

// ---- CDC ioctl header constants ----
const CDC_HDR_LEN: usize = 16;
const CDC_FLAG_SET: u32 = 0x0000; // SET ioctl
const CDC_FLAG_GET: u32 = 0x0002; // GET ioctl (bit 1)
const CDC_ID_SHIFT: u32 = 16; // transaction id in bits [31:16]

// ---- BRCMF ioctl command numbers ----
const BRCMF_C_GET_REVINFO: u32 = 202;
const BRCMF_C_SET_INFRA: u32 = 20;
const BRCMF_C_SET_AUTH: u32 = 22;
const BRCMF_C_SET_WSEC: u32 = 134;
const BRCMF_C_SET_SSID: u32 = 26;
const BRCMF_C_DISASSOC: u32 = 52;
const BRCMF_C_SET_PM: u32 = 86;
const BRCMF_C_SET_VAR: u32 = 263; // set iovar
const BRCMF_C_GET_VAR: u32 = 262; // get iovar
const BRCMF_C_SET_SCAN_CHANNEL_TIME: u32 = 185;
const BRCMF_C_SET_SCAN_UNASSOC_TIME: u32 = 187;

// ---- BRCMF event types (received on SDPCM_CHAN_EVENT) ----
const BRCMF_E_SET_SSID: u32 = 0;
const BRCMF_E_AUTH: u32 = 3;
const BRCMF_E_DEAUTH: u32 = 5;
const BRCMF_E_ASSOC: u32 = 7;
const BRCMF_E_DISASSOC: u32 = 11;
const BRCMF_E_LINK: u32 = 16;
const BRCMF_E_PSK_SUP: u32 = 46;
const BRCMF_E_STATUS_SUCCESS: u32 = 0;

// ---- WiFi security ----
const WSEC_AES: u32 = 4; // WPA2-AES
const WSEC_PASSPHRASE: u16 = 0x0200; // flag: key is ASCII passphrase (firmware converts to PSK)

// ---- Backplane / clock registers ----
const SBSDIO_FUNC1_CHIPCLKCSR: u32 = 0x1000E;
const SBSDIO_HT_AVAIL: u8 = 0x80; // HT clock available
const SBSDIO_HT_AVAIL_REQ: u8 = 0x10; // request HT clock
const SBSDIO_ALP_AVAIL_REQ: u8 = 0x08;
const SBSDIO_FUNC1_SBADDRLOW: u32 = 0x1000A;
const SBSDIO_FUNC1_SBADDRMID: u32 = 0x1000B;
const SBSDIO_FUNC1_SBADDRHIGH: u32 = 0x1000C;

// ---- CYW43455 chip addresses ----
const BRCMF_BACKPLANE_ARM: u32 = 0x18003000;

// ---- WiFi state machine ----
#[derive(Copy, Clone, PartialEq, Eq)]
pub enum WifiState {
    Uninitialized,
    FirmwareLoaded,
    Idle,
    Scanning,
    Connecting,
    Connected,
    Error,
}

pub struct WifiNetwork {
    pub ssid: [u8; 32],
    pub ssid_len: usize,
    pub rssi: i8,
    pub channel: u8,
    pub security: WifiSecurity,
    pub bssid: [u8; 6],
}

#[derive(Copy, Clone, PartialEq, Eq)]
pub enum WifiSecurity {
    Open,
    Wpa2Personal,
    Unknown,
}

// ---- Global state ----
static WIFI_STATE: Mutex<WifiState> = Mutex::new(WifiState::Uninitialized);
static mut WIFI_IP: [u8; 4] = [0; 4];
static mut WIFI_SSID: [u8; 32] = [0; 32];
static mut WIFI_SSID_LEN: usize = 0;
static mut TX_SEQ: u8 = 0;
static mut TX_ID: u16 = 0;

// ---- Utilities ----

fn delay_ms(ms: u32) {
    for _ in 0..(ms * 50_000) {
        unsafe { core::arch::asm!("nop") };
    }
}

// ============================================================
// Backplane access (via SDIO function 1)
// ============================================================

unsafe fn set_backplane_window(addr: u32) -> bool {
    let a0 = ((addr >> 8) & 0xFF) as u8;
    let a1 = ((addr >> 16) & 0xFF) as u8;
    let a2 = ((addr >> 24) & 0xFF) as u8;
    sdio::cmd52(sdio::SDIO_FUNC_BACKPLANE, true, SBSDIO_FUNC1_SBADDRLOW, a0).is_some()
        && sdio::cmd52(sdio::SDIO_FUNC_BACKPLANE, true, SBSDIO_FUNC1_SBADDRMID, a1).is_some()
        && sdio::cmd52(sdio::SDIO_FUNC_BACKPLANE, true, SBSDIO_FUNC1_SBADDRHIGH, a2).is_some()
}

unsafe fn backplane_read32(addr: u32) -> Option<u32> {
    let window = addr & !0x7FFF;
    if !set_backplane_window(window) {
        return None;
    }
    let offset = (addr & 0x7FFF) | 0x08000; // 32-bit access flag
    let mut buf = [0u8; 4];
    if !sdio::cmd53(sdio::SDIO_FUNC_BACKPLANE, false, offset, &mut buf) {
        return None;
    }
    Some(u32::from_le_bytes(buf))
}

unsafe fn backplane_write32(addr: u32, val: u32) -> bool {
    let window = addr & !0x7FFF;
    if !set_backplane_window(window) {
        return false;
    }
    let offset = (addr & 0x7FFF) | 0x08000;
    sdio::cmd53(
        sdio::SDIO_FUNC_BACKPLANE,
        true,
        offset,
        &mut val.to_le_bytes().clone(),
    )
}

// ============================================================
// Firmware upload and start
// ============================================================

unsafe fn upload_firmware(firmware: &[u8]) -> bool {
    crate::info!("[NET] WiFi: Uploading firmware ({} bytes)...", firmware.len());
    const CHUNK: usize = 64;
    let mut offset = 0usize;
    while offset < firmware.len() {
        let chunk_len = (firmware.len() - offset).min(CHUNK);
        let addr = offset as u32;
        if !set_backplane_window(addr & !0x7FFF) {
            return false;
        }
        let reg = addr & 0x7FFF;
        let mut buf = [0u8; CHUNK];
        buf[..chunk_len].copy_from_slice(&firmware[offset..offset + chunk_len]);
        if !sdio::cmd53(sdio::SDIO_FUNC_BACKPLANE, true, reg, &mut buf[..chunk_len]) {
            crate::info!("[NET] WiFi: firmware write failed at offset 0x{:X}", offset);
            return false;
        }
        offset += chunk_len;
        if offset.is_multiple_of(64 * 1024) {
            crate::info!(
                "WiFi: firmware {}KB / {}KB",
                offset / 1024,
                firmware.len() / 1024
            );
        }
    }
    crate::info!("[NET] WiFi: Firmware uploaded.");
    true
}

unsafe fn start_firmware() -> bool {
    // Release ARM core reset
    if !backplane_write32(BRCMF_BACKPLANE_ARM + 0x800, 0) {
        crate::info!("[NET] WiFi: ARM reset write failed");
        return false;
    }
    delay_ms(50);
    crate::info!("[NET] WiFi: ARM core reset released.");
    true
}

// ============================================================
// Wait for HT clock (firmware ready signal)
// ============================================================

unsafe fn wait_ht_clock() -> bool {
    // Request HT (High Throughput) clock from chip
    sdio::cmd52(
        sdio::SDIO_FUNC_BACKPLANE,
        true,
        SBSDIO_FUNC1_CHIPCLKCSR,
        SBSDIO_HT_AVAIL_REQ | SBSDIO_ALP_AVAIL_REQ,
    );
    // Wait up to 1000ms
    for _ in 0..1000 {
        let clk =
            sdio::cmd52(sdio::SDIO_FUNC_BACKPLANE, false, SBSDIO_FUNC1_CHIPCLKCSR, 0).unwrap_or(0);
        if (clk & SBSDIO_HT_AVAIL) != 0 {
            crate::info!("[NET] WiFi: HT clock available (CSR=0x{:02X})", clk);
            return true;
        }
        delay_ms(1);
    }
    crate::info!("[NET] WiFi: HT clock timeout");
    false
}

// ============================================================
// SDPCM frame transport (over SDIO function 2 WLAN FIFO)
// ============================================================

// Build and send an SDPCM frame on the specified channel.
// `payload` = bytes after the SDPCM header (CDC header + data for control channel).
unsafe fn sdpcm_send(channel: u8, payload: &[u8]) -> bool {
    let total = SDPCM_HDR_LEN + payload.len();
    if total > 512 {
        crate::info!("[NET] WiFi: sdpcm_send frame too large: {}", total);
        return false;
    }
    let mut frame = [0u8; 512];
    // SDPCM header (12 bytes)
    let len = total as u16;
    frame[0] = (len & 0xFF) as u8;
    frame[1] = (len >> 8) as u8;
    frame[2] = !(len as u8);
    frame[3] = !((len >> 8) as u8);
    frame[4] = TX_SEQ;
    TX_SEQ = TX_SEQ.wrapping_add(1);
    frame[5] = channel & 0x0F;
    frame[6] = 0; // next_len
    frame[7] = SDPCM_HDR_LEN as u8; // data offset
    frame[8] = 0; // flow control mask
    frame[9] = 0; // credit
    frame[10] = 0;
    frame[11] = 0;
    frame[SDPCM_HDR_LEN..total].copy_from_slice(payload);
    // Send via CMD53 to WLAN FIFO (func 2, addr 0, no-increment)
    sdio::cmd53(sdio::SDIO_FUNC_WLAN, true, 0, &mut frame[..total])
}

// Try to receive one SDPCM frame from the WLAN FIFO.
// Returns (channel, data_offset, frame_bytes) on success.
unsafe fn sdpcm_recv() -> Option<(u8, usize, Vec<u8>)> {
    // Check if F2 interrupt is pending (CCCR_INT_PENDING reg 0x05, bit 2 = F2)
    let int_pending = sdio::cmd52(sdio::SDIO_FUNC_BUS, false, 0x05, 0)?;
    if (int_pending & 0x04) == 0 {
        return None;
    }

    // Read first 4 bytes to discover frame length
    let mut hdr4 = [0u8; 4];
    if !sdio::cmd53(sdio::SDIO_FUNC_WLAN, false, 0, &mut hdr4) {
        return None;
    }
    let frame_len = u16::from_le_bytes([hdr4[0], hdr4[1]]) as usize;
    let frame_chk = u16::from_le_bytes([hdr4[2], hdr4[3]]) as usize;
    // Validate header checksum
    if frame_len == 0 || (frame_len ^ frame_chk) != 0xFFFF || frame_len < SDPCM_HDR_LEN {
        return None;
    }

    let mut frame = alloc::vec![0u8; frame_len];
    frame[..4].copy_from_slice(&hdr4);

    // Read the rest of the frame (up to 508 bytes per CMD53 call)
    let remaining = frame_len - 4;
    if remaining > 0 {
        let mut offset = 4usize;
        while offset < frame_len {
            let chunk = (frame_len - offset).min(508);
            if !sdio::cmd53(
                sdio::SDIO_FUNC_WLAN,
                false,
                0,
                &mut frame[offset..offset + chunk],
            ) {
                return None;
            }
            offset += chunk;
        }
    }

    let channel = frame[5] & 0x0F;
    let data_off = frame[7] as usize;
    if data_off > frame_len {
        return None;
    }
    Some((channel, data_off, frame))
}

// ============================================================
// CDC ioctl layer
// ============================================================

// Send a CDC ioctl and wait for a matching response.
// `get`: true = GET, false = SET
// `payload_buf`: for SET it contains data to send; for GET it will receive the response.
// Returns number of response bytes written to `payload_buf` on success.
unsafe fn ioctl(cmd: u32, get: bool, payload_buf: &mut [u8]) -> Option<usize> {
    TX_ID = TX_ID.wrapping_add(1);
    let id = TX_ID as u32;
    let flags = if get { CDC_FLAG_GET } else { CDC_FLAG_SET } | (id << CDC_ID_SHIFT);
    let plen = payload_buf.len();

    // Build CDC frame: 16-byte header + payload
    let mut cdc = alloc::vec![0u8; CDC_HDR_LEN + plen];
    cdc[0..4].copy_from_slice(&cmd.to_le_bytes());
    cdc[4..8].copy_from_slice(&(plen as u32).to_le_bytes());
    cdc[8..12].copy_from_slice(&flags.to_le_bytes());
    cdc[12..16].copy_from_slice(&0u32.to_le_bytes()); // status = 0 (request)
    if plen > 0 && !get {
        cdc[CDC_HDR_LEN..].copy_from_slice(payload_buf);
    }

    if !sdpcm_send(SDPCM_CHAN_CONTROL, &cdc) {
        crate::info!("[NET] WiFi: ioctl cmd={} send failed", cmd);
        return None;
    }

    // Poll for matching response (up to 2000ms)
    for _ in 0..2000 {
        if let Some((chan, data_off, frame)) = sdpcm_recv() {
            if chan == SDPCM_CHAN_CONTROL {
                if data_off + CDC_HDR_LEN > frame.len() {
                    continue;
                }
                let cdc_resp = &frame[data_off..];
                let resp_flags =
                    u32::from_le_bytes([cdc_resp[8], cdc_resp[9], cdc_resp[10], cdc_resp[11]]);
                let resp_id = resp_flags >> CDC_ID_SHIFT;
                if resp_id != (id & 0xFFFF) {
                    continue;
                } // id mismatch, keep polling
                let status =
                    u32::from_le_bytes([cdc_resp[12], cdc_resp[13], cdc_resp[14], cdc_resp[15]]);
                if status != 0 {
                    crate::info!("[NET] WiFi: ioctl cmd={} status={}", cmd, status);
                    return None;
                }
                if get {
                    let resp_plen =
                        u32::from_le_bytes([cdc_resp[4], cdc_resp[5], cdc_resp[6], cdc_resp[7]])
                            as usize;
                    let copy_len = resp_plen.min(plen);
                    let payload_start = data_off + CDC_HDR_LEN;
                    if payload_start + copy_len <= frame.len() {
                        payload_buf[..copy_len]
                            .copy_from_slice(&frame[payload_start..payload_start + copy_len]);
                        return Some(copy_len);
                    }
                }
                return Some(0);
            }
            // Discard non-control frames silently (events, data)
        }
        delay_ms(1);
    }
    crate::info!("[NET] WiFi: ioctl cmd={} timeout", cmd);
    None
}

unsafe fn ioctl_set_u32(cmd: u32, val: u32) -> bool {
    ioctl(cmd, false, &mut val.to_le_bytes().clone()).is_some()
}

// ---- Iovar helpers ----

// Build iovar payload: null-terminated name + value bytes
fn build_iovar(name: &str, data: &[u8]) -> Vec<u8> {
    let nb = name.as_bytes();
    let mut v = alloc::vec![0u8; nb.len() + 1 + data.len()];
    v[..nb.len()].copy_from_slice(nb);
    // byte[nb.len()] = 0 (null terminator, already zeroed)
    v[nb.len() + 1..].copy_from_slice(data);
    v
}

unsafe fn iovar_set(name: &str, data: &[u8]) -> bool {
    let mut buf = build_iovar(name, data);
    ioctl(BRCMF_C_SET_VAR, false, &mut buf).is_some()
}

unsafe fn iovar_set_u32(name: &str, val: u32) -> bool {
    iovar_set(name, &val.to_le_bytes())
}

// bsscfg-scoped iovar: name\0 + bsscfg_idx(u32) + value
unsafe fn iovar_set_bsscfg(name: &str, bsscfg: u32, data: &[u8]) -> bool {
    let nb = name.as_bytes();
    let mut buf = alloc::vec![0u8; nb.len() + 1 + 4 + data.len()];
    buf[..nb.len()].copy_from_slice(nb);
    buf[nb.len() + 1..nb.len() + 5].copy_from_slice(&bsscfg.to_le_bytes());
    buf[nb.len() + 5..].copy_from_slice(data);
    ioctl(BRCMF_C_SET_VAR, false, &mut buf).is_some()
}

// ============================================================
// Event reception
// ============================================================

// Poll for a specific event type. Returns event status on success.
// brcmf_event_msg (big-endian) is at: data_off + 4 (BDC) + 14 (Ethernet header)
unsafe fn wait_event(want_event: u32, timeout_ms: u32) -> Option<u32> {
    for _ in 0..timeout_ms {
        if let Some((chan, data_off, frame)) = sdpcm_recv() {
            if chan == SDPCM_CHAN_EVENT {
                // BDC header: 4 bytes at data_off
                // Ethernet header: 14 bytes
                // brcmf_event_msg: starts at data_off + 18
                let msg_off = data_off + 4 + 14; // skip BDC + Ethernet
                if msg_off + 20 > frame.len() {
                    continue;
                }
                // event_type at offset 4, status at offset 8 — big-endian
                let evt = u32::from_be_bytes([
                    frame[msg_off + 4],
                    frame[msg_off + 5],
                    frame[msg_off + 6],
                    frame[msg_off + 7],
                ]);
                let status = u32::from_be_bytes([
                    frame[msg_off + 8],
                    frame[msg_off + 9],
                    frame[msg_off + 10],
                    frame[msg_off + 11],
                ]);
                crate::info!("[NET] WiFi: event={} status={}", evt, status);
                if evt == want_event {
                    return Some(status);
                }
            }
            // Also handle control channel responses during event wait (discard)
        }
        delay_ms(1);
    }
    None
}

// Enable reception of specific events via "event_msgs" iovar (16-byte bitmask)
unsafe fn enable_events() -> bool {
    let mut mask = [0u8; 16];
    // Enable: E_SET_SSID(0), E_AUTH(3), E_DEAUTH(5), E_ASSOC(7), E_DISASSOC(11), E_LINK(16), E_PSK_SUP(46)
    let events = [0u32, 3, 5, 7, 11, 16, 46];
    for e in &events {
        let byte = (*e / 8) as usize;
        let bit = *e % 8;
        if byte < 16 {
            mask[byte] |= 1 << bit;
        }
    }
    iovar_set("event_msgs", &mask)
}

// ============================================================
// HLP (High-Level Protocol) initialization
// ============================================================

unsafe fn hlp_init() -> bool {
    // 1. Wait for HT clock (signals firmware is ready)
    if !wait_ht_clock() {
        return false;
    }

    // 2. Verify chip via GET_REVINFO (48-byte struct)
    let mut revinfo = [0u8; 48];
    if ioctl(BRCMF_C_GET_REVINFO, true, &mut revinfo).is_some() {
        let chip_id = u32::from_le_bytes([revinfo[0], revinfo[1], revinfo[2], revinfo[3]]);
        crate::info!("[NET] WiFi: chip_id=0x{:04X}", chip_id);
    } else {
        crate::info!("[NET] WiFi: GET_REVINFO failed (firmware may still be booting)");
        // Non-fatal: continue
    }

    // 3. Set country code (US, rev 0)
    // wl_country_t: country_abbrev[4] + rev(i32) + ccode[4]
    let mut country = [0u8; 12];
    country[0] = b'U';
    country[1] = b'S'; // country_abbrev
                       // rev = 0 (auto), ccode = "US"
    country[8] = b'U';
    country[9] = b'S';
    if !iovar_set("country", &country) {
        crate::info!("[NET] WiFi: country set failed (non-fatal)");
    }

    // 4. Disable power management for faster response
    if !ioctl_set_u32(BRCMF_C_SET_PM, 0) {
        crate::info!("[NET] WiFi: SET_PM failed (non-fatal)");
    }

    // 5. Enable event reception
    if !enable_events() {
        crate::info!("[NET] WiFi: enable_events failed (non-fatal)");
    }

    crate::info!("[NET] WiFi: HLP protocol initialized.");
    true
}

// ============================================================
// Public API
// ============================================================

/// Initialize WiFi: load firmware, start chip, run HLP init.
pub fn init() -> bool {
    if !sdio::is_initialized() {
        return false;
    }
    crate::info!("[NET] WiFi: Initializing CYW43455...");

    // Load firmware from SylFS
    let fw_bytes = match crate::kernel::fs::read_file("/boot/brcmfmac43455-sdio.bin") {
        Some(d) => d,
        None => {
            crate::info!("[NET] WiFi: firmware not found at /boot/brcmfmac43455-sdio.bin");
            crate::info!("[NET] WiFi: copy brcmfmac43455-sdio.bin to SD card /boot/ folder.");
            return false;
        }
    };
    crate::info!("[NET] WiFi: firmware {} bytes loaded.", fw_bytes.len());

    unsafe {
        // Request ALP clock before firmware upload
        sdio::cmd52(
            sdio::SDIO_FUNC_BACKPLANE,
            true,
            SBSDIO_FUNC1_CHIPCLKCSR,
            SBSDIO_ALP_AVAIL_REQ,
        );
        delay_ms(20);

        if !upload_firmware(&fw_bytes) {
            return false;
        }
        if !start_firmware() {
            return false;
        }
    }

    delay_ms(500); // wait for firmware to boot

    *WIFI_STATE.lock() = WifiState::FirmwareLoaded;

    unsafe {
        if !hlp_init() {
            crate::info!("[NET] WiFi: HLP init failed");
            *WIFI_STATE.lock() = WifiState::Error;
            return false;
        }
    }

    *WIFI_STATE.lock() = WifiState::Idle;
    crate::info!("[NET] WiFi: CYW43455 ready.");
    true
}

/// Connect to WPA2-PSK network.
/// Returns true when associated; IP must then be acquired via DHCP by the caller.
pub fn connect(ssid: &str, password: &str) -> bool {
    {
        let state = *WIFI_STATE.lock();
        if state != WifiState::Idle {
            return false;
        }
    }
    if ssid.is_empty() || ssid.len() > 32 {
        return false;
    }

    crate::info!("[NET] WiFi: Connecting to '{}'...", ssid);
    *WIFI_STATE.lock() = WifiState::Connecting;

    unsafe {
        WIFI_SSID_LEN = ssid.len();
        WIFI_SSID[..ssid.len()].copy_from_slice(ssid.as_bytes());
    }

    unsafe {
        // 1. Set WPA2-AES security
        if !ioctl_set_u32(BRCMF_C_SET_WSEC, WSEC_AES) {
            crate::info!("[NET] WiFi: SET_WSEC failed");
            return fail_connect();
        }

        // 2. Enable supplicant (firmware handles 4-way handshake)
        if !iovar_set_bsscfg("sup_wpa", 0, &1u32.to_le_bytes()) {
            crate::info!("[NET] WiFi: sup_wpa failed (non-fatal)");
        }
        // WPA2-EAPOL version = -1 (auto)
        if !iovar_set_bsscfg("sup_wpa2_eapver", 0, &0xFFFF_FFFFu32.to_le_bytes()) {
            crate::info!("[NET] WiFi: sup_wpa2_eapver failed (non-fatal)");
        }
        // Supplicant timeout 2500ms
        if !iovar_set_bsscfg("sup_wpa_tmo", 0, &2500u32.to_le_bytes()) {
            crate::info!("[NET] WiFi: sup_wpa_tmo failed (non-fatal)");
        }

        // 3. Provide passphrase (wsec_pmk_t: key_len[2] + flags[2] + key[64])
        {
            let pw = password.as_bytes();
            let pw_len = pw.len().min(64);
            let mut pmk = [0u8; 68]; // 2 + 2 + 64
            pmk[0] = pw_len as u8;
            pmk[1] = (pw_len >> 8) as u8;
            pmk[2] = (WSEC_PASSPHRASE & 0xFF) as u8;
            pmk[3] = (WSEC_PASSPHRASE >> 8) as u8;
            pmk[4..4 + pw_len].copy_from_slice(&pw[..pw_len]);
            if !iovar_set("wsec_pmk", &pmk) {
                crate::info!("[NET] WiFi: wsec_pmk failed");
                return fail_connect();
            }
        }

        // 4. Infrastructure mode (BSS)
        if !ioctl_set_u32(BRCMF_C_SET_INFRA, 1) {
            crate::info!("[NET] WiFi: SET_INFRA failed");
            return fail_connect();
        }

        // 5. Open auth (firmware does WPA2 internally)
        if !ioctl_set_u32(BRCMF_C_SET_AUTH, 0) {
            crate::info!("[NET] WiFi: SET_AUTH failed");
            return fail_connect();
        }

        // 6. Join: SET_SSID — wlc_ssid_t: ssid_len(u32) + ssid[32]
        {
            let mut ssid_buf = [0u8; 36];
            let slen = ssid.len().min(32);
            ssid_buf[0] = slen as u8;
            ssid_buf[4..4 + slen].copy_from_slice(&ssid.as_bytes()[..slen]);
            if ioctl(BRCMF_C_SET_SSID, false, &mut ssid_buf).is_none() {
                crate::info!("[NET] WiFi: SET_SSID failed");
                return fail_connect();
            }
        }

        // 7. Wait for BRCMF_E_SET_SSID event (up to 15s)
        crate::info!("[NET] WiFi: Waiting for association (up to 15s)...");
        match wait_event(BRCMF_E_SET_SSID, 15_000) {
            Some(BRCMF_E_STATUS_SUCCESS) => {
                crate::info!("[NET] WiFi: Associated with '{}'!", ssid);
                *WIFI_STATE.lock() = WifiState::Connected;
                true
            }
            Some(status) => {
                crate::info!("[NET] WiFi: Association failed (status={})", status);
                fail_connect()
            }
            None => {
                crate::info!("[NET] WiFi: Association timeout");
                fail_connect()
            }
        }
    }
}

fn fail_connect() -> bool {
    *WIFI_STATE.lock() = WifiState::Error;
    false
}

/// Disconnect from current AP.
pub fn disconnect() -> bool {
    if *WIFI_STATE.lock() != WifiState::Connected {
        return false;
    }
    unsafe {
        ioctl_set_u32(BRCMF_C_DISASSOC, 0);
    }
    *WIFI_STATE.lock() = WifiState::Idle;
    true
}

/// Store IP address (called by DHCP module after lease obtained).
pub fn set_ip(ip: [u8; 4]) {
    unsafe {
        WIFI_IP = ip;
    }
}

// ---- Accessors ----
pub fn state() -> WifiState {
    *WIFI_STATE.lock()
}
pub fn is_connected() -> bool {
    *WIFI_STATE.lock() == WifiState::Connected
}
pub fn get_ip() -> [u8; 4] {
    unsafe { WIFI_IP }
}
pub fn get_ssid() -> &'static str {
    unsafe { core::str::from_utf8(&WIFI_SSID[..WIFI_SSID_LEN]).unwrap_or("") }
}