kbf/include/kbf/wifi.h

#ifndef KBF_WIFI_H
#define KBF_WIFI_H

#include <string>
#include <memory>
#include <map>

#include <esp_wifi.h>
#include <esp_log.h>

#include "net.h"

using std::string;
using std::shared_ptr;
using std::map;

/** @brief WiFi functions */
namespace kbf::wifi {
    /** @brief Tag used for logging */
    static constexpr const char *TAG = "kbf::wifi";

    class AP;

    class STA;

    /**
     * @brief Starts WiFi in soft-AP mode.
     *
     * @see #kbf::wifi::AP::create()
     *
     * @param ap AP pointer
     */
    void start(shared_ptr<AP> ap);

    /**
     * @brief Starts WiFi in STA mode.
     *
     * @param sta STA pointer; if nullptr, kbf::wifi::STA::create() will be called.
     */
    void start(shared_ptr<STA> sta = nullptr);

    /**
     * @brief Starts WiFi in AP+STA dual mode.
     *
     * @warning not yet implemented
     *
     * @param ap
     * @param sta
     */
    void start(shared_ptr<AP> ap, shared_ptr<STA> sta);

    /**
     * @brief Stops WiFi.
     */
    void stop();

    /**
     * @brief Returns the currently used AP instance.
     *
     * @return AP pointer or nullptr if not in AP mode
     */
    shared_ptr<AP> getAP();

    /**
     * @brief Returns the currently used STA instance.
     *
     * @return STA pointer or nullptr if not in STA mode
     */
    shared_ptr<STA> getSTA();

    /**
     * @brief WiFi station (STA) mode.
     */
    class STA {
        /** @see #kbf::wifi::start() */
        friend void kbf::wifi::start(shared_ptr<STA>);

        /** @see #kbf::wifi::start() */
        friend void kbf::wifi::start(shared_ptr<AP>, shared_ptr<STA>);

        /** @see #kbf::wifi::start() */
        friend void kbf::wifi::stop();

    public:
        /** @brief Tag used for logging. */
        static constexpr const char *TAG = "kbf::wifi::STA";

        /**
         * @brief Creates STA instance.
         *
         * @note WiFi requires NVS. If NVS hasn't been initialized yet, the constructor will call #kbf::nvs::init()
         * and try again.
         *
         * @see #kbf::wifi::start()
         *
         * @return shared_ptr<STA>
         */
        static shared_ptr<STA> create();

        /**
         * Destructor. Unregisters event handlers from the default event loop.
         */
        ~STA();

        /**
         * @brief Initiates connection to an AP.
         *
         * @param ssid SSID of AP
         * @param password Password
         * @param maxRetryAttempts maximum number of retry attempts if the connection fails
         */
        void connect(const string &ssid, const string &password, int maxRetryAttempts = 3);

        /**
         * @brief Disconnects from the AP.
         */
        void disconnect();

        /**
         * @brief Called when the WiFi has started.
         */
        void (*onStart)(){};

        /**
         * @brief Called when the WiFi has stopped.
         */
        void (*onStop)(){};

        /**
         * @brief Called when the connection to the AP was successful.
         *
         * @note Connection may be useless until an IP address is assigned by the AP. See onIp().
         * @see onIp()
         */
        void (*onConnect)(){};

        /**
         * @brief Called once an IP address is assigned by the AP. The connection should be ready to use at this point.
         */
        void (*onIp)(){};

        /**
         * @brief Called before reconnect attempts.
         *
         * Reconnect is attempted after disconnect events if retryMax hasn't been reached yet.
         *
         * @note The retry attempt counter is reset after successful connection.
         *
         * @param attempt number of current attempt
         * @param max maximum number of attempts
         * @return if false is returned, reconnect will be cancelled
         */
        bool (*onReconnecting)(int attempt, int max){};

        /**
         * @brief Called if disconnected and retryMax has been reached.
         */
        void (*onDisconnect)(){};

        /** @brief Returns the IP address for the connection. */
        [[nodiscard]] const kbf::net::IP &ip() const { return m_ip; };

    private:
        STA();

        int retryNum = 0;
        int retryMax{};

        void registerEventHandlers();

        void unregisterEventHandlers();

        static void handleStart(void *arg, esp_event_base_t baseType, int32_t eventId, void *);

