#pragma once

#include <ESP8266WiFi.h>
#include "TaskScheduler.h"
#include "ha.h"

extern Scheduler ts;

using namespace Ha;

namespace HaESP {
    struct {
        bool heapStats = false;
        bool restartInfo = false;
        bool wifiInfo = false;
    } activeSensors;

    Task tHeap(5 * TASK_MINUTE, TASK_FOREVER, []() {
        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);
        GenericSensor::mapSensors["heap_fragmentation"]->updateState(value);
        GenericSensor::mapSensors["heap_free"]->updateState(to_string(hfree).c_str());
        GenericSensor::mapSensors["heap_max_free_block"]->updateState(to_string(hmax).c_str());
        }, &ts);

    Task tRestartInfo(TASK_IMMEDIATE, TASK_ONCE, []() {
        GenericSensor::mapSensors["restart_reason"]->updateState(ESP.getResetReason().c_str());
        }, &ts);

    Task tRestart(1 * TASK_SECOND, TASK_ONCE, []() {
        ESP.restart();
        }, &ts);

    Task tWifi(3 * TASK_MINUTE, TASK_FOREVER, []() {
        GenericSensor::mapSensors["wifi_signal_strength"]->updateState(to_string(WiFi.RSSI()).c_str());
        }, &ts);

    template <class T>
    Builder<T>& heapStats(Builder<T>& builder) {
        builder.addDiagnostic(Builder<Sensor>(new Sensor{ "Heap fragmentation", "heap_fragmentation" })
            .withUnitMeasure("%")
            .withPrecision(0)
            .withValueTemplate("{{ value_json.fragmentation }}")
            .withJsonAttributes("{{ value_json.heap | tojson }}")
            .build());
        builder.addDiagnostic(Builder<Sensor>(new Sensor{ "Heap free", "heap_free" })
            .withDeviceClass("data_size")
            .withUnitMeasure("B")
            .withPrecision(0)
            .build()
        );
        builder.addDiagnostic(Builder<Sensor>(new Sensor{ "Heap max free block", "heap_max_free_block" })
            .withDeviceClass("data_size")
            .withUnitMeasure("B")
            .withPrecision(0)
            .build()
        );
        activeSensors.heapStats = true;
        return builder;
    }

    template <class T>
    Builder<T>& restartInfo(Builder<T>& builder) {
        builder.addDiagnostic((new Sensor{ "Restart reason", "restart_reason" }));
        activeSensors.restartInfo = true;
        return builder;
    }

    template <class T>
    Builder<T>& wifiInfo(Builder<T>& builder) {
        builder.addDiagnostic(Builder<Sensor>((
            new Sensor{ "WiFi Signal (RSSI)", "wifi_signal_strength" }))
            .withDeviceClass("signal_strength")
            .withUnitMeasure("dBm")
            .withPrecision(0)
            .build()
        );
        activeSensors.wifiInfo = true;
        return builder;
    }

    Builder<Button>& restartButton() {
        return (new Builder<Button>(new Button{"Restart", "restart",
            [](const char* msg) {
                if (strcmp("PRESS", msg) == 0) tRestart.restartDelayed();
            }
        }))
        ->withIcon("mdi:restart");
    }

    void enableSensors() {
        if (activeSensors.heapStats) tHeap.enable();
        if (activeSensors.restartInfo) tRestartInfo.enable();
        if (activeSensors.wifiInfo) tWifi.enable();
    }
}