// Arduino & ESP headers
#include "Dtypes.h"
#include "SwimTrackerConfig.h"

#include <ArduinoJson.h>
#include <HTTPClient.h>

#include "esp_https_ota.h"
#include "esp_ota_ops.h"
#include "ESPmDNS.h"

// Own libs
#include "WifiManager.h"
#include "MockScale.h"
#include "Scale.h"
#include "MeasurementSession.h"
#include "SessionManager.h"
#include "SpiffsStorage.h"
#include "SimpleMeasurementSession.h"
#include "EspHttp.h"
#include "WebDAV.h"
#include "WebsocketServer.h"
#include "UserDB.h"

using Session_T = SimpleMeasurementSession<MeasurementT, CONFIG_SESSION_MAX_SIZE>;
SessionManager<Session_T> sessionManager;

UserStorage userStorage;

EspHttp espHttpServer;
WebsocketServer<Session_T> webSocketServer(sessionManager, userStorage, 81);

WifiManager wifiManager;

extern const uint8_t certificate_pem[] asm("_binary_certificate_pem_start");

bool firmwareUpdate()
{
    esp_http_client_config_t config;
    Serial.println((char *)certificate_pem);
    memset(&config, 0, sizeof(esp_http_client_config_t));
    config.url = UPDATE_URL;
    config.cert_pem = (char *)certificate_pem;
    Serial.println("Starting firmware upgrade");
    esp_err_t ret = esp_https_ota(&config);
    if (ret == ESP_OK)
    {
        Serial.println("Firmware upgrade successful - restarting");
        esp_restart();
    }
    else
    {
        Serial.println("Firmware upgrade failed");
        return false;
    }
    return true;
}

void sessionManagerSetup()
{
    Preferences scalePrefs;
    scalePrefs.begin("st_prefs");
    int valueRightShift = scalePrefs.getInt("valueRightShift", CONFIG_VALUE_RIGHT_SHIFT);
    uint8_t tareAvgCount = scalePrefs.getUInt("tareAvgCount", CONFIG_TARE_AVG_COUNT);
    MeasurementT autoStartMinThreshold = scalePrefs.getUInt("aStartMinTh", CONFIG_AUTO_START_MIN_THRESHOLD);
    MeasurementT autoStartMaxThreshold = scalePrefs.getUInt("aStartMaxTh", CONFIG_AUTO_START_MAX_THRESHOLD);
    uint32_t autoStartMaxMeasurementsBetweenPeaks = scalePrefs.getUInt("aStartCount", CONFIG_AUTO_START_MAX_MEASUREMENTS_BETWEEN_PEAKS);
    MeasurementT autoStopThreshold = scalePrefs.getUInt("aStopTh", CONFIG_AUTO_STOP_THRESHOLD);
    uint32_t autoStopNumMeasurements = scalePrefs.getUInt("aStopCount", CONFIG_AUTO_STOP_NUM_MEASUREMENTS);

    bool autoStartEnabled = scalePrefs.getBool("aStartEnabled", true);
    bool autoStopEnabled = scalePrefs.getBool("aStopEnabled", true);
    if (wifiManager.inProvisioningMode())
        autoStartEnabled = false;

    // Session
    sessionManager.begin(CONFIG_SCALE_DOUT_PIN, CONFIG_SCALE_SCK_PIN,
                         tareAvgCount, valueRightShift,
                         autoStartMinThreshold, autoStartMaxThreshold, autoStartMaxMeasurementsBetweenPeaks,
                         autoStopThreshold, autoStopNumMeasurements);
    sessionManager.enableAutoStart(autoStartEnabled);
    sessionManager.enableAutoStop(autoStopEnabled);
}

