#include "scramble/fibonacci_scrambler.h"

#include <gtest/gtest.h>

#include <cstdint>
#include <unordered_map>
#include <vector>

namespace {

using radiantwave::scramble::ScrambleConfig;
using radiantwave::scramble::canvasSize;
using radiantwave::scramble::scramble;
using radiantwave::scramble::unscramble;

// W x H RGBA8888 image whose every pixel has a distinct 32-bit value, so a
// permutation can be verified by counting values. Pixel (x, y) gets the value
// y*width + x + 1; the +1 keeps every source value distinct from the
// transparent-black padding (0).
std::vector<uint8_t> uniqueImage(int width, int height)
{
    std::vector<uint8_t> img(static_cast<size_t>(width) * height * 4);
    for (int y = 0; y < height; ++y)
    {
        for (int x = 0; x < width; ++x)
        {
            const uint32_t v = static_cast<uint32_t>(y) * width + x + 1;
            uint8_t *p = img.data() + (static_cast<size_t>(y) * width + x) * 4;
            p[0] = static_cast<uint8_t>(v & 0xFF);
            p[1] = static_cast<uint8_t>((v >> 8) & 0xFF);
            p[2] = static_cast<uint8_t>((v >> 16) & 0xFF);
            p[3] = static_cast<uint8_t>((v >> 24) & 0xFF);
        }
    }
    return img;
}

bool roundTrips(int width, int height, int rounds, uint64_t seed)
{
    const auto img = uniqueImage(width, height);
    const int n = canvasSize(width, height);
    std::vector<uint8_t> scrambled(static_cast<size_t>(n) * n * 4);
    std::vector<uint8_t> recovered(static_cast<size_t>(width) * height * 4);

    ScrambleConfig cfg;
    cfg.rounds = rounds;
    cfg.seed = seed;
    if (!scramble(img.data(), scrambled.data(), width, height, cfg))
    {
        return false;
    }
    if (!unscramble(scrambled.data(), recovered.data(), width, height, cfg))
    {
        return false;
    }
    return recovered == img;
}

} // namespace

TEST(FibonacciScramblerTest, CanvasSizeIsMaxOfDimensions)
{
    EXPECT_EQ(canvasSize(1, 1), 1);
    EXPECT_EQ(canvasSize(8, 8), 8);
    EXPECT_EQ(canvasSize(8, 4), 8);
    EXPECT_EQ(canvasSize(4, 8), 8);
    EXPECT_EQ(canvasSize(1, 8), 8);
    EXPECT_EQ(canvasSize(8, 1), 8);
    EXPECT_EQ(canvasSize(0, 8), 0);
    EXPECT_EQ(canvasSize(8, 0), 0);
    EXPECT_EQ(canvasSize(0, 0), 0);
    EXPECT_EQ(canvasSize(-3, 8), 0);
}

TEST(FibonacciScramblerTest, RoundTripIsExactForSquareImages)
{
    const int sizes[] = {1, 2, 3, 8, 32};
    for (int n : sizes)
    {
        for (int t = 1; t <= 6; ++t)
        {
            EXPECT_TRUE(roundTrips(n, n, t, radiantwave::scramble::kDefaultSeed))
                << "N=" << n << " rounds=" << t;
        }
    }
}

TEST(FibonacciScramblerTest, RoundTripIsExactForWideAndTallImages)
{
    const int dims[][2] = {{8, 4}, {4, 8}, {1, 8}, {8, 1}, {3, 17}, {31, 3}};
    for (const auto &d : dims)
    {
        for (int t : {1, 4})
        {
            EXPECT_TRUE(roundTrips(d[0], d[1], t, radiantwave::scramble::kDefaultSeed))
                << "W=" << d[0] << " H=" << d[1] << " rounds=" << t;
        }
    }
}

TEST(FibonacciScramblerTest, RoundTripIsExactForGlyphShapedTiles)
{
    // Punctuation glyphs are narrow or short, so N is max(W, H): these exercise
    // the non-square padding path with glyph-like proportions (e.g. an
    // apostrophe tile is 6 wide x 38 tall at 32px).
    const int dims[][2] = {{1, 5}, {2, 7}, {6, 1}, {1, 1},
                           {17, 37}, {9, 37}, {37, 9}};
    for (const auto &d : dims)
    {
        for (int t : {1, 4})
        {
            EXPECT_TRUE(roundTrips(d[0], d[1], t, radiantwave::scramble::kDefaultSeed))
                << "W=" << d[0] << " H=" << d[1] << " rounds=" << t;
        }
    }
}

