// This file is part of the 64k demo project.
// Tests for the ANS (rANS) entropy coder in util/ans.h.
// Encoder is enabled in tests via ANS_ENABLE_ENCODER.

#include "util/ans.h"

#include <cassert>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <random>
#include <string>
#include <vector>

namespace {

bool RoundtripCheck(const std::vector<uint8_t>& input,
                    const uint32_t* initial_counts, const char* label) {
  std::vector<uint8_t> compressed;
  if (!ans::Encode(input.data(), input.size(), &compressed, initial_counts)) {
    fprintf(stderr, "[%s] Encode failed\n", label);
    return false;
  }
  std::vector<uint8_t> decoded(input.size());
  size_t decoded_size = 0;
  if (!ans::Decode(compressed.data(), compressed.size(), decoded.data(),
                   decoded.size(), &decoded_size, initial_counts)) {
    fprintf(stderr, "[%s] Decode failed\n", label);
    return false;
  }
  if (decoded_size != input.size() ||
      std::memcmp(decoded.data(), input.data(), input.size()) != 0) {
    fprintf(stderr, "[%s] payload mismatch\n", label);
    return false;
  }
  const double ratio =
      input.empty() ? 0.0 : (double)compressed.size() / (double)input.size();
  fprintf(stderr, "[%s] OK: %zu -> %zu bytes (%.3f x)\n", label, input.size(),
          compressed.size(), ratio);
  return true;
}

// Covers: empty / single byte / single-symbol run / all-zeros / random uniform
// / random skewed / repeated ASCII. Each is a roundtrip with default (uniform)
// initial counts.
void TestRoundtripVariants() {
  std::mt19937 rng_uniform(12345);
  std::vector<uint8_t> random_uniform(64 * 1024);
  for (auto& b : random_uniform)
    b = (uint8_t)(rng_uniform() & 0xff);

  std::mt19937 rng_skewed(67890);
  std::vector<uint8_t> random_skewed(32 * 1024);
  for (auto& b : random_skewed) {
    // 90% byte 'A', 10% other ASCII letters.
    const uint32_t r = rng_skewed();
    b = (r % 10 == 0) ? (uint8_t)('B' + (r >> 8) % 25) : (uint8_t)'A';
  }

  const char* ascii =
      "@group(0) @binding(0) var smplr: sampler;\n"
      "@group(0) @binding(1) var tex: texture_2d<f32>;\n"
      "@fragment fn fs_main(@builtin(position) p: vec4f) -> @location(0) "
      "vec4f {\n"
      "  let uv = p.xy / vec2f(1280.0, 720.0);\n"
      "  return textureSample(tex, smplr, uv);\n"
      "}\n";
  std::vector<uint8_t> ascii_block;
  for (int i = 0; i < 50; ++i) { // cross chunk boundary
    ascii_block.insert(ascii_block.end(), ascii, ascii + std::strlen(ascii));
  }

  struct Case {
    const char* label;
    std::vector<uint8_t> data;
  };
  const Case cases[] = {
      {"empty", {}},
      {"single-byte", {0x42}},
      {"single-symbol-run", std::vector<uint8_t>(4096, 'A')},
      {"all-zeros", std::vector<uint8_t>(8192, 0)},
      {"random-uniform", random_uniform},
      {"random-skewed", random_skewed},
      {"ascii-shader-text", ascii_block},
  };
  for (const Case& c : cases) {
    assert(RoundtripCheck(c.data, nullptr, c.label));
  }
}

// Seeded initial counts should compress a matching corpus at least as well
// as uniform init, and roundtrip identically.
void TestSeededInitialCounts() {
  std::vector<uint8_t> corpus;
  const char* sample = "fn main() { let x = vec3f(1.0, 2.0, 3.0); }\n";
  for (int i = 0; i < 200; ++i) {
    corpus.insert(corpus.end(), sample, sample + std::strlen(sample));
  }
  uint32_t hist[256] = {};
  ans::Histogram(corpus.data(), corpus.size(), hist);

  std::vector<uint8_t> seeded_out, uniform_out;
  assert(ans::Encode(corpus.data(), corpus.size(), &seeded_out, hist));
  assert(ans::Encode(corpus.data(), corpus.size(), &uniform_out, nullptr));

  std::vector<uint8_t> decoded(corpus.size());
  size_t decoded_size = 0;
  assert(ans::Decode(seeded_out.data(), seeded_out.size(), decoded.data(),
                     decoded.size(), &decoded_size, hist));
  assert(decoded_size == corpus.size());
  assert(std::memcmp(decoded.data(), corpus.data(), corpus.size()) == 0);

  fprintf(stderr, "[seeded-vs-uniform] seeded=%zu uniform=%zu raw=%zu\n",
          seeded_out.size(), uniform_out.size(), corpus.size());
  assert(seeded_out.size() <= uniform_out.size());
}

// The chunk-end state check must reject: a mismatched model (decoder uses a
// different initial distribution), a single bit-flip in the payload, and a
// truncated stream.
void TestRejection() {
  std::mt19937 rng(1);

  // 1) Mismatched models.
  {
    std::vector<uint8_t> v(4096);
    for (auto& b : v)
      b = (uint8_t)('a' + (rng() % 26));
    uint32_t hist[256] = {};
    ans::Histogram(v.data(), v.size(), hist);

    std::vector<uint8_t> encoded;
    assert(ans::Encode(v.data(), v.size(), &encoded, hist));

    std::vector<uint8_t> decoded(v.size());
    size_t decoded_size = 0;
    assert(!ans::Decode(encoded.data(), encoded.size(), decoded.data(),
                        decoded.size(), &decoded_size, nullptr));
    fprintf(stderr, "[rejection] mismatched-counts OK\n");
  }

  // 2) Corruption.
  {
    std::vector<uint8_t> v(2048);
    for (auto& b : v)
      b = (uint8_t)(rng() & 0xff);
    std::vector<uint8_t> encoded;
    assert(ans::Encode(v.data(), v.size(), &encoded, nullptr));
    encoded[encoded.size() / 2] ^= 0x55; // flip a payload byte

    std::vector<uint8_t> decoded(v.size());
    size_t decoded_size = 0;
    assert(!ans::Decode(encoded.data(), encoded.size(), decoded.data(),
                        decoded.size(), &decoded_size, nullptr));
    fprintf(stderr, "[rejection] corruption OK\n");
  }

  // 3) Truncation.
  {
    std::vector<uint8_t> v(4096);
    for (size_t i = 0; i < v.size(); ++i)
      v[i] = (uint8_t)i;
    std::vector<uint8_t> encoded;
    assert(ans::Encode(v.data(), v.size(), &encoded, nullptr));
    encoded.resize(encoded.size() - 8);

    std::vector<uint8_t> decoded(v.size());
    size_t decoded_size = 0;
    assert(!ans::Decode(encoded.data(), encoded.size(), decoded.data(),
                        decoded.size(), &decoded_size, nullptr));
    fprintf(stderr, "[rejection] truncation OK\n");
  }
}

void TestPeekSize() {
  std::vector<uint8_t> v(1234, 'Q');
  std::vector<uint8_t> encoded;
  assert(ans::Encode(v.data(), v.size(), &encoded, nullptr));
  assert(ans::PeekUncompressedSize(encoded.data(), encoded.size()) == v.size());
}

} // namespace

int main() {
  TestRoundtripVariants();
  TestSeededInitialCounts();
  TestRejection();
  TestPeekSize();
  fprintf(stderr, "All ANS tests passed.\n");
  return 0;
}