#ifndef STATEMACHINE_POOL_HPP
#define STATEMACHINE_POOL_HPP

#include "state_machine.hpp"

template<typename StateOwnerType>
class StateMachinePool {
    public:
        void RegisterEntity(const std::string &id, StateOwnerType &entity) {
            auto machine = std::make_unique<StateMachine<StateOwnerType>>(entity, factory);
            machines[id] = std::move(machine);
        }

        void UnregisterEntity(const std::string &id) {
            machines.erase(id);
        }

        void UpdateAll() {
            for (const auto &[id, machine] : machines)
                machine->Update();
        }

        void ChangeState(const std::string &id, const std::string &stateId) {
            auto it = machines.find(id);

            if (it != machines.end())
                it->second->ChangeState(stateId);
        }

        void RevertState(const std::string &id) {
            auto it = machines.find(id);

            if (it != machines.end())
                it->second->RevertState();
        }

        void RegisterState(const std::string &id, std::shared_ptr<State<StateOwnerType>> state) {
            factory.RegisterState(id, state);
        }

        StateMachine<StateOwnerType> *GetMachine(const std::string &id) {
            auto it = machines.find(id);

            return it != machines.end() ? it->second.get() : nullptr;
        }

        json Serialize() const {
            json j;

            for (const auto &[id, machine] : machines) {
                j["machines"][id] = {
                    {"current", machine->GetCurrentStateId()},
                    {"history", machine->GetHistory()}
                };
            }

            j["settings"] = {
                {"auto_transitions", auto_transition ? "true" : "false"},
                {"strict_transitions", strict_transition ? "true" : "false"}
            };

            return j;
        }

        void Deserialize(const json &j, std::function<StateOwnerType &(const std::string&)> getEntity) {
            SetStrictTransitions(j.value("strict_transitions", false));
            SetAutoTransitions(j.value("auto_transitions", false));

            const json &machines_json = j["machines"];

            for (auto it = machines_json.begin(); it != machines_json.end(); ++it) {
                std::string id = it.key();
                const json &entry = it.value();

                std::string current = entry.value("current", "null");
                std::vector<std::string> history = entry.value("history", std::vector<std::string>{});

                StateOwnerType& entity = getEntity(id);
                RegisterEntity(id, entity);

                auto *machine = GetMachine(id);

                if (machine) {
                    machine->SetHistory(history);
                    machine->SetState(current);
                }
            }
        }

        void SetStrictTransitions(bool enabled) {
            strict_transition = enabled;

            for (auto &[_, machine] : machines)
                machine->SetStrictTransitions(enabled);
        }

        void SetAutoTransitions(bool enabled) {
            auto_transition = enabled;

            for (const auto &[_, machine] : machines)
                machine->SetAutoTransitions(enabled);
        }

    private:
        StateFactory<StateOwnerType> factory;
        std::unordered_map<std::string, std::unique_ptr<StateMachine<StateOwnerType>>> machines;
        bool auto_transition = false;
        bool strict_transition = false;
};

#endif // STATEMACHINE_POOL_HPP