TEST(FibonacciScramblerTest, ScrambleIsABijectionOnTheCanvas)
{
    const int dims[][2] = {{5, 7}, {8, 8}, {3, 1}, {1, 4}};
    for (const auto &d : dims)
    {
        const int width = d[0];
        const int height = d[1];
        const int n = canvasSize(width, height);
        const auto img = uniqueImage(width, height);
        std::vector<uint8_t> scrambled(static_cast<size_t>(n) * n * 4);

        ScrambleConfig cfg;
        cfg.rounds = 3;
        ASSERT_TRUE(scramble(img.data(), scrambled.data(), width, height, cfg));

        std::unordered_map<uint32_t, int> counts;
        for (const uint8_t *p = scrambled.data(); p != scrambled.data() + scrambled.size(); p += 4)
        {
            const uint32_t v = static_cast<uint32_t>(p[0]) |
                               (static_cast<uint32_t>(p[1]) << 8) |
                               (static_cast<uint32_t>(p[2]) << 16) |
                               (static_cast<uint32_t>(p[3]) << 24);
            counts[v]++;
        }

        // Every source pixel value appears exactly once (the map restricted to
        // the W x H region is injective)...
        const size_t pixels = static_cast<size_t>(width) * height;
        for (size_t i = 0; i < pixels; ++i)
        {
            const uint32_t v = static_cast<uint32_t>(i) + 1;
            EXPECT_EQ(counts[v], 1)
                << "W=" << width << " H=" << height << " value=" << v;
        }
        // ...and the remaining canvas positions are transparent-black padding.
        EXPECT_EQ(counts[0], static_cast<int>(static_cast<size_t>(n) * n - pixels))
            << "W=" << width << " H=" << height;
    }
}

TEST(FibonacciScramblerTest, ScrambleChangesTheImage)
{
    const auto img = uniqueImage(8, 8);
    const int n = 8;
    std::vector<uint8_t> scrambled(static_cast<size_t>(n) * n * 4);

    ScrambleConfig cfg;
    cfg.rounds = 1;
    ASSERT_TRUE(scramble(img.data(), scrambled.data(), 8, 8, cfg));

    bool changed = false;
    for (size_t i = 0; i < img.size(); ++i)
    {
        if (scrambled[i] != img[i])
        {
            changed = true;
            break;
        }
    }
    EXPECT_TRUE(changed);
}

TEST(FibonacciScramblerTest, OutputIsDeterministic)
{
    const auto img = uniqueImage(6, 9);
    const int n = 9;
    std::vector<uint8_t> a(static_cast<size_t>(n) * n * 4);
    std::vector<uint8_t> b(static_cast<size_t>(n) * n * 4);

    ScrambleConfig cfg;
    cfg.rounds = 3;
    ASSERT_TRUE(scramble(img.data(), a.data(), 6, 9, cfg));
    ASSERT_TRUE(scramble(img.data(), b.data(), 6, 9, cfg));
    EXPECT_EQ(a, b);
}

TEST(FibonacciScramblerTest, DifferentSeedChangesOutput)
{
    const auto img = uniqueImage(8, 8);
    const int n = 8;
    std::vector<uint8_t> a(static_cast<size_t>(n) * n * 4);
    std::vector<uint8_t> b(static_cast<size_t>(n) * n * 4);

    ScrambleConfig cfgA;
    cfgA.rounds = 2;
    cfgA.seed = radiantwave::scramble::kDefaultSeed;
    ScrambleConfig cfgB = cfgA;
    cfgB.seed = cfgA.seed + 1;

    ASSERT_TRUE(scramble(img.data(), a.data(), 8, 8, cfgA));
    ASSERT_TRUE(scramble(img.data(), b.data(), 8, 8, cfgB));
    EXPECT_NE(a, b);
}