        static void handleStop(void *arg, esp_event_base_t baseType, int32_t eventId, void *);

        static void handleConnect(void *arg, esp_event_base_t baseType, int32_t eventId, void *);

        static void handleDisconnect(void *arg, esp_event_base_t baseType, int32_t eventId, void *);

        static void handleGotIp(void *arg, esp_event_base_t baseType, int32_t eventId, void *pEventData);

        esp_netif_obj                *netif            = nullptr;
        wifi_sta_config_t            config            = {};
        esp_event_handler_instance_t startHandler      = {};
        esp_event_handler_instance_t stopHandler       = {};
        esp_event_handler_instance_t connectHandler    = {};
        esp_event_handler_instance_t disconnectHandler = {};
        esp_event_handler_instance_t gotIpHandler      = {};

        kbf::net::IP m_ip;
    };

    /**
     * @brief WiFi Software-enabled AccessPoint (AP) mode
     *
     * To start AP, use the #create() method to create an instance of this class, then pass it to #kbf::wifi::start().
     */
    class AP {
        /** @see #kbf::wifi::start() */
        friend void kbf::wifi::start(shared_ptr<AP>);

        /** @see #kbf::wifi::start() */
        friend void kbf::wifi::start(shared_ptr<AP>, shared_ptr<STA>);

        /** @see #kbf::wifi::start() */
        friend void kbf::wifi::stop();

    public:
        /** @brief Tag used for logging. */
        static constexpr const char *TAG = "kbf::wifi::AP";

        /**
         * @brief Creates an instance of the class.
         *
         * @param ssid
         * @param password
         * @return shared_ptr<AP>
         */
        static shared_ptr<AP> create(const string &ssid, const string &password);

        // TODO figure out why this doesn't work then use make_shared
//        template<typename ...Arg> std::shared_ptr<AP> static create(Arg&&...arg) {
//            struct EnableMakeShared : public AP {
//                explicit EnableMakeShared(Arg&&...arg) :AP(std::forward<Arg>(arg)...) {}
//            };
//            return std::make_shared<EnableMakeShared>(std::forward<Arg>(arg)...);
//        }

        /**
         * @brief Destructor; will unregister event handlers from the default event loop.
         */
        ~AP();

        /** @brief SSID to use. */
        const string ssid;

        /** @brief WPA2/PSK password. */
        const string password;

        /**
         * @brief Represents a station connected to this AP.
         */
        class STA {
        public:
            STA() : aid(0), mac{} {};

            explicit STA(wifi_event_ap_staconnected_t &);

            const int           aid;
            const kbf::net::MAC mac;
        };

        /** @brief aid mapped to STA instances */
        map<int, STA> stations{};

        /**
         * @brief Called when the AP has started.
         *
         * @param ap
         */
        void (*onStart)(AP &ap){};

        /**
         * @brief Called when the AP has stopped.
         *
         * @param ap
         */
        void (*onStop)(AP &ap){};

        /**
         * @brief Called when a station connects to the AP.
         *
         * @param ap
         * @param sta
         */
        void (*onConnect)(AP &ap, STA &sta){};

        /**
         * @brief Called when a station disconnects from the AP.
         *
         * @param ap
         * @param sta
         */
        void (*onDisconnect)(AP &ap, STA &sta){};

    private:
        AP() { ESP_LOGW(TAG, "running default constructor!"); }

        AP(string ssid, string password);

        void registerEventHandlers();

        void unregisterEventHandlers();

        static void handleStart(void *arg, esp_event_base_t baseType, int32_t eventId, void *);

        static void handleStop(void *arg, esp_event_base_t baseType, int32_t eventId, void *);

        static void handleConnect(void *arg, esp_event_base_t baseType, int32_t eventId, void *);

        static void handleDisconnect(void *arg, esp_event_base_t baseType, int32_t eventId, void *);

        esp_netif_obj                *netif            = nullptr;
        wifi_ap_config_t             config            = {};
        esp_event_handler_instance_t startHandler      = {};
        esp_event_handler_instance_t stopHandler       = {};
        esp_event_handler_instance_t connectHandler    = {};
        esp_event_handler_instance_t disconnectHandler = {};
    };

}

#endif //KBF_WIFI_H