template <typename SessionT>
void httpSetup(SessionManager<SessionT> *sessionManager, WifiManager *wifiManager)
{
    auto cbStartSession = [sessionManager](httpd_req_t *req) {
        httpd_resp_set_hdr(req, "Access-Control-Allow-Origin", "*");
        httpd_resp_send(req, "Session started", -1);
        sessionManager->startMeasurements();
        Serial.println("Started session");
    };
    auto cbStopSession = [sessionManager](httpd_req_t *req) {
        httpd_resp_set_hdr(req, "Access-Control-Allow-Origin", "*");
        httpd_resp_send(req, "Session stopped", -1);
        sessionManager->stopMeasurements();
        Serial.println("Stopped session");
    };
    auto cbRestart = [](httpd_req_t *req) {
        Serial.println("Restarted requested");
        ESP.restart();
    };
    auto cbTare = [sessionManager](httpd_req_t *req) {
        Serial.println("Tare");
        sessionManager->tare();
    };
    auto cbFirmwareUpdate = [](httpd_req_t *req) {
        httpd_resp_set_hdr(req, "Access-Control-Allow-Origin", "*");
        httpd_resp_send(req, "OK", -1);
        firmwareUpdate();
    };
    auto cbStatus = [sessionManager](httpd_req_t *req) {
        httpd_resp_set_hdr(req, "Access-Control-Allow-Origin", "*");
        httpd_resp_set_hdr(req, "Content-Type", "application/json");

        StaticJsonDocument<1024> json;
        // session
        {
            JsonObject sessionObj = json.createNestedObject("session");
            if (sessionManager->isMeasuring())
            {
                const auto &session = sessionManager->session();
                sessionObj["started"] = session.getStartTime();
                sessionObj["num_measurements"] = session.numMeasurements();
            }
            sessionObj["auto_start"] = sessionManager->autoStartEnabled();
            sessionObj["auto_stop"] = sessionManager->autoStopEnabled();
        }
        // scale
        {
            JsonObject scaleObj = json.createNestedObject("scale");
            scaleObj["tare_offset"] = sessionManager->tareOffset();
            scaleObj["value_right_shift"] = sessionManager->valueRightShift();
        }
        // flash
        {
            JsonObject fsObj = json.createNestedObject("flash");
            fsObj["used"] = portablefs::usedBytes();
            fsObj["free"] = portablefs::totalBytes() - portablefs::usedBytes();
        }
        // RAM
        {
            JsonObject ramObj = json.createNestedObject("ram");
            ramObj["used"] = ESP.getHeapSize() - ESP.getFreeHeap();
            ramObj["free"] = ESP.getFreeHeap();
        }
        // PSRAM
        {
            JsonObject psramObj = json.createNestedObject("psram");
            psramObj["used"] = ESP.getPsramSize() - ESP.getFreePsram();
            psramObj["free"] = ESP.getFreePsram();
        }
        // firmware
        {
            auto descr = esp_ota_get_app_description();

            JsonObject firmware = json.createNestedObject("firmware");
            firmware["name"] = descr->project_name;
            firmware["version"] = descr->version;
            firmware["idf_version"] = descr->idf_ver;
            firmware["compile_date"] = descr->date;
            firmware["compile_time"] = descr->time;
        }
        char jsonText[512];
        auto bytesWritten = serializeJson(json, jsonText);
        httpd_resp_send(req, jsonText, bytesWritten);
    };
    auto cbGetData = [sessionManager](httpd_req_t *req) {
        auto sessionId = sessionManager->session().getStartTime();
        uint32_t startIdx = getUrlQueryParameter(req, "startIdx", 0);
        //Serial.printf("Data request, start index: %d\n", startIdx);
        if (startIdx >= sessionManager->session().numMeasurements())
        {
            httpd_resp_send_404(req);
            return;
        }

        //headers
        httpd_resp_set_hdr(req, "Content-Type", "application/x-msgpack");
        httpd_resp_set_hdr(req, "Access-Control-Allow-Origin", "*");
        String contentDisp = "attachment; filename=\"" + String(sessionId) + ".st\"";
        httpd_resp_set_hdr(req, "content-disposition", contentDisp.c_str());

        //data
        StreamingMsgPackEncoder<DummyWriter> encoderToDetermineSize(nullptr);
        encoderToDetermineSize.setSizeCountMode(true);
        sessionManager->session().serialize(encoderToDetermineSize, startIdx);
        auto totalSize = encoderToDetermineSize.getContentLength();

        char *buf = (char *)malloc(totalSize);
        CopyWriter copyWriter((uint8_t *)buf);
        StreamingMsgPackEncoder<CopyWriter> encoder(&copyWriter);
        sessionManager->session().serialize(encoder, startIdx);
        httpd_resp_send(req, buf, totalSize);
        free(buf);
    };
    auto cbWifiGet = [wifiManager](httpd_req_t *req) {
        httpd_resp_set_hdr(req, "Access-Control-Allow-Origin", "*");
        StaticJsonDocument<128> json;
        json["state"] = wifiManager->stateStr();
        char jsonText[128];
        auto bytesWritten = serializeJson(json, jsonText);
        httpd_resp_send(req, jsonText, bytesWritten);
    };
    auto cbWifiPost = [wifiManager](httpd_req_t *req) {
        httpd_resp_set_hdr(req, "Access-Control-Allow-Origin", "*");
        StaticJsonDocument<1024> json;
        char content[512];
        size_t recvSize = min(req->content_len, sizeof(content)); //truncate if too long
        int ret = httpd_req_recv(req, content, recvSize);
        if (ret <= 0)
        {
            if (ret == HTTPD_SOCK_ERR_TIMEOUT)
                httpd_resp_send_408(req);
            return;
        }
        DeserializationError err = deserializeJson(json, content);
        if (err)
        {
            httpd_resp_set_status(req, "400 Bad Request");
            httpd_resp_send(req, "JSON parse error", -1);
            return;
        }
        if (json.containsKey("reset_to_provisioning") && json["reset_to_provisioning"].as<bool>())
        {
            wifiManager->resetToApProvisioning();
            httpd_resp_send(req, "OK", -1);
            ESP.restart();
            return;
        }
        if (json.containsKey("ap_password"))
        {
            httpd_resp_send(req, "OK", -1);
            wifiManager->setApCredentials(json["ap_password"].as<const char *>());
            ESP.restart();
            return;
        }
        if (json.containsKey("sta_ssid") && json.containsKey("sta_password"))
        {
            wifiManager->setStaCredentials(json["sta_ssid"].as<const char *>(), //
                                           json["sta_password"].as<const char *>());
            httpd_resp_send(req, "OK", -1);
            ESP.restart();
            return;
        }
        httpd_resp_set_status(req, "400 Bad Request");
        httpd_resp_send(req, "Invalid keys in JSON", -1);
    };
    auto cbSettingsGet = [](httpd_req_t *req) {
        httpd_resp_set_hdr(req, "Access-Control-Allow-Origin", "*");
        httpd_resp_set_hdr(req, "Content-Type", "application/json");

        StaticJsonDocument<1024> json;
        Preferences prefs;
        prefs.begin("st_prefs");

        json["valueRightShift"] = prefs.getInt("valueRightShift", CONFIG_VALUE_RIGHT_SHIFT);
        json["tareAvgCount"] = prefs.getUInt("tareAvgCount", CONFIG_TARE_AVG_COUNT);
        json["autoStartMinThreshold"] = prefs.getUInt("aStartMinTh", CONFIG_AUTO_START_MIN_THRESHOLD);
        json["autoStartMaxThreshold"] = prefs.getUInt("aStartMaxTh", CONFIG_AUTO_START_MAX_THRESHOLD);
        json["autoStartMaxMeasurementsBetweenPeaks"] = prefs.getUInt("aStartCount", CONFIG_AUTO_START_MAX_MEASUREMENTS_BETWEEN_PEAKS);
        json["autoStopThreshold"] = prefs.getUInt("aStopTh", CONFIG_AUTO_STOP_THRESHOLD);
        json["autoStopNumMeasurements"] = prefs.getUInt("aStopCount", CONFIG_AUTO_STOP_NUM_MEASUREMENTS);
        json["autoStartEnabled"] = prefs.getBool("aStartEnabled", true);
        json["autoStopEnabled"] = prefs.getBool("aStopEnabled", true);
        char jsonText[1024];
        auto bytesWritten = serializeJson(json, jsonText);
        httpd_resp_send(req, jsonText, bytesWritten);
    };
    auto cbSettingsPost = [](httpd_req_t *req) {
        httpd_resp_set_hdr(req, "Access-Control-Allow-Origin", "*");
        StaticJsonDocument<1024> json;
        char content[512];
        size_t recvSize = min(req->content_len, sizeof(content)); //truncate if too long
        int ret = httpd_req_recv(req, content, recvSize);
        if (ret <= 0)
        {
            if (ret == HTTPD_SOCK_ERR_TIMEOUT)
                httpd_resp_send_408(req);
            return;
        }
        DeserializationError err = deserializeJson(json, content);
        if (err)
        {
            httpd_resp_set_status(req, "400 Bad Request");
            httpd_resp_send(req, "JSON parse error", -1);
            return;
        }

        Preferences prefs;
        prefs.begin("st_prefs");
        if (json.containsKey("valueRightShift"))
            prefs.putInt("valueRightShift", json["valueRightShift"].as<int>());
        if (json.containsKey("tareAvgCount"))
            prefs.putUInt("tareAvgCount", json["tareAvgCount"].as<unsigned int>());
        if (json.containsKey("autoStartMinThreshold"))
            prefs.putUInt("aStartMinTh", json["autoStartMinThreshold"].as<unsigned int>());
        if (json.containsKey("autoStartMaxThreshold"))
            prefs.putUInt("aStartMaxTh", json["autoStartMaxThreshold"].as<unsigned int>());
        if (json.containsKey("autoStartMaxMeasurementsBetweenPeaks"))
            prefs.putUInt("aStartCount", json["autoStartMaxMeasurementsBetweenPeaks"].as<unsigned int>());
        if (json.containsKey("autoStopThreshold"))
            prefs.putUInt("aStopTh", json["autoStopThreshold"].as<unsigned int>());
        if (json.containsKey("autoStopNumMeasurements"))
            prefs.putUInt("aStopCount", json["autoStopNumMeasurements"].as<unsigned int>());
        if (json.containsKey("autoStartEnabled"))
            prefs.putBool("aStartEnabled", json["autoStartEnabled"].as<bool>());
        if (json.containsKey("autoStopEnabled"))
            prefs.putBool("aStopEnabled", json["autoStopEnabled"].as<bool>());

        sessionManagerSetup();
        httpd_resp_send(req, "OK", -1);
    };
    auto cbSettingsDelete = [](httpd_req_t *req) {
        httpd_resp_set_hdr(req, "Access-Control-Allow-Origin", "*");
        Preferences prefs;
        prefs.begin("st_prefs");
        prefs.putInt("valueRightShift", CONFIG_VALUE_RIGHT_SHIFT);
        prefs.putUInt("tareAvgCount", CONFIG_TARE_AVG_COUNT);
        prefs.putUInt("aStartMinTh", CONFIG_AUTO_START_MIN_THRESHOLD);
        prefs.putUInt("aStartMaxTh", CONFIG_AUTO_START_MAX_THRESHOLD);
        prefs.putUInt("aStartCount", CONFIG_AUTO_START_MAX_MEASUREMENTS_BETWEEN_PEAKS);
        prefs.putUInt("aStopTh", CONFIG_AUTO_STOP_THRESHOLD);
        prefs.putUInt("aStopCount", CONFIG_AUTO_STOP_NUM_MEASUREMENTS);
        prefs.putBool("aStartEnabled", true);
        prefs.putBool("aStopEnabled", true);
        httpd_resp_send(req, "OK", -1);
    };

    espHttpServer.start();

    espHttpServer.on("/api/wifi", HTTP_GET, cbWifiGet);
    espHttpServer.on("/api/wifi", HTTP_POST, cbWifiPost);
    espHttpServer.on("/api/restart", HTTP_GET, cbRestart);

    if (!wifiManager->inProvisioningMode())
    {
        espHttpServer.on("/api/session/start", HTTP_GET, cbStartSession);
        espHttpServer.on("/api/session/start", HTTP_POST, cbStartSession);

        espHttpServer.on("/api/session/stop", HTTP_GET, cbStopSession);
        espHttpServer.on("/api/session/stop", HTTP_POST, cbStopSession);

        espHttpServer.on("/api/session/data", HTTP_GET, cbGetData);
        espHttpServer.on("/api/status", HTTP_GET, cbStatus);
        espHttpServer.on("/api/tare", HTTP_GET, cbTare);
        espHttpServer.on("/api/firmwareupdate", HTTP_GET, cbFirmwareUpdate);

        espHttpServer.on("/api/settings", HTTP_GET, cbSettingsGet);
        espHttpServer.on("/api/settings", HTTP_POST, cbSettingsPost);
        espHttpServer.on("/api/settings", HTTP_DELETE, cbSettingsDelete);

        auto webdav = webdavHandler("/webdav/", "/dat");
        espHttpServer.on("/webdav/*?", HTTP_GET, webdav);
        espHttpServer.on("/webdav/*?", HTTP_PROPFIND, webdav);
        espHttpServer.on("/webdav/*?", HTTP_DELETE, webdav);
        espHttpServer.on("/webdav/*?", HTTP_OPTIONS, webdav);
        Serial.println("HTTP setup done");
    }
    else
    {
        Serial.println("HTTP setup with limited API in provisioning mode");
    }
}

