/*
    * Wpa.hpp
    * WPA2/WPA3-PSK supplicant: PMK derivation and the EAPOL-Key 4-way
    * handshake that unlocks the link after association.
    *
    * The supplicant lives in the kernel because the handshake sits between
    * association and the first IP packet: nothing above the driver can send or
    * receive until the pairwise key is installed in the firmware.  It drives
    * the exchange but does not touch the hardware itself -- transmitting and
    * key installation are provided by the MLME through the three hooks at the
    * bottom of this header.
    *
    * Copyright (c) 2026 Daniel Hammer
*/

#pragma once
#include <cstdint>

namespace Drivers::Net::Wifi {

    enum class WpaState : uint8_t {
        Idle = 0,
        WaitMsg1,        // associated, waiting for the AP to start the exchange
        WaitMsg3,        // msg 2 sent, waiting for the GTK
        Complete,        // keys installed, link is usable
        Failed,
    };

    struct WpaConfig {
        uint8_t  OwnMac[6];
        uint8_t  Bssid[6];
        uint8_t  Ssid[32];
        uint8_t  SsidLen;
        char     Passphrase[64];
        uint8_t  PassLen;
        uint8_t  Akm;               // RSN_AKM_*
        uint8_t  PairwiseCipher;    // RSN_CIPHER_*
        uint8_t  GroupCipher;
        bool     Mfp;               // management frame protection negotiated
    };

    // Derive the PMK and arm the handshake.  Returns false when the
    // configuration names a cipher or AKM this supplicant cannot do.
    bool WpaStart(const WpaConfig& cfg);
    void WpaReset();

    // Feed the 802.1X payload of an inbound EAPOL frame (everything after the
    // LLC/SNAP header).  Returns true when the frame was consumed.
    bool WpaOnEapol(const uint8_t* data, uint32_t len);

    // Re-send the last outbound message if the AP has gone quiet, and fail the
    // handshake once it has been silent for too long.  Called from the idle
    // loop; `nowMs` is a monotonic millisecond clock.
    void WpaService(uint64_t nowMs);

    WpaState WpaGetState();
    bool WpaIsComplete();

    // Build the RSN information element advertising what WpaStart() was
    // configured with.  Returns the number of bytes written, 0 on error.
    uint32_t WpaBuildRsnIe(uint8_t* out, uint32_t cap);

    // Pick the pairwise/group cipher and AKM out of an AP's RSN IE.  `ie`
    // points at the element body (after id/len).  Returns false when nothing
    // in the IE is supported.
    bool WpaParseApRsn(const uint8_t* ie, uint32_t len, WpaConfig& cfg);

    // -------------------------------------------------------------------------
    // Hooks implemented by the MLME (IwxConnect.cpp)
    // -------------------------------------------------------------------------

    // Transmit an EAPOL frame body (802.1X header included) to the AP.
    bool WpaTxEapol(const uint8_t* body, uint32_t len);

    // Install the pairwise temporal key.  `cipher` is an RSN_CIPHER_* value.
    bool WpaInstallPtk(const uint8_t* tk, uint32_t tkLen, uint8_t cipher);

    // Install a group temporal key at `keyIdx`.  `rsc` is the EAPOL key RSC
    // field: 8 bytes, of which the low 6 are the AP's packet number.
    bool WpaInstallGtk(const uint8_t* gtk, uint32_t gtkLen, uint8_t keyIdx,
                       uint8_t cipher, const uint8_t* rsc);
}