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

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

#include <memory>

#include <ctime>

#include "http_head.h"
#include "platform/https_get.h"
#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_;
};

// A release being downloaded from Gitea: the same bytes a push or a card would give.
class GiteaSource : public UpdateSource {
   public:
    GiteaSource(HttpsGet& get, long cutAfter, long flipAt) : get_(get), cut_(cutAfter), flip_(flipAt) {}
    int read(uint8_t* buf, size_t len) override {
        if (cut_ >= 0 && pos_ >= cut_) return -1;  // Debug Builds: the connection "breaks" here
        int n = get_.read(buf, len);
        if (n > 0 && flip_ >= pos_ && flip_ < pos_ + n) buf[flip_ - pos_] ^= 1;  // and a byte "arrives wrong"
        if (n > 0) pos_ += n;
        return n;
    }

   private:
    HttpsGet& get_;
    long cut_, flip_, pos_ = 0;
};

constexpr time_t kClockSetAfter = 1700000000;  // anything earlier is the clock's default

}  // 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_failed", pending);  // not announced again until there's a newer one (Q168)
        store_.putString("ota_pending", "");
    }

    if (!task_) xTaskCreate(taskEntry, "update", 7168, this, 1, &task_);  // peak 5.4 KB: TLS to Gitea and the signature check (#6)
}

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) {
    scheduleDailyCheck(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();
    });
}

std::string UpdateService::failedVersion() const {
    std::string failed;
    store_.getString("ota_failed", failed);
    return failed;
}

std::string UpdateService::requestInstall(const std::string& tag, bool force) {
    if (phase_ != Phase::Idle || giteaBusy()) return "Busy with something else";
    if (!force && release::isDebugBuild(runningVersion())) return "Debug Build: update from the PC";  // short: it is drawn in one line at the screen's foot
    if (wifi_.state() != WifiController::State::Connected) return "No Wi-Fi";
    release::Release r;
    if (!gitea_.find(tag, r)) return "Look for releases first";
    if (!release::trustedAssetUrl(r.otaUrl)) return "That download isn't on the project's server";
    installTag_ = tag;
    request_ = Request::Install;
    return "";
}

// Once a day while Wi-Fi is up and the Clock is set (Q165). Skipped, and tried again later, when
// something else is going on or memory is short; never during Probation or an install.
void UpdateService::scheduleDailyCheck(uint32_t nowMs) {
    if (!dailyCheck_ || !settings_.getBool(Setting::CheckUpdates)) return;
    if (static_cast<int32_t>(nowMs - nextCheckMs_) < 0 || probation_ || phase_ != Phase::Idle || giteaBusy()) return;
    if (wifi_.state() != WifiController::State::Connected) return;
    time_t now = time(nullptr);
    if (now < kClockSetAfter) return;
    int32_t today = static_cast<int32_t>(now / 86400), last = 0;
    store_.getInt("rel_day", last);
    if (today == last) {
        nextCheckMs_ = nowMs + 3600000;  // look again in an hour, in case the day has turned
        return;
    }
    // Never at IRC's expense: with IRC connected there isn't the room (Q172), so this waits.
    if (esp_get_free_heap_size() < HttpsGet::kNeedFree) {
        nextCheckMs_ = nowMs + 600000;
        return;
    }
    nextCheckMs_ = nowMs + 1800000;  // if this one fails
    request_ = Request::BackgroundCheck;
}

// What a person asked for (a check, a list, an install) may take IRC offline for a few seconds:
// a TLS connection needs about 80 KB and IRC's own holds 40 KB of what there is (R1, Q172).
bool UpdateService::makeRoom() {
    if (esp_get_free_heap_size() >= HttpsGet::kNeedFree) return true;
    if (!holdMemory) return false;
    held_ = true;
    holdMemory(true);
    for (int i = 0; i < 40 && esp_get_free_heap_size() < HttpsGet::kNeedFree; i++) delay(100);  // its TLS session goes
    return esp_get_free_heap_size() >= HttpsGet::kNeedFree;
}

void UpdateService::serve() {
    Request r = request_;
    if (r == Request::None) return;
    request_ = Request::None;
    serving_ = true;
    held_ = false;
    if (r != Request::BackgroundCheck && r != Request::None) {
        bool ok = makeRoom();
        if (!ok) {
            gitea_.fail("Not enough memory: close a Gemini page");
            if (r == Request::Install) notify("Update refused: not enough memory", NotificationLevel::Warning);
            r = Request::None;
        }
    }
    switch (r) {
        case Request::Check:
        case Request::BackgroundCheck: {
            if (!gitea_.fetchLatest()) break;
            time_t now = time(nullptr);
            if (now >= kClockSetAfter) store_.putInt("rel_day", static_cast<int32_t>(now / 86400));
            release::Release latest;
            if (r == Request::BackgroundCheck && gitea_.latest(latest) &&
                release::shouldAnnounce(latest, runningVersion(), failedVersion(), announced_)) {
                announced_ = latest.tag;
                notify(latest.tag + " is out: see Settings > Firmware", NotificationLevel::Info);  // 48 bytes at most
            }
            break;
        }
        case Request::List: gitea_.fetchList(); break;
        case Request::Install: {
            release::Release release;
            if (gitea_.find(installTag_, release)) installFromGitea(release);
            break;
        }
#ifdef RORO_DEBUG
        case Request::Probe: {
            HttpsGet get(net::User::Updates);
            std::string why = get.open(probeHost_, probePath_, "*/*");
            probeResult_ = why.empty() ? "accepted, the server answered 200" : why;
            break;
        }
#endif
        case Request::None: break;
    }
    // A check or a list someone asked for that failed says so; the daily one stays quiet.
    if ((r == Request::Check || r == Request::List) && gitea_.status() == GiteaReleases::Status::Failed)
        notify(gitea_.error(), NotificationLevel::Warning);
    // IRC comes back, unless the device is about to restart into an update.
    if (held_ && phase_ != Phase::Installed && holdMemory) holdMemory(false);
    held_ = false;
    serving_ = false;
}

void UpdateService::installFromGitea(const release::Release& r) {
    release::Url url = release::parseUrl(r.otaUrl);
    incoming_ = r.tag;
    percent_ = 0;
    phase_ = Phase::Receiving;
    HttpsGet get(net::User::Updates);
    std::string why = get.open(url.host, url.path, "application/octet-stream");
    if (why.empty() && r.otaSize && get.contentLength() >= 0 && static_cast<uint32_t>(get.contentLength()) != r.otaSize)
        why = "The file isn't the size the server listed";
    if (!why.empty()) {
        phase_ = Phase::Idle;
        notify("Update refused: " + why, NotificationLevel::Warning);
        return;
    }
#ifdef RORO_DEBUG
    GiteaSource source(get, damageCut_, damageFlip_);
    damageCut_ = damageFlip_ = -1;
#else
    GiteaSource source(get, -1, -1);
#endif
    install(source, "Gitea");
}

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 (phase_ == Phase::Idle) serve();
        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