#include <unity.h>

#include <string>
#include <vector>

#include "file_receiver.h"
#include "sha256.h"

using namespace roro;

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

static std::string hexSha(const std::vector<uint8_t>& data) {
    uint8_t d[32];
    Sha256::hash(data.data(), data.size(), d);
    static const char* hex = "0123456789abcdef";
    std::string s;
    for (uint8_t b : d) {
        s += hex[b >> 4];
        s += hex[b & 15];
    }
    return s;
}

static std::vector<uint8_t> bytes(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;
}

static std::string args(const std::string& path, const std::vector<uint8_t>& data) {
    return path + " " + std::to_string(data.size()) + " " + hexSha(data);
}

void test_begin_accepts_path_size_and_checksum() {
    FileReceiver r;
    auto data = bytes(10);
    TEST_ASSERT_EQUAL_STRING("", r.begin(args("/updates/a.ota", data), 0).c_str());
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Receiving);
    TEST_ASSERT_EQUAL_STRING("/updates/a.ota", r.path().c_str());
    TEST_ASSERT_EQUAL_STRING("/updates/a.ota.part", r.partPath().c_str());
    TEST_ASSERT_EQUAL_UINT32(10, r.size());
}

void test_begin_refuses_bad_arguments() {
    FileReceiver r;
    std::string sha(64, 'a');
    TEST_ASSERT_NOT_EQUAL(0, r.begin("updates/a.ota 10 " + sha, 0).size());   // not absolute
    TEST_ASSERT_NOT_EQUAL(0, r.begin("/a/../b.ota 10 " + sha, 0).size());     // climbs out
    TEST_ASSERT_NOT_EQUAL(0, r.begin("/a.ota 0 " + sha, 0).size());           // empty
    TEST_ASSERT_NOT_EQUAL(0, r.begin("/a.ota 99999999 " + sha, 0).size());    // too big
    TEST_ASSERT_NOT_EQUAL(0, r.begin("/a.ota 1x " + sha, 0).size());          // not a number
    TEST_ASSERT_NOT_EQUAL(0, r.begin("/a.ota 10 abc", 0).size());             // short checksum
    TEST_ASSERT_NOT_EQUAL(0, r.begin("/a.ota 10 " + std::string(64, 'g'), 0).size());
    TEST_ASSERT_NOT_EQUAL(0, r.begin("/a.ota 10", 0).size());
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Idle);
}

void test_bytes_are_handed_out_one_chunk_at_a_time() {
    FileReceiver r;
    auto data = bytes(FileReceiver::kChunk * 2 + 100);
    r.begin(args("/f.bin", data), 0);

    // Everything arrives at once: only the first chunk is taken.
    TEST_ASSERT_EQUAL(FileReceiver::kChunk, r.feed(data.data(), data.size(), 1));
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Writing);
    TEST_ASSERT_EQUAL(FileReceiver::kChunk, r.chunk().size());
    TEST_ASSERT_EQUAL(0, r.feed(data.data() + FileReceiver::kChunk, 10, 2));  // nothing more while writing

    r.chunkWritten(true, 3);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Receiving);
    TEST_ASSERT_EQUAL_UINT32(FileReceiver::kChunk, r.received());

    r.feed(data.data() + FileReceiver::kChunk, FileReceiver::kChunk, 4);
    r.chunkWritten(true, 5);
    // The last, short chunk is complete as soon as the announced size is reached.
    TEST_ASSERT_EQUAL(100, r.feed(data.data() + 2 * FileReceiver::kChunk, 100, 6));
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Writing);
    TEST_ASSERT_EQUAL(100, r.chunk().size());
    TEST_ASSERT_EQUAL_UINT8(data[2 * FileReceiver::kChunk], r.chunk()[0]);

    r.chunkWritten(true, 7);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Finishing);
    r.finished(true);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Done);
}

void test_chunk_split_across_reads() {
    FileReceiver r;
    auto data = bytes(30);
    r.begin(args("/f.bin", data), 0);
    TEST_ASSERT_EQUAL(10, r.feed(data.data(), 10, 1));
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Receiving);
    TEST_ASSERT_EQUAL(20, r.feed(data.data() + 10, 20, 2));
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Writing);
    TEST_ASSERT_EQUAL(30, r.chunk().size());
}

void test_wrong_checksum_fails_before_the_last_write() {
    FileReceiver r;
    auto data = bytes(50);
    r.begin(args("/f.bin", data), 0);
    data[20] ^= 1;  // corrupted on the way
    r.feed(data.data(), data.size(), 1);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Failed);
    TEST_ASSERT_EQUAL_STRING("checksum mismatch", r.error().c_str());
}

