#include "update_service.h"
#include "platform/counted_client.h"

#include <SD.h>
#include <WiFi.h>
#include <esp_ota_ops.h>

#include <memory>

#include "platform/ota_device.h"
#include "platform/system_info.h"
#include "probation.h"
#include "update_parser.h"
#include "version.h"

namespace roro {

// Where an Update File comes from: the network or a file on the SD card.
class UpdateSource {
   public:
    virtual ~UpdateSource() = default;
    // Bytes read, 0 at the end of the file, -1 on an error or a stall.
    virtual int read(uint8_t* buf, size_t len) = 0;
    virtual void reply(const std::string& line) { (void)line; }
};

namespace {

constexpr uint32_t kStallMs = 10000;
constexpr uint32_t kForceRestartMs = 90000;  // the main loop waits up to 60 s for someone typing
constexpr size_t kChunk = 4096;

class NetSource : public UpdateSource {
   public:
    explicit NetSource(NetworkClient& c) : c_(c) {}
    int read(uint8_t* buf, size_t len) override {
        uint32_t since = millis();
        while (!c_.available()) {
            if (!c_.connected()) return 0;  // the sender closed its side: end of file
            if (millis() - since > kStallMs) return -1;
            delay(5);
        }
        return c_.read(buf, len);
    }
    void reply(const std::string& line) override {
        c_.print((line + "\n").c_str());
        c_.flush();
    }

   private:
    NetworkClient& c_;
};

class FileSource : public UpdateSource {
   public:
    explicit FileSource(File& f) : f_(f) {}
    int read(uint8_t* buf, size_t len) override { return static_cast<int>(f_.read(buf, len)); }

