#include <unity.h>

#include <cstring>
#include <string>
#include <vector>

#include "probation.h"
#include "safe_mode.h"
#include "sha256.h"
#include "update_parser.h"
#include "version_compare.h"

using namespace roro;

void setUp() {}
void tearDown() {}

// ---- helpers ----------------------------------------------------------------

static std::string hex(const uint8_t* d, size_t n) {
    static const char* h = "0123456789abcdef";
    std::string s;
    for (size_t i = 0; i < n; i++) {
        s += h[d[i] >> 4];
        s += h[d[i] & 15];
    }
    return s;
}

// Accepts a signature equal to the first bytes of the digest it is asked about ("signed" by us).
struct FakeVerifier : SignatureVerifier {
    bool verify(const uint8_t digest[32], const uint8_t* sig, size_t len) override {
        calls++;
        return len == 8 && std::memcmp(digest, sig, 8) == 0;
    }
    int calls = 0;
};

struct MemorySink : UpdateSink {
    bool begin(size_t size) override {
        began = true;
        expected = size;
        return true;
    }
    bool write(const uint8_t* d, size_t n) override {
        data.insert(data.end(), d, d + n);
        return true;
    }
    bool finish() override { return finished = true; }
    void abort() override { aborted = true; }
    bool began = false, finished = false, aborted = false;
    size_t expected = 0;
    std::vector<uint8_t> data;
};

// Builds an Update File the way make_ota.py does, signing with FakeVerifier's scheme.
static std::vector<uint8_t> makeFile(const std::string& version, const std::vector<uint8_t>& image, bool goodSig = true) {
    std::vector<uint8_t> h(update::kHeaderSize, 0);
    std::memcpy(h.data(), update::kMagic, 8);
    h[8] = update::kFormat & 0xFF;
    h[10] = update::kHeaderSize & 0xFF;
    h[11] = update::kHeaderSize >> 8;
    uint32_t n = static_cast<uint32_t>(image.size());
    for (int i = 0; i < 4; i++) h[12 + i] = (n >> (8 * i)) & 0xFF;
    Sha256::hash(image.data(), image.size(), &h[16]);
    std::strncpy(reinterpret_cast<char*>(&h[48]), version.c_str(), 31);
    uint8_t digest[32];
    Sha256::hash(h.data(), update::kSignedBytes, digest);
    h[80] = 8;
    std::memcpy(&h[82], digest, 8);
    if (!goodSig) h[82] ^= 0xFF;
    h.insert(h.end(), image.begin(), image.end());
    return h;
}

static std::vector<uint8_t> image(size_t n) {
    std::vector<uint8_t> v(n);
    for (size_t i = 0; i < n; i++) v[i] = static_cast<uint8_t>(i * 7 + 3);
    return v;
}

// ---- SHA-256 ----------------------------------------------------------------

void test_sha256_known_vectors() {
    uint8_t d[32];
    Sha256::hash(reinterpret_cast<const uint8_t*>(""), 0, d);
    TEST_ASSERT_EQUAL_STRING("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", hex(d, 32).c_str());
    Sha256::hash(reinterpret_cast<const uint8_t*>("abc"), 3, d);
    TEST_ASSERT_EQUAL_STRING("ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad", hex(d, 32).c_str());
    const char* two = "abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq";
    Sha256::hash(reinterpret_cast<const uint8_t*>(two), std::strlen(two), d);
    TEST_ASSERT_EQUAL_STRING("248d6a61d20638b8e5c026930c3e6039a33ce45964ff2167f6ecedd419db06c1", hex(d, 32).c_str());
}

void test_sha256_streaming_matches_one_shot() {
    auto img = image(1000);
    uint8_t a[32], b[32];
    Sha256::hash(img.data(), img.size(), a);
    Sha256 s;
    for (size_t i = 0; i < img.size(); i += 37) s.update(img.data() + i, std::min<size_t>(37, img.size() - i));
    s.finish(b);
    TEST_ASSERT_EQUAL_MEMORY(a, b, 32);
}

// ---- parser -----------------------------------------------------------------

void test_valid_file_in_one_chunk_is_installed() {
    FakeVerifier v;
    MemorySink sink;
    auto img = image(5000);
    auto file = makeFile("v0.3.0", img);
    UpdateParser p(v, sink, 100000, "v0.2.1");
    p.feed(file.data(), file.size());
    TEST_ASSERT_TRUE(p.end());
    TEST_ASSERT_EQUAL(static_cast<int>(UpdateParser::State::Done), static_cast<int>(p.state()));
    TEST_ASSERT_TRUE(sink.finished);
    TEST_ASSERT_TRUE(sink.data == img);
    TEST_ASSERT_EQUAL_STRING("v0.3.0", p.version().c_str());
    TEST_ASSERT_FALSE(p.isDowngrade());
}

void test_valid_file_byte_by_byte() {
    FakeVerifier v;
    MemorySink sink;
    auto img = image(700);
    auto file = makeFile("v0.3.0", img);
    UpdateParser p(v, sink, 100000, "v0.2.1");
    for (auto b : file) p.feed(&b, 1);
    TEST_ASSERT_TRUE(p.end());
    TEST_ASSERT_TRUE(sink.data == img);
}

