// Copyright (C) 2024 The Qt Company Ltd.
// SPDX-License-Identifier: LicenseRef-Qt-Commercial
#ifndef STATEMACHINE_H
#define STATEMACHINE_H

#include <etl/map.h>
#include <etl/multimap.h>
#include <qul/timer.h>

template<typename TStateId,
         const size_t MaxStates = 10,
         typename TLayerId = int,
         const size_t NumLayers = 1,
         const size_t ChangeStateBufferSize = NumLayers * 10,
         const size_t UpdateStateBufferSize = NumLayers * 10>
class StateMachine
{
public:
    struct AbstractState
    {
        typedef TStateId Id;
        typedef StateMachine SM;
        typedef TLayerId LayerId;

        virtual ~AbstractState() {}
        virtual void onEnter(const TStateId &prevState, const TLayerId &layerId, StateMachine &sm) = 0;
        virtual void onUpdate(uint32_t tick, const TLayerId &layerId, StateMachine &sm) = 0;
        virtual void onLeave(const TStateId &nextState, const TLayerId &layerId, StateMachine &sm) = 0;
    };

    typedef AbstractState BaseStateType;
    typedef TStateId StateId;

    StateMachine(const TStateId &idleStateId);

    bool registerState(const TStateId &id, BaseStateType *state);
    void changeState(const TStateId &nextStateId, const TLayerId &layerId = TLayerId(), uint32_t delay = 0);
    void update(uint32_t tick, const TLayerId &layerId = TLayerId());

    bool requestUpdate(uint32_t interval, bool repeat, const TLayerId &layerId = TLayerId());
    bool dismissUpdate(const TLayerId &layerId = TLayerId());
    bool dismissScheduledStateChanges(const TLayerId &layerId = TLayerId());

    TStateId getCurrentStateId(const TLayerId &layerId = TLayerId()) const;

private:
    struct DelayedTask
    {
        DelayedTask(const TLayerId &layerId_, StateMachine &sm_)
            : layerId(layerId_)
            , sm(&sm_)
        {
            TimerElapsed e = {this};
            timer.onTimeout(e);
        }

        DelayedTask(const DelayedTask &task)
        {
            layerId = task.layerId;
            sm = task.sm;

            // Don't copy the timer
            timer.stop();
            timer.setInterval(task.timer.interval());
            timer.setSingleShot(task.timer.isSingleShot());
            TimerElapsed e = {this};
            timer.onTimeout(e);
        }

        DelayedTask &operator=(const DelayedTask &task)
        {
            layerId = task.layerId;
            sm = task.sm;

            // Don't copy the timer
            timer.stop();

            timer.setInterval(task.timer.interval());
            timer.setSingleShot(task.timer.isSingleShot());
            return *this;
        }

        virtual ~DelayedTask() { timer.stop(); }

        void schedule(bool repeat, uint32_t timeout)
        {
            timer.setSingleShot(!repeat);
            timer.start(timeout);
        }

        virtual void doTask() = 0;

        struct TimerElapsed
        {
            DelayedTask *self;
            void operator()() { self->doTask(); }
        };

    protected:
        TLayerId layerId;
        StateMachine *sm;
        Qul::Timer timer;
    };

    struct UpdateStateTask : DelayedTask
    {
        UpdateStateTask(const TLayerId &layerId_, StateMachine &sm_)
            : DelayedTask(layerId_, sm_)
        {}

        void doTask()
        {
            // TODO: Can try to calculate real time delta
            this->sm->update(this->timer.interval(), this->layerId);

            if (this->timer.isSingleShot()) {
                this->sm->removeUpdateStateTask(this->layerId, this);
            }
        };
    };

    struct ChangeStateTask : DelayedTask
    {
        ChangeStateTask(const TStateId &nextStateId_, const TLayerId &layerId_, StateMachine &sm_)
            : DelayedTask(layerId_, sm_)
            , nextStateId(nextStateId_)
        {}