   private:
    File& f_;
};

}  // namespace

UpdateService::UpdateService(KeyValueStore& store, WifiService& wifi, SavedNetworks& saved, StorageService& storage,
                             EventBus& bus, const Settings& settings)
    : store_(store), wifi_(wifi), saved_(saved), storage_(storage), bus_(bus), settings_(settings) {
}

void UpdateService::notify(const std::string& text, NotificationLevel level) {
    bus_.publish(Event::withText(EventType::Notification, text.c_str(), static_cast<int32_t>(level)));
}

std::string UpdateService::incomingVersion() const { return incoming_; }

void UpdateService::start() {
    // Probation: new firmware boots "pending verify" until it marks itself valid.
    esp_ota_img_states_t state;
    probation_ = esp_ota_get_state_partition(esp_ota_get_running_partition(), &state) == ESP_OK &&
                 state == ESP_OTA_IMG_PENDING_VERIFY;

    // A pending update that isn't the running version means the bootloader rolled it back.
    std::string pending, from;
    store_.getString("ota_pending", pending);
    store_.getString("ota_from", from);
    if (!pending.empty() && pending != versionString()) {
        notify("Update to " + pending + " failed, back on " + versionString(), NotificationLevel::Warning);
        store_.putString("ota_pending", "");
    }

    if (!task_) xTaskCreate(taskEntry, "update", 5120, this, 1, &task_);  // peak 3.4 KB (ECDSA check, M2)
}

void UpdateService::bootGuard(KeyValueStore& store) {
    esp_ota_img_states_t state;
    bool probation = esp_ota_get_state_partition(esp_ota_get_running_partition(), &state) == ESP_OK &&
                     state == ESP_OTA_IMG_PENDING_VERIFY;
    system_info::recordSlotVersion(store, esp_ota_get_running_partition(), versionString());
    int32_t attempts = 0;
    store.getInt("ota_attempts", attempts);
    if (Probation::rollBackAtBoot(probation, attempts)) {
        store.putInt("ota_attempts", 0);
        esp_ota_mark_app_invalid_rollback_and_reboot();  // does not return when there's a previous image
    }
    store.putInt("ota_attempts", probation ? attempts + 1 : 0);
}

void UpdateService::tick(uint32_t nowMs) {
    if (!probation_) return;
    bool wifiConfigured = settings_.getBool(Setting::WifiEnabled) && saved_.count() > 0;
    bool wifiUp = wifi_.state() == WifiController::State::Connected;
    switch (Probation::judge(nowMs, firstFrame_, wifiConfigured, wifiUp)) {
        case Probation::Verdict::Wait: break;
        case Probation::Verdict::Confirm:
            esp_ota_mark_app_valid_cancel_rollback();
            probation_ = false;
            store_.putString("ota_pending", "");
            store_.putInt("ota_attempts", 0);
            notify(std::string("Updated to ") + versionString(), NotificationLevel::Info);
            break;
        case Probation::Verdict::RollBack:
            // ota_pending still names this version: the previous firmware will report the failure.
            store_.putInt("ota_attempts", 0);
            delay(200);
            esp_ota_mark_app_invalid_rollback_and_reboot();
            break;
    }
}

void UpdateService::install(UpdateSource& source, const char* via) {
    EcdsaVerifier verifier;
    EspOtaSink sink;
    UpdateParser parser(verifier, sink, sink.capacity(), versionString());
    std::unique_ptr<uint8_t[]> buf(new uint8_t[kChunk]);
    incoming_.clear();
    percent_ = 0;
    phase_ = Phase::Receiving;

    bool ok = false;
    for (;;) {
        int n = source.read(buf.get(), kChunk);
        if (n < 0) break;  // stalled or broken: aborted below
        if (n == 0) {
            ok = parser.end();
            break;
        }
        parser.feed(buf.get(), n);
        if (incoming_.empty() && parser.state() == UpdateParser::State::Image) incoming_ = parser.version();
        percent_ = parser.percent();
        if (parser.state() == UpdateParser::State::Failed) break;
        if (parser.complete()) {  // all announced bytes are here: answer while the connection is open
            ok = parser.end();
            break;
        }
    }
    if (!ok && parser.error().empty()) parser.end();  // e.g. a stall: abort the slot

    if (ok) {
        store_.putString("ota_pending", parser.version());
        system_info::recordSlotVersion(store_, esp_ota_get_boot_partition(), parser.version().c_str());
        store_.putString("ota_from", versionString());
        source.reply("OK " + parser.version() + (parser.isDowngrade() ? " (older than the installed one)" : ""));
        notify("Update " + parser.version() + " installed (" + via + "), restarting", NotificationLevel::Info);
        phase_ = Phase::Installed;
    } else {
        std::string why = parser.error().empty() ? "transfer interrupted" : parser.error();
        source.reply("ERR " + why);
        notify("Update refused: " + why, NotificationLevel::Warning);
        phase_ = Phase::Idle;
    }
}

void UpdateService::installFromSd(const std::string& path) {
    if (phase_ != Phase::Idle) return;
    storage_.runJob([this, path]() {
        File f = SD.open(path.c_str());
        if (!f) {
            notify("Can't open " + path, NotificationLevel::Warning);
            return;
        }
        FileSource src(f);
        install(src, "SD card");
        f.close();
    });
}

void UpdateService::taskEntry(void* self) { static_cast<UpdateService*>(self)->listen(); }

void UpdateService::listen() {
    NetworkServer server(kPort);
    bool listening = false;
    uint32_t installedAt = 0;
    for (;;) {
        bool connected = wifi_.state() == WifiController::State::Connected;
        if (connected && !listening) {
            server.begin();
            listening = true;
        } else if (!connected && listening) {
            server.end();
            listening = false;
        }
        // The main loop restarts into an installed update when it's safe. If it never does (stuck,
        // or waiting on someone typing for too long), restart from here: the update must not wait.
        if (phase_ == Phase::Installed) {
            if (!installedAt) installedAt = millis();
            if (millis() - installedAt > kForceRestartMs) {
                ESP_LOGW("update", "the main loop never restarted into the update: restarting");
                esp_restart();
            }
        }
        if (listening && phase_ == Phase::Idle) {
            Counted<NetworkClient> client(server.accept(), net::User::Updates);
            if (client) {
                client.setNoDelay(true);
                NetSource src(client);
                install(src, "Wi-Fi");
                delay(50);
                client.stop();
            }
        }
        vTaskDelay(pdMS_TO_TICKS(200));
    }
}

}  // namespace roro