void test_failed_write_fails_the_transfer() {
    FileReceiver r;
    auto data = bytes(50);
    r.begin(args("/f.bin", data), 0);
    r.feed(data.data(), data.size(), 1);
    r.chunkWritten(false, 2);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Failed);
    TEST_ASSERT_EQUAL_STRING("write failed", r.error().c_str());
}

void test_failed_rename_fails_the_transfer() {
    FileReceiver r;
    auto data = bytes(5);
    r.begin(args("/f.bin", data), 0);
    r.feed(data.data(), data.size(), 1);
    r.chunkWritten(true, 2);
    r.finished(false);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Failed);
}

void test_silence_times_out() {
    FileReceiver r;
    auto data = bytes(FileReceiver::kChunk * 2);
    r.begin(args("/f.bin", data), 1000);
    r.tick(1000 + FileReceiver::kTimeoutMs - 1);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Receiving);
    r.feed(data.data(), 10, 5000);  // bytes reset the clock
    r.tick(5000 + FileReceiver::kTimeoutMs - 1);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Receiving);
    r.tick(5000 + FileReceiver::kTimeoutMs);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Failed);
    TEST_ASSERT_EQUAL_STRING("timed out", r.error().c_str());
}

void test_a_write_that_never_completes_times_out() {
    // With no card mounted, the storage task drops the job and never answers.
    FileReceiver r;
    auto data = bytes(10);
    r.begin(args("/f.bin", data), 0);
    r.feed(data.data(), data.size(), 100);
    r.tick(100 + FileReceiver::kTimeoutMs);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Failed);
}

void test_wanted_counts_down_within_a_chunk() {
    FileReceiver r;
    auto data = bytes(FileReceiver::kChunk + 10);
    TEST_ASSERT_EQUAL(0, r.wanted());
    r.begin(args("/f.bin", data), 0);
    TEST_ASSERT_EQUAL(FileReceiver::kChunk, r.wanted());
    r.feed(data.data(), 24, 1);
    TEST_ASSERT_EQUAL(FileReceiver::kChunk - 24, r.wanted());
    r.feed(data.data() + 24, FileReceiver::kChunk - 24, 2);
    TEST_ASSERT_EQUAL(0, r.wanted());  // writing
    r.chunkWritten(true, 3);
    TEST_ASSERT_EQUAL(10, r.wanted());  // the short last chunk
}

void test_a_rename_that_never_completes_times_out() {
    FileReceiver r;
    auto data = bytes(10);
    r.begin(args("/f.bin", data), 0);
    r.feed(data.data(), data.size(), 1);
    r.chunkWritten(true, 2);
    r.tick(2 + FileReceiver::kTimeoutMs);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Failed);
}

void test_reset_returns_to_idle() {
    FileReceiver r;
    auto data = bytes(10);
    r.begin(args("/f.bin", data), 0);
    r.reset();
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Idle);
    TEST_ASSERT_FALSE(r.active());
}

// The received bytes can be right and the card's copy wrong: read back, it must hash the same.
void test_card_check() {
    auto data = bytes(10);
    FileReceiver r;
    r.begin(args("/a.bin", data), 0);
    r.feed(data.data(), data.size(), 0);
    r.chunkWritten(true, 0);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Finishing);
    uint8_t good[32];
    Sha256::hash(data.data(), data.size(), good);
    r.cardChecked(good);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Finishing);
    r.finished(true);
    TEST_ASSERT_TRUE(r.state() == FileReceiver::State::Done);

    FileReceiver bad;
    bad.begin(args("/a.bin", data), 0);
    bad.feed(data.data(), data.size(), 0);
    bad.chunkWritten(true, 0);
    auto zeroed = data;
    zeroed[4] = 0;
    uint8_t wrong[32];
    Sha256::hash(zeroed.data(), zeroed.size(), wrong);
    bad.cardChecked(wrong);
    TEST_ASSERT_TRUE(bad.state() == FileReceiver::State::Failed);
    TEST_ASSERT_EQUAL_STRING("the copy on the card differs", bad.error().c_str());
}

int main() {
    UNITY_BEGIN();
    RUN_TEST(test_begin_accepts_path_size_and_checksum);
    RUN_TEST(test_begin_refuses_bad_arguments);
    RUN_TEST(test_bytes_are_handed_out_one_chunk_at_a_time);
    RUN_TEST(test_chunk_split_across_reads);
    RUN_TEST(test_wrong_checksum_fails_before_the_last_write);
    RUN_TEST(test_failed_write_fails_the_transfer);
    RUN_TEST(test_failed_rename_fails_the_transfer);
    RUN_TEST(test_silence_times_out);
    RUN_TEST(test_a_write_that_never_completes_times_out);
    RUN_TEST(test_wanted_counts_down_within_a_chunk);
    RUN_TEST(test_a_rename_that_never_completes_times_out);
    RUN_TEST(test_reset_returns_to_idle);
    RUN_TEST(test_card_check);
    return UNITY_END();
}