void test_progress_counts_image_bytes() {
    FakeVerifier v;
    MemorySink sink;
    auto file = makeFile("v0.3.0", image(1000));
    UpdateParser p(v, sink, 100000, "v0.2.1");
    p.feed(file.data(), update::kHeaderSize + 250);
    TEST_ASSERT_EQUAL(25, p.percent());
}

void test_complete_once_the_declared_image_size_has_arrived() {
    FakeVerifier v;
    MemorySink sink;
    auto file = makeFile("v0.3.0", image(300));
    UpdateParser p(v, sink, 100000, "v0.2.1");
    p.feed(file.data(), file.size() - 1);
    TEST_ASSERT_FALSE(p.complete());
    p.feed(file.data() + file.size() - 1, 1);
    TEST_ASSERT_TRUE(p.complete());
    TEST_ASSERT_TRUE(p.end());
}

void test_not_an_update_file_is_refused_before_writing() {
    FakeVerifier v;
    MemorySink sink;
    auto file = makeFile("v0.3.0", image(100));
    file[0] = 'X';
    UpdateParser p(v, sink, 100000, "v0.2.1");
    p.feed(file.data(), file.size());
    TEST_ASSERT_EQUAL(static_cast<int>(UpdateParser::State::Failed), static_cast<int>(p.state()));
    TEST_ASSERT_FALSE(sink.began);
    TEST_ASSERT_FALSE(p.error().empty());
}

void test_bad_signature_is_refused_before_writing() {
    FakeVerifier v;
    MemorySink sink;
    auto file = makeFile("v0.3.0", image(100), false);
    UpdateParser p(v, sink, 100000, "v0.2.1");
    p.feed(file.data(), file.size());
    TEST_ASSERT_EQUAL(static_cast<int>(UpdateParser::State::Failed), static_cast<int>(p.state()));
    TEST_ASSERT_FALSE(sink.began);
    TEST_ASSERT_TRUE(p.error().find("signature") != std::string::npos);
}

void test_tampered_version_breaks_the_signature() {
    FakeVerifier v;
    MemorySink sink;
    auto file = makeFile("v0.3.0", image(100));
    file[48 + 3] = '9';  // v0.3.0 -> v0.9.0 after signing
    UpdateParser p(v, sink, 100000, "v0.2.1");
    p.feed(file.data(), file.size());
    TEST_ASSERT_EQUAL(static_cast<int>(UpdateParser::State::Failed), static_cast<int>(p.state()));
}

void test_corrupted_image_is_aborted_at_the_end() {
    FakeVerifier v;
    MemorySink sink;
    auto file = makeFile("v0.3.0", image(500));
    file[update::kHeaderSize + 123] ^= 1;
    UpdateParser p(v, sink, 100000, "v0.2.1");
    p.feed(file.data(), file.size());
    TEST_ASSERT_FALSE(p.end());
    TEST_ASSERT_TRUE(sink.aborted);
    TEST_ASSERT_FALSE(sink.finished);
}

void test_truncated_file_is_aborted() {
    FakeVerifier v;
    MemorySink sink;
    auto file = makeFile("v0.3.0", image(500));
    UpdateParser p(v, sink, 100000, "v0.2.1");
    p.feed(file.data(), file.size() - 10);
    TEST_ASSERT_FALSE(p.end());
    TEST_ASSERT_TRUE(sink.aborted);
}

void test_trailing_bytes_are_refused() {
    FakeVerifier v;
    MemorySink sink;
    auto file = makeFile("v0.3.0", image(100));
    file.push_back(0);
    UpdateParser p(v, sink, 100000, "v0.2.1");
    p.feed(file.data(), file.size());
    TEST_ASSERT_FALSE(p.end());
    TEST_ASSERT_FALSE(sink.finished);
}

void test_image_larger_than_the_slot_is_refused_before_writing() {
    FakeVerifier v;
    MemorySink sink;
    auto file = makeFile("v0.3.0", image(2000));
    UpdateParser p(v, sink, 1000, "v0.2.1");
    p.feed(file.data(), file.size());
    TEST_ASSERT_EQUAL(static_cast<int>(UpdateParser::State::Failed), static_cast<int>(p.state()));
    TEST_ASSERT_FALSE(sink.began);
}

void test_downgrade_is_flagged_but_allowed() {
    FakeVerifier v;
    MemorySink sink;
    auto file = makeFile("v0.2.0", image(100));
    UpdateParser p(v, sink, 100000, "v0.2.1");
    p.feed(file.data(), file.size());
    TEST_ASSERT_TRUE(p.end());
    TEST_ASSERT_TRUE(p.isDowngrade());
}

// ---- versions ---------------------------------------------------------------