TEST(FibonacciScramblerTest, DifferentRoundsChangeOutput)
{
    const auto img = uniqueImage(8, 8);
    const int n = 8;
    std::vector<uint8_t> a(static_cast<size_t>(n) * n * 4);
    std::vector<uint8_t> b(static_cast<size_t>(n) * n * 4);

    ScrambleConfig cfgA;
    cfgA.rounds = 1;
    ScrambleConfig cfgB;
    cfgB.rounds = 2;

    ASSERT_TRUE(scramble(img.data(), a.data(), 8, 8, cfgA));
    ASSERT_TRUE(scramble(img.data(), b.data(), 8, 8, cfgB));
    EXPECT_NE(a, b);
}

TEST(FibonacciScramblerTest, AlphaChannelIsPreservedByRoundTrip)
{
    // Pixel values exercise all 32 bits, alpha included.
    EXPECT_TRUE(roundTrips(11, 7, 5, 12345));
}

TEST(FibonacciScramblerTest, PaddingIsTransparentAndRoundTrips)
{
    // Opaque white source: after scrambling, every canvas pixel must be either
    // white (a real source pixel) or transparent black (padding), and the
    // unscramble must restore pure white.
    const int width = 5;
    const int height = 7;
    const int n = canvasSize(width, height);
    std::vector<uint8_t> img(static_cast<size_t>(width) * height * 4, 255);
    std::vector<uint8_t> scrambled(static_cast<size_t>(n) * n * 4);
    std::vector<uint8_t> recovered(static_cast<size_t>(width) * height * 4);

    ScrambleConfig cfg;
    cfg.rounds = 4;
    ASSERT_TRUE(scramble(img.data(), scrambled.data(), width, height, cfg));
    for (size_t i = 0; i < scrambled.size(); i += 4)
    {
        const bool isWhite = scrambled[i] == 255 && scrambled[i + 1] == 255 &&
                             scrambled[i + 2] == 255 && scrambled[i + 3] == 255;
        const bool isPadding = scrambled[i] == 0 && scrambled[i + 1] == 0 &&
                               scrambled[i + 2] == 0 && scrambled[i + 3] == 0;
        EXPECT_TRUE(isWhite || isPadding)
            << "canvas pixel " << (i / 4) << " is neither white nor padding";
    }
    ASSERT_TRUE(unscramble(scrambled.data(), recovered.data(), width, height, cfg));
    EXPECT_EQ(recovered, img);
}

TEST(FibonacciScramblerTest, RoundsBelowOneClampToOne)
{
    const auto img = uniqueImage(7, 5);
    const int n = 7;
    std::vector<uint8_t> clamped(static_cast<size_t>(n) * n * 4);
    std::vector<uint8_t> single(static_cast<size_t>(n) * n * 4);

    ScrambleConfig cfgClamped;
    cfgClamped.rounds = 0;
    ScrambleConfig cfgSingle;
    cfgSingle.rounds = 1;

    ASSERT_TRUE(scramble(img.data(), clamped.data(), 7, 5, cfgClamped));
    ASSERT_TRUE(scramble(img.data(), single.data(), 7, 5, cfgSingle));
    EXPECT_EQ(clamped, single);
    EXPECT_TRUE(roundTrips(7, 5, 0, radiantwave::scramble::kDefaultSeed));
}

TEST(FibonacciScramblerTest, RejectsInvalidArguments)
{
    std::vector<uint8_t> dummy(16 * 16 * 4, 0);
    ScrambleConfig cfg;

    EXPECT_FALSE(scramble(nullptr, dummy.data(), 4, 4, cfg));
    EXPECT_FALSE(scramble(dummy.data(), nullptr, 4, 4, cfg));
    EXPECT_FALSE(scramble(dummy.data(), dummy.data(), 0, 4, cfg));
    EXPECT_FALSE(scramble(dummy.data(), dummy.data(), 4, -1, cfg));
    EXPECT_FALSE(unscramble(nullptr, dummy.data(), 4, 4, cfg));
    EXPECT_FALSE(unscramble(dummy.data(), nullptr, 4, 4, cfg));
    EXPECT_FALSE(unscramble(dummy.data(), dummy.data(), 0, 4, cfg));
}