        void doTask()
        {
            this->sm->doStateChange(this->nextStateId, this->layerId);
            this->sm->removeStateChangeTask(this->layerId, this);
        };

    private:
        TStateId nextStateId;
    };

    struct Tasks
    {
        typedef etl::multimap<TLayerId, ChangeStateTask, ChangeStateBufferSize> ChangeStateTasks;
        typedef etl::multimap<TLayerId, UpdateStateTask, UpdateStateBufferSize> UpdateStateTasks;

        ChangeStateTasks changeState;
        UpdateStateTasks updateState;
    };

    typedef etl::map<TStateId, AbstractState *, MaxStates> StatesMap;
    typedef std::pair<TStateId, AbstractState *> StateEntry;
    typedef etl::map<TLayerId, StateEntry, NumLayers> Layers;

    bool doStateChange(const TStateId &newStateId, const TLayerId &layerId);
    bool scheduleStateChange(const TStateId &newStateId, const TLayerId &layerId, uint32_t timeout);
    bool removeStateChangeTask(const TLayerId &layerId, ChangeStateTask *task);
    bool removeUpdateStateTask(const TLayerId &layerId, UpdateStateTask *task);

    StatesMap _stateRegistry;
    Layers _layers;
    Tasks _tasks;

    const TStateId _idleStateId;
};

template<typename TStateId,
         const size_t MaxStates,
         typename TLayerId,
         const size_t NumLayers,
         const size_t ChangeStateBufferSize,
         const size_t UpdateStateBufferSize>
StateMachine<TStateId, MaxStates, TLayerId, NumLayers, ChangeStateBufferSize, UpdateStateBufferSize>::StateMachine(
    const TStateId &idleStateId)
    : _idleStateId(idleStateId)
{}

template<typename TStateId,
         const size_t MaxStates,
         typename TLayerId,
         const size_t NumLayers,
         const size_t ChangeStateBufferSize,
         const size_t UpdateStateBufferSize>
bool StateMachine<TStateId, MaxStates, TLayerId, NumLayers, ChangeStateBufferSize, UpdateStateBufferSize>::registerState(
    const TStateId &id, BaseStateType *state)
{
    if (_stateRegistry.size() >= MaxStates) {
        return false;
    }

    if (_stateRegistry.count(id) > 0) {
        return false;
    }

    _stateRegistry[id] = state;
    return true;
}

template<typename TStateId,
         const size_t MaxStates,
         typename TLayerId,
         const size_t NumLayers,
         const size_t ChangeStateBufferSize,
         const size_t UpdateStateBufferSize>
void StateMachine<TStateId, MaxStates, TLayerId, NumLayers, ChangeStateBufferSize, UpdateStateBufferSize>::changeState(
    const TStateId &nextStateId, const TLayerId &layer, uint32_t delay)
{
    if (delay == 0) {
        doStateChange(nextStateId, layer);
    } else {
        scheduleStateChange(nextStateId, layer, delay);
    }
}

template<typename TStateId,
         const size_t MaxStates,
         typename TLayerId,
         const size_t NumLayers,
         const size_t ChangeStateBufferSize,
         const size_t UpdateStateBufferSize>
void StateMachine<TStateId, MaxStates, TLayerId, NumLayers, ChangeStateBufferSize, UpdateStateBufferSize>::update(
    uint32_t tick, const TLayerId &layerId)
{
    typename Layers::iterator layerIt = _layers.find(layerId);
    if (layerIt != _layers.end()) {
        layerIt->second.second->onUpdate(tick, layerId, *this);
    }
}

template<typename TStateId,
         const size_t MaxStates,
         typename TLayerId,
         const size_t NumLayers,
         const size_t ChangeStateBufferSize,
         const size_t UpdateStateBufferSize>
