#include <Arduino.h>
#include "appcontext.hpp"
#include "config.h"
#include "mqtt.h"
#include "pin.h"
#include "statemachine.hpp"

enum class StateId
{
    StartState,
    NetworkInitializeState,
    RunningState
};

enum class EventId
{
    PinInitialized,
    WifiConnected
};

class StartState : public State
{
public:
    StartState(AppContext *appContext) : State("StartState", StateId::StartState), appContext(appContext) {}

    void onEnter(StateMachineBase &stateMachine) override {
        if (appContext) {
            appContext->pinR->setLedLevel(0);
            appContext->pinG->setLedLevel(0);
            appContext->pinB->setLedLevel(0);
            appContext->pinCW->setLedLevel(0);
            appContext->pinWW->setLedLevel(0);
        }
        stateMachine.postEvent(EventId::PinInitialized);
    }

    void onExit(StateMachineBase &stateMachine) override {
    }

    void onUpdate(StateMachineBase &stateMachine) override {
    }

private:
    AppContext *appContext = nullptr;
};

class NetworkInitializeState : public State
{
public:
    NetworkInitializeState(AppContext *appContext) : State("NetworkInitializeState", StateId::NetworkInitializeState), appContext(appContext) {}

    void onEnter(StateMachineBase &stateMachine) override {
        if (appContext) {
            appContext->network = new Network(hostName, friendlyName);
        }
    }

    void onExit(StateMachineBase &stateMachine) override {
        if (appContext && appContext->network) {
            appContext->network->registerMDNS();
            appContext->otaHandler = new OTAHandler(hostName);
        }
        Mqtt::connect(mqttBroker.data(), 1883, "smart_rgb_dev", "mqtt", "mqtt");
    }

    void onUpdate(StateMachineBase &stateMachine) override {
        if (appContext && appContext->network && appContext->network->isConnected()) {
            stateMachine.postEvent(EventId::WifiConnected);
        }
    }

private:
    AppContext *appContext = nullptr;
};

class RunningState : public State
{
public:
    RunningState(AppContext *appContext) : State("RunningState", StateId::RunningState), appContext(appContext) {}
    void onEnter(StateMachineBase &stateMachine) override {
        lastOtaPollMs = millis();
        lastMqttPollMs = millis();
        lastMqttCheckConnectionPollSecond = millis() / 1000;
    }

    void onExit(StateMachineBase &stateMachine) override {
    }

    void onUpdate(StateMachineBase &stateMachine) override {
        if ((millis() - lastOtaPollMs) >= otaPollInterval) {
            lastOtaPollMs = millis();
            if (appContext && appContext->otaHandler) {
                appContext->otaHandler->poll();
            }
        }
        if ((millis() - lastMqttPollMs) >= mqttPollInterval) {
            lastMqttPollMs = millis();
            Mqtt::poll();
        }
        if ((millis() - lastMqttCheckConnectionPollSecond) >= mqttCheckConnectionPollIntervalSecond * 1000) {
            lastMqttCheckConnectionPollSecond = millis() / 1000;
            Mqtt::checkConnection();
        }
    }

private:
    AppContext *appContext = nullptr;
    uint32_t lastOtaPollMs = 0;
    static constexpr uint32_t otaPollInterval = 1000; // Poll every second
    uint32_t lastMqttPollMs = 0;
    static constexpr uint32_t mqttPollInterval = 100; // Poll every 100 milliseconds
    uint32_t lastMqttCheckConnectionPollSecond = 0;
    static constexpr uint32_t mqttCheckConnectionPollIntervalSecond = 30; // Poll every 30 seconds
};