#include <unity.h>

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

#include "png_rgb332.h"

using namespace roro::png;

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

static uint32_t be(const std::vector<uint8_t>& d, size_t at) {
    return (uint32_t(d[at]) << 24) | (uint32_t(d[at + 1]) << 16) | (uint32_t(d[at + 2]) << 8) | d[at + 3];
}

void test_the_checksums_match_their_reference_values() {
    const char* nine = "123456789";
    TEST_ASSERT_EQUAL_HEX32(0xCBF43926, crc32(0, reinterpret_cast<const uint8_t*>(nine), 9));
    TEST_ASSERT_EQUAL_HEX32(0x11E60398, adler32(1, reinterpret_cast<const uint8_t*>("Wikipedia"), 9));
    // Running: in two pieces, the same.
    uint32_t c = crc32(0, reinterpret_cast<const uint8_t*>(nine), 4);
    TEST_ASSERT_EQUAL_HEX32(0xCBF43926, crc32(c, reinterpret_cast<const uint8_t*>(nine + 4), 5));
}

// Walks the chunks of what was written: every length adds up and every CRC is right.
void test_a_small_image_is_a_well_formed_png() {
    std::vector<uint8_t> out;
    const int w = 4, h = 3;
    Rgb332Writer writer(w, h, [&](const uint8_t* d, size_t n) {
        out.insert(out.end(), d, d + n);
        return true;
    });
    TEST_ASSERT_TRUE(writer.begin());
    const uint8_t rows[3][4] = {{0x00, 0xE0, 0x1C, 0x03}, {0xFF, 0x80, 0x10, 0x02}, {1, 2, 3, 4}};
    for (auto& r : rows) TEST_ASSERT_TRUE(writer.row(r));
    TEST_ASSERT_FALSE(writer.row(rows[0]));  // no more rows than the height
    TEST_ASSERT_TRUE(writer.end());
    TEST_ASSERT_EQUAL(Rgb332Writer::fileSize(w, h), out.size());

    const uint8_t signature[] = {0x89, 'P', 'N', 'G', '\r', '\n', 0x1A, '\n'};
    TEST_ASSERT_EQUAL_MEMORY(signature, out.data(), 8);
    std::vector<std::string> kinds;
    std::vector<uint8_t> idat;
    for (size_t at = 8; at < out.size();) {
        uint32_t len = be(out, at);
        std::string kind(out.begin() + at + 4, out.begin() + at + 8);
        kinds.push_back(kind);
        TEST_ASSERT_EQUAL_HEX32(crc32(0, out.data() + at + 4, len + 4), be(out, at + 8 + len));
        if (kind == "IDAT") idat.assign(out.begin() + at + 8, out.begin() + at + 8 + len);
        if (kind == "IHDR") {
            TEST_ASSERT_EQUAL(w, be(out, at + 8));
            TEST_ASSERT_EQUAL(h, be(out, at + 12));
            TEST_ASSERT_EQUAL(8, out[at + 16]);
            TEST_ASSERT_EQUAL(3, out[at + 17]);  // indexed
        }
        if (kind == "PLTE") {
            TEST_ASSERT_EQUAL(768, len);
            const uint8_t* white = out.data() + at + 8 + 255 * 3;
            TEST_ASSERT_EQUAL(255, white[0] + 0);
            TEST_ASSERT_EQUAL(255, white[1] + 0);
            TEST_ASSERT_EQUAL(255, white[2] + 0);
            const uint8_t* red = out.data() + at + 8 + 0xE0 * 3;
            TEST_ASSERT_EQUAL(255, red[0] + 0);
            TEST_ASSERT_EQUAL(0, red[1] + 0);
        }
        at += 12 + len;
    }
    TEST_ASSERT_EQUAL(4, kinds.size());
    TEST_ASSERT_EQUAL_STRING("IHDR", kinds[0].c_str());
    TEST_ASSERT_EQUAL_STRING("PLTE", kinds[1].c_str());
    TEST_ASSERT_EQUAL_STRING("IDAT", kinds[2].c_str());
    TEST_ASSERT_EQUAL_STRING("IEND", kinds[3].c_str());

    // The zlib stream: one stored, final block of (w + 1) * h bytes, then their Adler-32.
    const size_t raw = (w + 1) * h;
    TEST_ASSERT_EQUAL(2 + 5 + raw + 4, idat.size());
    TEST_ASSERT_EQUAL_HEX8(0x78, idat[0]);
    TEST_ASSERT_EQUAL(0, ((idat[0] << 8) | idat[1]) % 31);  // the header's own check
    TEST_ASSERT_EQUAL_HEX8(0x01, idat[2]);                   // final, stored
    TEST_ASSERT_EQUAL(raw, idat[3] | (idat[4] << 8));
    TEST_ASSERT_EQUAL(static_cast<uint16_t>(~raw), idat[5] | (idat[6] << 8));
    TEST_ASSERT_EQUAL(0, idat[7]);          // the first row's filter byte: none
    TEST_ASSERT_EQUAL_HEX8(0xE0, idat[9]);  // its second pixel
    TEST_ASSERT_EQUAL_HEX32(adler32(1, idat.data() + 7, raw), be(idat, 7 + raw));
}

void test_the_screen_fits_and_a_bigger_image_is_refused() {
    TEST_ASSERT_EQUAL(8 + 25 + 780 + 12 + (2 + 5 + 241 * 135 + 4) + 12, Rgb332Writer::fileSize(240, 135));
    int calls = 0;
    Rgb332Writer screen(240, 135, [&](const uint8_t*, size_t) { return ++calls > 0; });
    TEST_ASSERT_TRUE(screen.begin());
    Rgb332Writer big(320, 240, [](const uint8_t*, size_t) { return true; });  // 77,040 bytes: more than one block holds
    TEST_ASSERT_FALSE(big.begin());
}

void test_a_sink_that_fails_stops_it_and_too_few_rows_dont_end() {
    Rgb332Writer failing(4, 3, [](const uint8_t*, size_t) { return false; });
    TEST_ASSERT_FALSE(failing.begin());
    Rgb332Writer shortOne(4, 3, [](const uint8_t*, size_t) { return true; });
    TEST_ASSERT_TRUE(shortOne.begin());
    const uint8_t row[4] = {};
    TEST_ASSERT_TRUE(shortOne.row(row));
    TEST_ASSERT_FALSE(shortOne.end());
}

int main(int, char**) {
    UNITY_BEGIN();
    RUN_TEST(test_the_checksums_match_their_reference_values);
    RUN_TEST(test_a_small_image_is_a_well_formed_png);
    RUN_TEST(test_the_screen_fits_and_a_bigger_image_is_refused);
    RUN_TEST(test_a_sink_that_fails_stops_it_and_too_few_rows_dont_end);
    return UNITY_END();
}