TStateId
StateMachine<TStateId, MaxStates, TLayerId, NumLayers, ChangeStateBufferSize, UpdateStateBufferSize>::getCurrentStateId(
    const TLayerId &layerId) const
{
    typename Layers::const_iterator layerIt = _layers.find(layerId);
    if (layerIt == _layers.end()) {
        // No layer found (layer was idle)
        return _idleStateId;
    }

    return layerIt->second.first;
}

template<typename TStateId,
         const size_t MaxStates,
         typename TLayerId,
         const size_t NumLayers,
         const size_t ChangeStateBufferSize,
         const size_t UpdateStateBufferSize>
bool StateMachine<TStateId, MaxStates, TLayerId, NumLayers, ChangeStateBufferSize, UpdateStateBufferSize>::doStateChange(
    const TStateId &nextStateId, const TLayerId &layer)
{
    // Check if state is present on any layer
    typename Layers::iterator layerIt;
    for (layerIt = _layers.begin(); layerIt != _layers.end(); ++layerIt) {
        if (layerIt->second.first == nextStateId) {
            // next state is already present on other layer
            break;
        }
    }

    if (layerIt != _layers.end()) {
        // nextState already exist on different layer
        if (layerIt->first == layer) {
            // Change for the current state requested
            // Nothing to do here
            return false;
        } else {
            // Move state between layers
            // Get current state from target layer
            typename Layers::iterator targetLayerIt = _layers.find(layer);
            TStateId prevStateId = _idleStateId;
            if (targetLayerIt != _layers.end()) {
                // Notify old state from target layer if any
                AbstractState *currentStateFromTargetLayer = targetLayerIt->second.second;
                prevStateId = targetLayerIt->second.first;
                currentStateFromTargetLayer->onLeave(nextStateId, layer, *this);
            }

            assert(nextStateId == layerIt->second.first);

            // Leave old layer
            AbstractState *currentState = layerIt->second.second;
            currentState->onLeave(_idleStateId, layerIt->first, *this);

            // Enter new layer
            currentState->onEnter(prevStateId, layer, *this);
            _layers[layer] = layerIt->second;

            // Clear old layer
            _layers.erase(layerIt);
            return true;
        }
    }

    // Find next state
    typename StatesMap::iterator nextStateIt = _stateRegistry.end();
    if (nextStateId != _idleStateId) {
        nextStateIt = _stateRegistry.find(nextStateId);
        if (nextStateIt == _stateRegistry.end()) {
            // Target state not found! This should not happen.
            assert(false);
            return false;
        }
    }

    // Find target layer
    typename Layers::iterator targetLayerIt = _layers.find(layer);
    if (targetLayerIt == _layers.end()) {
        // No state on this layer yet
        if (nextStateIt == _stateRegistry.end()) {
            // No new state is idle - nothing to do
            return false;
        }

        nextStateIt->second->onEnter(_idleStateId, layer, *this);
        _layers[layer] = *nextStateIt;
    } else {
        // State already exists on this layer
        if (nextStateIt == _stateRegistry.end()) {
            // Next state is idle
            targetLayerIt->second.second->onLeave(_idleStateId, layer, *this);
            _layers.erase(targetLayerIt);
        } else {
            // Switch states
            targetLayerIt->second.second->onLeave(nextStateIt->first, layer, *this);
            nextStateIt->second->onEnter(targetLayerIt->second.first, layer, *this);
            _layers[layer] = *nextStateIt;
        }
    }
    return true;
}

template<typename TStateId,
         const size_t MaxStates,
         typename TLayerId,
         const size_t NumLayers,
         const size_t ChangeStateBufferSize,
         const size_t UpdateStateBufferSize>
