/*
    * mlme_harness.cpp
    * Host harness for the 802.11 MLME and data path.
    *
    * Compiles the real IwxConnect.cpp (and the supplicant behind it) against a
    * stubbed transport, so everything the driver puts on the air and everything
    * it makes of what comes back can be checked without the adapter.  The
    * firmware/radio interaction is what remains untestable here; the frame
    * construction, parsing, encapsulation and state machine are not.
    *
    * Speaks hex over stdio; ap_mlme.py is the peer.
    *
    * Copyright (c) 2026 Daniel Hammer
*/

#include <cstdio>
#include <cstring>
#include <cstdint>
#include <string>
#include <iostream>

#include "Drivers/Net/Wifi/Iwx.hpp"
#include "Drivers/Net/Wifi/Ieee80211.hpp"
#include "Drivers/Net/Wifi/Wpa.hpp"

namespace Timekeeping {
    uint64_t g_ms = 1000;
    uint64_t GetMilliseconds() { return g_ms; }
}

// A host command is a round trip to the firmware: it takes real time, and the
// clock moves on while the service loop is inside one.  Modelling that is what
// catches elapsed-time arithmetic that samples the clock once and then compares
// it against timestamps taken later in the same pass (an unsigned underflow
// that reads as an instant timeout).  A frozen clock hides that class of bug
// entirely, which is why this is not simply left at a constant.
static constexpr uint64_t CMD_ROUND_TRIP_MS = 1;

using namespace Drivers::Net::Wifi;

static void puthex(const char* tag, const uint8_t* p, uint32_t n) {
    printf("%s ", tag);
    for (uint32_t i = 0; i < n; i++) printf("%02x", p[i]);
    printf("\n");
}

// =============================================================================
// Stubbed transport
// =============================================================================

namespace Drivers::Net::Wifi {

    IwxState g_iwx;

    // Host commands: record the opcode and the exact bytes, always succeed.
    // The payload matters -- a struct that does not match the version the
    // firmware advertises asserts it on real hardware.
    static void DumpCmd(uint32_t id, const void* data, uint32_t len) {
        printf("CMD %u ", id);
        const uint8_t* p = (const uint8_t*)data;
        for (uint32_t i = 0; i < len; i++) printf("%02x", p[i]);
        printf("\n");
        Timekeeping::g_ms += CMD_ROUND_TRIP_MS;
    }
    bool IwxSendCmdPdu(uint32_t id, const void* data, uint32_t len) {
        DumpCmd(id, data, len);
        return true;
    }
    bool IwxSendCmdStatus(uint32_t id, const void* data, uint32_t len,
                          uint32_t* statusOut) {
        DumpCmd(id, data, len);
        // ADD_STA reports success in the low byte; everything else uses 0.
        if (statusOut) *statusOut = (id == IWX_ADD_STA || id == IWX_ADD_STA_KEY)
                                    ? IWX_ADD_STA_SUCCESS : 0;
        return true;
    }
    bool IwxSendCmd(IwxHostCmd& cmd) { (void)cmd; return true; }

    int IwxLookupCmdVer(uint8_t group, uint8_t cmd) {
        // Match what AX211 firmware 89 advertises for the versions the MLME
        // branches on.
        if (group == IWX_DATA_PATH_GROUP && cmd == IWX_RLC_CONFIG_CMD) return 2;
        if (group == IWX_DATA_PATH_GROUP && cmd == IWX_SCD_QUEUE_CONFIG_CMD) return 3;
        return -1;
    }
    int IwxLookupNotifVer(uint8_t, uint8_t) { return 7; }

    bool IwxAbortScan() { return true; }

    bool IwxEnableTxq(IwxTxRing& ring, int staId, int qid, int tid) {
        (void)staId; (void)tid;
        ring.Qid = qid;
        ring.Active = true;
        ring.StageSlots = IWX_TX_STAGE_SLOTS;
        printf("TXQ-UP %d\n", qid);
        return true;
    }
    void IwxDisableTxq(IwxTxRing& ring, int staId, int tid) {
        (void)staId; (void)tid;
        ring.Active = false;
        printf("TXQ-DOWN\n");
    }

    // Capture what the driver wants to put on the air.
    bool IwxTxFrame(IwxTxRing& ring, const uint8_t* hdr, uint32_t hdrLen,
                    const uint8_t* payload, uint32_t payloadLen,
                    bool encrypt, bool fixedRate) {
        if (!ring.Active) { printf("TX-DROP queue-down\n"); return false; }
        printf("TX enc=%d rate=%d ", encrypt ? 1 : 0, fixedRate ? 1 : 0);
        for (uint32_t i = 0; i < hdrLen; i++) printf("%02x", hdr[i]);
        printf(" ");
        for (uint32_t i = 0; i < payloadLen; i++) printf("%02x", payload[i]);
        printf("\n");
        return true;
    }

    bool IwxSetKey(const uint8_t* key, uint32_t keyLen, uint8_t keyIdx,
                   bool pairwise, uint8_t cipher, const uint8_t* rsc) {
        printf("KEY pairwise=%d idx=%u cipher=%u ", pairwise ? 1 : 0, keyIdx, cipher);
        for (uint32_t i = 0; i < keyLen; i++) printf("%02x", key[i]);
        printf(" ");
        if (rsc) for (int i = 0; i < 6; i++) printf("%02x", rsc[i]);
        printf("\n");
        return true;
    }

