diff --git a/library.json b/library.json
index 04fd2d0..8e8ed86 100644
--- a/library.json
+++ b/library.json
@@ -1,7 +1,7 @@
 {
   "$schema": "https://raw.githubusercontent.com/platformio/platformio-core/develop/platformio/assets/schema/library.json",
   "name": "ha-mqtt",
-  "version": "1.3.1",
+  "version": "1.4.0",
   "description": "Home Assistant classes for integration with MQTT auto discovery",
   "repository": {
     "type": "git",
diff --git a/src/esp.h b/src/esp.h
index 0fe1095..5cfb8a7 100644
--- a/src/esp.h
+++ b/src/esp.h
@@ -12,9 +12,16 @@ namespace HaESP {
     } activeSensors;
 
     Task tHeap(5 * TASK_MINUTE, TASK_FOREVER, []() {
-        Sensor::mapSensors["heap_fragmentation"]->updateState(to_string(ESP.getHeapFragmentation()).c_str());
-        Sensor::mapSensors["heap_free"]->updateState(to_string(ESP.getFreeHeap()).c_str());
-        Sensor::mapSensors["heap_max_free_block"]->updateState(to_string(ESP.getMaxFreeBlockSize()).c_str());
+        uint32_t hfree;
+        uint32_t hmax;
+        uint8_t hfrag;
+        ESP.getHeapStats(&hfree, &hmax, &hfrag);
+
+        char value[128];
+        snprintf_P(value, sizeof(value),  PSTR("{\"fragmentation\":%d,\"heap\":{\"Heap free\":\"%d B\",\"Heap max free block\":\"%d B\"}}"), hfrag, hfree, hmax);
+        Sensor::mapSensors["heap_fragmentation"]->updateState(value);
+        Sensor::mapSensors["heap_free"]->updateState(to_string(hfree).c_str());
+        Sensor::mapSensors["heap_max_free_block"]->updateState(to_string(hmax).c_str());
         }, &ts);
 
     Task tRestartInfo(TASK_IMMEDIATE, TASK_ONCE, []() {
@@ -23,9 +30,12 @@ namespace HaESP {
 
     template <class T>
     Builder<T>& heapStats(Builder<T>& builder) {
-        builder.addDiagnostic(Builder<Sensor>::instance(new Sensor{ "Heap fragmentation", "heap_fragmentation" })
+        auto heap = new Sensor{ "Heap fragmentation", "heap_fragmentation" };
+        builder.addDiagnostic(Builder<Sensor>::instance(heap)
             .withUnitMeasure("%")
             .withPrecision(0)
+            .withValueTemplate("{{ value_json.fragmentation }}")
+            .withJsonAttributes("{{ value_json.heap | tojson }}")
             .build());
         builder.addDiagnostic(Builder<Sensor>::instance(new Sensor{ "Heap free", "heap_free" })
             .withDeviceClass("data_size")
diff --git a/src/ha.h b/src/ha.h
index 7103b45..111ebc1 100644
--- a/src/ha.h
+++ b/src/ha.h
@@ -7,14 +7,17 @@ using namespace std;
 
 #define JSON_SIZE   512
 #define TOPIC_SIZE  255
+#define BASE_TOPIC  "homeassistant/%s/%s/%s"
 
 namespace Ha {
     uint16_t(*publisher)(const char* topic, const char* message);
     typedef void (*onMessage)(const char* msg);
 
-    struct Component;
-    struct DeviceConfig {
-        friend struct Component;
+    struct Config {
+        virtual void buildConfig(JsonDocument& jsonDoc) = 0;
+    };
+
+    struct DeviceConfig : Config {
         const char* id = nullptr;
         const char* name = nullptr;
         const char* model = nullptr;
@@ -51,10 +54,7 @@ namespace Ha {
             return *this;
         }
 
-    protected:
-        DeviceConfig(const char* id) : id(id) {}
-
-        void buildConfig(JsonDocument& jsonDoc) {
+        void buildConfig(JsonDocument& jsonDoc) override {
             JsonObject device = jsonDoc.createNestedObject("device");
             if (name) device["name"] = name;
             if (model) device["model"] = model;
@@ -64,40 +64,24 @@ namespace Ha {
             JsonArray identifiers = device.createNestedArray("identifiers");
             identifiers.add(id);
         }
+
+    private:
+        DeviceConfig(const char* id) : id(id) {}
     };
 
-    struct Component {
-        
-        template <typename V>
-        struct JsonPairs {
-            typedef pair<const char*, V> JsonPair;
-            unordered_map<const char*, V> jsonPairs;
+    struct Component : Config {
 
-            void add(JsonPair pair) {
-                jsonPairs.insert(pair);
-            }
-
-            void addToJson(JsonDocument& jsonDoc) {
-                for (JsonPair el : jsonPairs) {
-                    jsonDoc[el.first] = el.second;
-                }
-            }
-        };
+        const char* id = nullptr;
+        const char* type = nullptr;
 
+        const char* name = nullptr;
         const char* entityCategory = nullptr;
         const char* deviceClass = nullptr;
-        const char* name = nullptr;
-        char* id = nullptr;
-        const char* type = nullptr;
-        char configTopic[TOPIC_SIZE];
-        JsonPairs<bool> jsonBooleans;
-        JsonPairs<float> jsonNumbers;
-        JsonPairs<const char*> jsonStrings;
         DeviceConfig* mainDevice = nullptr;
         static List<Component> components;
         bool multiValueComponent = false;
 
-        Component(const char* name, const char* id, const char* type) : name(name), id((char*)id), type(type) {
+        Component(const char* id, const char* name, const char* type) : id(id), type(type), name(name) {
             components.add(this);
         }
 
@@ -116,31 +100,34 @@ namespace Ha {
         }
 
     protected:
-        virtual void buildUniqueId(char* uniqueId) {
-            sprintf(uniqueId, "%s_%s", MAIN_DEVICE_ID, id);
-        }
-
-        virtual void buildConfigTopic() {
-            sprintf(configTopic, "homeassistant/%s/%s/%s/config", type, MAIN_DEVICE_ID, id);
-        }
-
-        virtual void buildComponentConfig(JsonDocument& jsonDoc) = 0;
-
-        void buildConfig(JsonDocument& jsonDoc) {
+        void buildConfig(JsonDocument& jsonDoc) override {
             if (mainDevice) mainDevice->buildConfig(jsonDoc);
             if (entityCategory) jsonDoc["entity_category"] = entityCategory;
             if (deviceClass) jsonDoc["device_class"] = deviceClass;
             jsonDoc["name"] = name;
-            char uniqueId[50];
-            buildUniqueId(uniqueId);
-            jsonDoc["unique_id"] = uniqueId;
+           
+            buildUniqueId(jsonDoc);
             buildConfigTopic();
+        }
+    private:
+        char configTopic[TOPIC_SIZE];
 
-            buildComponentConfig(jsonDoc);
+        void buildUniqueId(JsonDocument& jsonDoc) {
+            char uniqueId[50];
+            if (multiValueComponent && deviceClass) {
+                snprintf(uniqueId, sizeof(uniqueId), "%s_%s_%s", MAIN_DEVICE_ID, deviceClass, id);
+            } else {
+                snprintf(uniqueId, sizeof(uniqueId), "%s_%s", MAIN_DEVICE_ID, id);
+            }
+            jsonDoc["unique_id"] = uniqueId;
+        }
 
-            jsonBooleans.addToJson(jsonDoc);
-            jsonNumbers.addToJson(jsonDoc);
-            jsonStrings.addToJson(jsonDoc);
+        void buildConfigTopic() {
+            if (multiValueComponent && deviceClass) {
+                snprintf(configTopic, sizeof(configTopic), BASE_TOPIC"_%s""/config", type, MAIN_DEVICE_ID, deviceClass, id);
+            } else {
+                snprintf(configTopic, sizeof(configTopic), BASE_TOPIC"/config", type, MAIN_DEVICE_ID, id);
+            }
         }
     };
 
@@ -200,19 +187,23 @@ namespace Ha {
             return *this;
         }
 
-        Builder& overrideConfig(const char* key, bool value) {
-            cmp->jsonBooleans.add({key, value});
+        Builder& withMin(unsigned int value) {
+            cmp->min = value;
             return *this;
         }
 
-        template <typename V>
-        Builder& overrideConfig(const char* key, V value) {
-            cmp->jsonNumbers.add({key, value});
+        Builder& withMax(unsigned int value) {
+            cmp->max = value;
             return *this;
         }
 
-        Builder& overrideConfig(const char* key, const char* value) {
-            cmp->jsonStrings.add({key, value});
+        Builder& withStep(unsigned int value) {
+            cmp->step = value;
+            return *this;
+        }
+
+        Builder& withJsonAttributes(const char* value) {
+            cmp->jsonAttributesTemplate = value;
             return *this;
         }
 
@@ -257,13 +248,13 @@ namespace Ha {
         }
     };
 
-    struct Command : Component {
+    struct Command : Config {
         bool retain = false;
-        static unordered_map<string, Command*> mapCommands;
+        static unordered_map<string, Command*> mapCommandTopics;
 
-        Command(const char* name, const char* id, const char* type, onMessage f) : Component(name, id, type), f(f) {
-            sprintf(commandTopic, "homeassistant/%s/%s/%s/set", type, MAIN_DEVICE_ID, id);
-            mapCommands.insert({ string(commandTopic), this });
+        Command(Component* cmp, onMessage f) : f(f), cmp(cmp) {
+            snprintf(commandTopic, sizeof(commandTopic),  BASE_TOPIC"/set", cmp->type, MAIN_DEVICE_ID, cmp->id);
+            mapCommandTopics.insert({ string(commandTopic), this });
         }
 
         virtual void onCommand(const char* msg) {
@@ -273,31 +264,33 @@ namespace Ha {
     protected:
         char commandTopic[TOPIC_SIZE];
         onMessage f;
-        virtual void buildCommandConfig(JsonDocument& jsonDoc) = 0;
 
-        void buildComponentConfig(JsonDocument& jsonDoc) override {
-            jsonDoc["command_topic"] = commandTopic;
+        void buildConfig(JsonDocument& jsonDoc) override {
+            jsonDoc["command_topic"] = (const char*)commandTopic;
             if (retain) jsonDoc["retain"] = true;
-            buildCommandConfig(jsonDoc);
+        }
+    private:
+        Component* cmp;
+    };
+
+    struct Button : Component, Command {
+
+        Button(const char* name, const char* id, onMessage f = nullptr) : Component(id, name, "button"), Command(this, f) {}
+        
+        void buildConfig(JsonDocument& jsonDoc) override {
+            Component::buildConfig(jsonDoc);
+            Command::buildConfig(jsonDoc);
         }
     };
 
-    struct Button : Command {
+    struct State : Config {
+        char stateTopic[TOPIC_SIZE];
+        const char* jsonAttributesTemplate = nullptr;
 
-        Button(const char* name, const char* id, onMessage f = nullptr) : Command(name, id, "button", f) {}
-
-    protected:
-        void buildCommandConfig(JsonDocument& jsonDoc) override {}
-
-    };
-
-    struct StateConfig {
-        static unordered_map<string, Command*> mapStateTopics;
-
-        StateConfig(Component* cmp) : cmp(cmp) {}
+        State(Component* cmp) : cmp(cmp) {}
 
         void withStateTopic() {
-            sprintf(stateTopic, "homeassistant/%s/%s/%s/state", cmp->type, MAIN_DEVICE_ID, cmp->id);
+            snprintf(stateTopic, sizeof(stateTopic),  BASE_TOPIC"/state", cmp->type, MAIN_DEVICE_ID, cmp->id);
         }
 
         void updateState(const char* message) {
@@ -305,86 +298,81 @@ namespace Ha {
         }
 
     protected:
-        char stateTopic[TOPIC_SIZE];
+        void buildConfig(JsonDocument& jsonDoc) override {
+            if (stateTopic[0]) {
+                jsonDoc["state_topic"] = (const char*)stateTopic;
+                if (jsonAttributesTemplate) {
+                    jsonDoc["json_attributes_template"] = jsonAttributesTemplate;
+                    jsonDoc["json_attributes_topic"] = (const char*)stateTopic;
+                }
+            }
+        }
+    private:
         Component* cmp;
-        void buildConfig(JsonDocument& jsonDoc) {
-            if (stateTopic[0]) jsonDoc["state_topic"] = stateTopic;
+    };
+
+    struct StatefulCommand : Command, State {
+        static unordered_map<string, Command*> mapRestoreStateTopics;
+
+        StatefulCommand(Component* cmp, onMessage f) : Command(cmp, f), State(cmp) {}
+
+        void restoreFromState() {
+            mapRestoreStateTopics.insert({stateTopic, this});
+        }
+
+    protected:
+        void buildConfig(JsonDocument& jsonDoc) override {
+            Command::buildConfig(jsonDoc);
+            State::buildConfig(jsonDoc);
         }
     };
 
-    struct Switch : Command, StateConfig {
+    struct Switch : Component, StatefulCommand {
 
-        Switch(const char* name, const char* id, onMessage f = nullptr) : Command(name, id, "switch", f), StateConfig(this) {}
+        Switch(const char* name, const char* id, onMessage f = nullptr) : Component(id, name, "switch"), StatefulCommand(this, f) {}
 
         void updateState(bool state) {
-            StateConfig::updateState(state ? "ON" : "OFF");
+            State::updateState(state ? "ON" : "OFF");
         }
 
-        void restoreFromState() {
-            mapStateTopics.insert({stateTopic, this});
-        }
-
-    protected:
-        void buildCommandConfig(JsonDocument& jsonDoc) override {
-            StateConfig::buildConfig(jsonDoc);
+        void buildConfig(JsonDocument& jsonDoc) override {
+            Component::buildConfig(jsonDoc);
+            StatefulCommand::buildConfig(jsonDoc);
         }
     };
 
-    struct Number : Command, StateConfig {
+    struct Number : Component, StatefulCommand {
         unsigned int min, max, step;
 
-        Number(const char* name, const char* id, unsigned int min, unsigned int max, unsigned int step, onMessage f)
-            : Command(name, id, "number", f), StateConfig(this), min(min), max(max), step(step) {}
+        Number(const char* name, const char* id, onMessage f) : Component(id, name, "number"), StatefulCommand(this, f) {}
 
         void updateState(unsigned int value) {
-            char message[32];
-            sprintf(message, "%u", value);
-            StateConfig::updateState(message);
+            State::updateState(to_string(value).c_str());
         }
 
-        void restoreFromState() {
-            mapStateTopics.insert({stateTopic, this});
-        }
-
-    protected:
-        void buildCommandConfig(JsonDocument& jsonDoc) override {
-            StateConfig::buildConfig(jsonDoc);
+        void buildConfig(JsonDocument& jsonDoc) override {
+            Component::buildConfig(jsonDoc);
+            StatefulCommand::buildConfig(jsonDoc);
             jsonDoc["min"] = min;
             jsonDoc["max"] = max;
             jsonDoc["step"] = step;
         }
     };
 
-    struct Sensor : Component, StateConfig {
+    struct Sensor : Component, State {
         const char* unitMeasure = nullptr;
         const char* valueTemplate = nullptr;
         unsigned int precision = 2;
         static unordered_map<string, Sensor*> mapSensors;
 
-        Sensor(const char* name, const char* id) : Component(name, id, "sensor"), StateConfig(this) {
+        Sensor(const char* name, const char* id) : Component(id, name, "sensor"), State(this) {
             withStateTopic();
             mapSensors.insert({ string(id), this });
         }
 
-    protected:
-        void buildUniqueId(char* uniqueId) override {
-            if (multiValueComponent && deviceClass) {
-                sprintf(uniqueId, "%s_%s_%s", MAIN_DEVICE_ID, deviceClass, id);
-            } else {
-                Component::buildUniqueId(uniqueId);
-            }
-        }
-
-        void buildConfigTopic() override {
-            if (multiValueComponent && deviceClass) {
-                sprintf(configTopic, "homeassistant/%s/%s/%s_%s/config", type, MAIN_DEVICE_ID, deviceClass, id);
-            } else {
-                Component::buildConfigTopic();
-            }
-        }
-
-        void buildComponentConfig(JsonDocument& jsonDoc) override {
-            StateConfig::buildConfig(jsonDoc);
+        void buildConfig(JsonDocument& jsonDoc) override {
+            Component::buildConfig(jsonDoc);
+            State::buildConfig(jsonDoc);
             if (unitMeasure) jsonDoc["unit_of_measurement"] = unitMeasure;
             if (valueTemplate) jsonDoc["value_template"] = valueTemplate;
             if (isNumericSensor()) jsonDoc["suggested_display_precision"] = precision;
@@ -435,7 +423,7 @@ namespace Ha {
 
     List<Component> Component::components;
     List<AbstractBuilder> AbstractBuilder::builders;
-    unordered_map<string, Command*> Command::mapCommands;
+    unordered_map<string, Command*> Command::mapCommandTopics;
     unordered_map<string, Sensor*> Sensor::mapSensors;
-    unordered_map<string, Command*> StateConfig::mapStateTopics;
+    unordered_map<string, Command*> StatefulCommand::mapRestoreStateTopics;
 }
diff --git a/src/mqtt.h b/src/mqtt.h
index 9f2be03..796aa78 100644
--- a/src/mqtt.h
+++ b/src/mqtt.h
@@ -47,13 +47,13 @@ namespace Mqtt {
         memcpy(msg, payload, len);
         msg[len] = 0;
         auto strTopic = string{ topic };
-        auto cmd = Command::mapCommands[strTopic];
+        auto cmd = Command::mapCommandTopics[strTopic];
         if (cmd) cmd->onCommand(msg);
 
-        cmd = StateConfig::mapStateTopics[strTopic];
+        cmd = StatefulCommand::mapRestoreStateTopics[strTopic];
         if (cmd) {
             cmd->onCommand(msg);
-            StateConfig::mapStateTopics.erase(strTopic);
+            StatefulCommand::mapRestoreStateTopics.erase(strTopic);
         }
     }