bool StateMachine<TStateId, MaxStates, TLayerId, NumLayers, ChangeStateBufferSize, UpdateStateBufferSize>::
    scheduleStateChange(const TStateId &newStateId, const TLayerId &layerId, uint32_t timeout)
{
    if (_tasks.changeState.size() >= _tasks.changeState.max_size()) {
        assert(!"Tasks capacity exceeded");
        return false;
    }
    typename Tasks::ChangeStateTasks::iterator result = _tasks.changeState.insert(
        typename Tasks::ChangeStateTasks::value_type(layerId, ChangeStateTask(newStateId, layerId, *this)));
    result->second.schedule(false, timeout);
    return true;
}

template<typename TStateId,
         const size_t MaxStates,
         typename TLayerId,
         const size_t NumLayers,
         const size_t ChangeStateBufferSize,
         const size_t UpdateStateBufferSize>
bool StateMachine<TStateId, MaxStates, TLayerId, NumLayers, ChangeStateBufferSize, UpdateStateBufferSize>::
    removeStateChangeTask(const TLayerId &layerId, ChangeStateTask *task)
{
    std::pair<typename Tasks::ChangeStateTasks::iterator, typename Tasks::ChangeStateTasks::iterator> scope
        = _tasks.changeState.equal_range(layerId);
    for (; scope.first != scope.second; ++scope.first) {
        if (&scope.first->second == task) {
            _tasks.changeState.erase(scope.first);
            return true;
        }
    }
    assert(!"Task not found");
    return false;
}

template<typename TStateId,
         const size_t MaxStates,
         typename TLayerId,
         const size_t NumLayers,
         const size_t ChangeStateBufferSize,
         const size_t UpdateStateBufferSize>
bool StateMachine<TStateId, MaxStates, TLayerId, NumLayers, ChangeStateBufferSize, UpdateStateBufferSize>::requestUpdate(
    uint32_t interval, bool repeat, const TLayerId &layerId)
{
    if (_tasks.updateState.size() >= _tasks.updateState.max_size()) {
        assert(!"Tasks capacity exceeded");
        return false;
    }
    typename Tasks::UpdateStateTasks::iterator result = _tasks.updateState.insert(
        typename Tasks::UpdateStateTasks::value_type(layerId, UpdateStateTask(layerId, *this)));
    result->second.schedule(repeat, interval);
    return true;
}

template<typename TStateId,
         const size_t MaxStates,
         typename TLayerId,
         const size_t NumLayers,
         const size_t ChangeStateBufferSize,
         const size_t UpdateStateBufferSize>
bool StateMachine<TStateId, MaxStates, TLayerId, NumLayers, ChangeStateBufferSize, UpdateStateBufferSize>::
    removeUpdateStateTask(const TLayerId &layerId, UpdateStateTask *task)
{
    std::pair<typename Tasks::UpdateStateTasks::iterator, typename Tasks::UpdateStateTasks::iterator> scope
        = _tasks.updateState.equal_range(layerId);
    for (; scope.first != scope.second; ++scope.first) {
        if (&scope.first->second == task) {
            _tasks.updateState.erase(scope.first);
            return true;
        }
    }
    assert(!"Task not found");
    return false;
}

template<typename TStateId,
         const size_t MaxStates,
         typename TLayerId,
         const size_t NumLayers,
         const size_t ChangeStateBufferSize,
         const size_t UpdateStateBufferSize>
bool StateMachine<TStateId, MaxStates, TLayerId, NumLayers, ChangeStateBufferSize, UpdateStateBufferSize>::dismissUpdate(
    const TLayerId &layerId)
{
    return _tasks.updateState.erase(layerId) > 0;
}
template<typename TStateId,
         const size_t MaxStates,
         typename TLayerId,
         const size_t NumLayers,
         const size_t ChangeStateBufferSize,
         const size_t UpdateStateBufferSize>
bool StateMachine<TStateId, MaxStates, TLayerId, NumLayers, ChangeStateBufferSize, UpdateStateBufferSize>::
    dismissScheduledStateChanges(const TLayerId &layerId)
{
    return _tasks.changeState.erase(layerId) > 0;
}

#endif // STATEMACHINE_H