    bool IwxRemoveKey(uint8_t keyIdx, bool pairwise, uint8_t cipher,
                      uint32_t keyLen) {
        printf("KEY-REMOVE pairwise=%d idx=%u cipher=%u len=%u\n",
               pairwise ? 1 : 0, keyIdx, cipher, keyLen);
        return true;
    }

    // Sink normally provided by Wifi.cpp.
    void WifiRxEthernet(const uint8_t* frame, uint32_t len) {
        puthex("ETH", frame, len);
    }
}

// =============================================================================
// Driver
// =============================================================================

static int unhex(const std::string& s, uint8_t* out) {
    int n = 0;
    for (size_t i = 0; i + 1 < s.size(); i += 2) {
        unsigned v; sscanf(s.c_str() + i, "%2x", &v); out[n++] = (uint8_t)v;
    }
    return n;
}

int main() {
    // A firmware state good enough for the MLME: alive, one antenna, a MAC.
    g_iwx.State = IwxFwState::Running;
    g_iwx.Fw.PhyConfig = (1u << IWX_FW_PHY_CFG_TX_CHAIN_POS)
                       | (1u << IWX_FW_PHY_CFG_RX_CHAIN_POS);
    g_iwx.Nvm.ValidTxAnt = 1;
    g_iwx.Nvm.ValidRxAnt = 1;

    std::string line;
    uint8_t buf[4096];

    while (std::getline(std::cin, line)) {
        if (line.rfind("MAC ", 0) == 0) {
            unhex(line.substr(4), g_iwx.Nvm.HwAddr);
            printf("DONE\n");
        } else if (line.rfind("CAPA ", 0) == 0) {
            unsigned bit, on;
            sscanf(line.c_str(), "CAPA %u %u", &bit, &on);
            if (on) g_iwx.Fw.Capa[bit / 8] |= (uint8_t)(1 << (bit % 8));
            else    g_iwx.Fw.Capa[bit / 8] &= (uint8_t)~(1 << (bit % 8));
            printf("DONE\n");
        } else if (line.rfind("CONNECT ", 0) == 0) {
            // CONNECT <bssid> <channel> <is5> <ssid> <pass|-> <rsnie|->
            char bssid[64], ssid[64], pass[128], rsn[256];
            unsigned chan, is5;
            sscanf(line.c_str(), "CONNECT %63s %u %u %63s %127s %255s",
                   bssid, &chan, &is5, ssid, pass, rsn);
            uint8_t bs[6]; unhex(bssid, bs);
            uint8_t ie[128]; int ieLen = 0;
            if (strcmp(rsn, "-") != 0) ieLen = unhex(rsn, ie);
            bool ok = IwxConnectStart(bs, (uint8_t)chan, is5 != 0, ssid,
                                      strcmp(pass, "-") == 0 ? nullptr : pass,
                                      ieLen ? ie : nullptr, (uint32_t)ieLen,
                                      100, 2);
            printf("CONNECT-OK %d STATE %d\n", ok ? 1 : 0, IwxConnectState());
        } else if (line.rfind("RXMGMT ", 0) == 0) {
            int n = unhex(line.substr(7), buf);
            IwxConnectRxMgmt(buf, (uint32_t)n);
            printf("DONE STATE %d\n", IwxConnectState());
        } else if (line.rfind("RXDATA ", 0) == 0) {
            int n = unhex(line.substr(7), buf);
            IwxConnectRxData(buf, (uint32_t)n);
            printf("DONE STATE %d\n", IwxConnectState());
        } else if (line.rfind("TXETH ", 0) == 0) {
            int n = unhex(line.substr(6), buf);
            bool ok = IwxConnectSendEthernet(buf, (uint32_t)n);
            printf("TXETH-OK %d\n", ok ? 1 : 0);
        } else if (line.rfind("TXSTATUS ", 0) == 0) {
            // TXSTATUS <acked> <count> -- report N transmit outcomes, the way
            // IwxTxComplete does off the firmware's TX response.  Drives the
            // link supervision that notices an access point which stopped
            // acknowledging without ever deauthenticating.
            unsigned acked, count;
            sscanf(line.c_str(), "TXSTATUS %u %u", &acked, &count);
            for (unsigned i = 0; i < count; i++) IwxConnectNoteTx(acked != 0);
            printf("DONE STATE %d LINK %d\n", IwxConnectState(), IwxLinkUp() ? 1 : 0);
        } else if (line.rfind("SERVICE", 0) == 0) {
            IwxConnectService();
            printf("DONE STATE %d LINK %d\n", IwxConnectState(), IwxLinkUp() ? 1 : 0);
        } else if (line.rfind("TICK ", 0) == 0) {
            unsigned ms; sscanf(line.c_str(), "TICK %u", &ms);
            Timekeeping::g_ms += ms;
            printf("DONE\n");
        } else if (line.rfind("ABORT", 0) == 0) {
            IwxConnectAbort();
            printf("DONE STATE %d\n", IwxConnectState());
        } else if (line.rfind("STATE", 0) == 0) {
            printf("STATE %d LINK %d\n", IwxConnectState(), IwxLinkUp() ? 1 : 0);
        } else if (line.rfind("QUIT", 0) == 0) {
            break;
        }
        fflush(stdout);
    }
    return 0;
}