String getIdSuffix()
{
    uint8_t baseMac[6];
    esp_read_mac(baseMac, ESP_MAC_WIFI_STA);
    char baseMacChr[18] = {0};
    sprintf(baseMacChr, "-%02X%02X%02X", baseMac[3], baseMac[4], baseMac[5]);
    return String(baseMacChr);
}

void mdnsSetup(const String &fullHostname)
{
    if (!MDNS.begin(fullHostname.c_str()))
    {
        Serial.println("Error setting up MDNS responder!");
        while (true)
        {
            delay(1000);
        }
    }
    Serial.printf("mDNS started %s\n", fullHostname.c_str());
    MDNS.addService("swimtracker", "tcp", 81);
}

void setup()
{
    // Serial
    Serial.begin(115200);
    while (!Serial)
    {
    }

    // File system
    bool spiffsResult = SPIFFS.begin(true);
    if (!spiffsResult)
        Serial.println("Failed to mount/format SPIFFS file system");

    ESP_ERROR_CHECK(esp_event_loop_create_default());

    userStorage.init();

    // WiFi
    String fullHostname = String(CONFIG_HOSTNAME) + getIdSuffix();
    wifiManager.begin(fullHostname);
    Serial.print("Connected to WiFi. IP:");
    Serial.println(WiFi.localIP());
    Serial.printf("WIFI state: %s\n", wifiManager.stateStr());

    mdnsSetup(fullHostname);

    sessionManagerSetup();
    sessionManager.tare();

    // HTTP & Websocket server
    httpSetup(&sessionManager, &wifiManager);
    if (!wifiManager.inProvisioningMode())
        webSocketServer.begin();
}

void loop()
{
    sessionManager.iteration();
    webSocketServer.iteration();
    wifiManager.wifiWatchdog();
}