void test_version_comparison() {
    TEST_ASSERT_TRUE(versionOlder("v0.2.0", "v0.2.1"));
    TEST_ASSERT_TRUE(versionOlder("v0.9.9", "v0.10.0"));
    TEST_ASSERT_TRUE(versionOlder("v0.2.1", "v1.0.0"));
    TEST_ASSERT_FALSE(versionOlder("v0.2.1", "v0.2.1"));
    TEST_ASSERT_FALSE(versionOlder("v0.3.0", "v0.2.1"));
    // build suffixes from git describe compare on the release part
    TEST_ASSERT_FALSE(versionOlder("v0.2.1-3-gabc1234", "v0.2.1"));
    TEST_ASSERT_TRUE(versionOlder("v0.2.1-3-gabc1234-dirty", "v0.3.0"));
    // unknown versions are never called older
    TEST_ASSERT_FALSE(versionOlder("unknown", "v0.2.1"));
    TEST_ASSERT_FALSE(versionOlder("v0.2.1", "4c0be89"));
}

// ---- probation --------------------------------------------------------------

static int judge(uint32_t up, bool frame, bool configured, bool connected) {
    return static_cast<int>(Probation::judge(up, frame, configured, connected));
}
static const int kWait = static_cast<int>(Probation::Verdict::Wait);
static const int kConfirm = static_cast<int>(Probation::Verdict::Confirm);
static const int kRollBack = static_cast<int>(Probation::Verdict::RollBack);

void test_probation_waits_30_seconds_and_a_first_frame() {
    TEST_ASSERT_EQUAL(kWait, judge(29999, true, false, false));
    TEST_ASSERT_EQUAL(kWait, judge(60000, false, false, false));
    TEST_ASSERT_EQUAL(kConfirm, judge(30000, true, false, false));
}

void test_probation_needs_wifi_when_it_is_configured() {
    TEST_ASSERT_EQUAL(kWait, judge(60000, true, true, false));
    TEST_ASSERT_EQUAL(kConfirm, judge(60000, true, true, true));
}

void test_probation_rolls_back_when_wifi_never_comes() {
    TEST_ASSERT_EQUAL(kWait, judge(179999, true, true, false));
    TEST_ASSERT_EQUAL(kRollBack, judge(180000, true, true, false));
    TEST_ASSERT_EQUAL(kConfirm, judge(500000, true, false, false));  // no Wi-Fi configured: fine
}

void test_rollback_at_boot_after_an_unconfirmed_start() {
    TEST_ASSERT_FALSE(Probation::rollBackAtBoot(true, 0));   // first start of new firmware
    TEST_ASSERT_TRUE(Probation::rollBackAtBoot(true, 1));    // it died before confirming
    TEST_ASSERT_FALSE(Probation::rollBackAtBoot(false, 3));  // confirmed firmware: never
}

void test_safe_mode_after_three_crashes_in_a_row() {
    int crashes = 0;
    crashes = SafeMode::countAtBoot(true, crashes);
    crashes = SafeMode::countAtBoot(true, crashes);
    TEST_ASSERT_FALSE(SafeMode::active(crashes));
    crashes = SafeMode::countAtBoot(true, crashes);
    TEST_ASSERT_TRUE(SafeMode::active(crashes));
}

void test_safe_mode_count_restarts_after_a_normal_start() {
    int crashes = SafeMode::countAtBoot(true, 2);
    TEST_ASSERT_TRUE(SafeMode::active(crashes));
    crashes = SafeMode::countAtBoot(false, crashes);  // e.g. `reboot` from the Debug Console
    TEST_ASSERT_EQUAL(0, crashes);
    TEST_ASSERT_FALSE(SafeMode::active(crashes));
}

int main() {
    UNITY_BEGIN();
    RUN_TEST(test_sha256_known_vectors);
    RUN_TEST(test_sha256_streaming_matches_one_shot);
    RUN_TEST(test_valid_file_in_one_chunk_is_installed);
    RUN_TEST(test_valid_file_byte_by_byte);
    RUN_TEST(test_progress_counts_image_bytes);
    RUN_TEST(test_complete_once_the_declared_image_size_has_arrived);
    RUN_TEST(test_not_an_update_file_is_refused_before_writing);
    RUN_TEST(test_bad_signature_is_refused_before_writing);
    RUN_TEST(test_tampered_version_breaks_the_signature);
    RUN_TEST(test_corrupted_image_is_aborted_at_the_end);
    RUN_TEST(test_truncated_file_is_aborted);
    RUN_TEST(test_trailing_bytes_are_refused);
    RUN_TEST(test_image_larger_than_the_slot_is_refused_before_writing);
    RUN_TEST(test_downgrade_is_flagged_but_allowed);
    RUN_TEST(test_version_comparison);
    RUN_TEST(test_probation_waits_30_seconds_and_a_first_frame);
    RUN_TEST(test_probation_needs_wifi_when_it_is_configured);
    RUN_TEST(test_probation_rolls_back_when_wifi_never_comes);
    RUN_TEST(test_rollback_at_boot_after_an_unconfirmed_start);
    RUN_TEST(test_safe_mode_after_three_crashes_in_a_row);
    RUN_TEST(test_safe_mode_count_restarts_after_a_normal_start);
    return UNITY_END();
}