#ifndef STATEMACHINE_POOL_HPP
#define STATEMACHINE_POOL_HPP

#include "state.hpp"
#include "state_machine.hpp"
#include <functional>
#include <iostream>
#include <ostream>
#include <unordered_map>

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 AddTransition(Transition<StateOwnerType> &transition) {
            std::string from = transition.from;
            std::string to = transition.target_state;
            std::function <bool (const StateOwnerType&)> condition;

            transitions[from].push_back(transition);
        }

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

        void Update(const std::string &id) {
            auto &machine = machines.find(id)->second;
            auto &owner = machine->GetOwner();
            const std::string &current = machine->GetCurrentStateId();

            auto it = transitions.find(current);
            if (it != transitions.end() && auto_transition == true)
                for (const auto &transition : it->second)
                    if (transition.condition(owner)) {
                        machine->ChangeState(transition.target_state);
                        break;
                    }
            machine->Update();
        }

        void ChangeState(const std::string &id, const std::string &stateId) {
            auto it = machines.find(id);
            const std::string current = it->second->GetCurrentStateId();
            auto &owner = it->second->GetOwner();

            bool allowed = true;
            auto transitions_it = transitions.find(current);
            if (strict_transition && (transitions_it != transitions.end())) {
                for (const auto &transition : transitions_it->second)
                    if (!transition.condition(owner))
                        allowed = false;
            }

            if (it != machines.end() && allowed) {
                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) {
            this->strict_transition = enabled;
        }

        bool GetStrictTransitions() const {
            return strict_transition;
        }

        void SetAutoTransitions(bool enabled) {
            this->auto_transition = enabled;
        }

        bool GetAutoTransitions() const {
            return auto_transition;
        }

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

#endif // STATEMACHINE_POOL_HPP