#include "fgc/JpegXlEncoder.h"

#include <doctest/doctest.h>
#include "jxl_test_decode.h"

#include <chrono>
#include <cmath>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <iterator>
#include <random>
#include <string>
#include <vector>

using fgc::JpegXlEncoder;
using fgc_test::decode;
namespace fs = std::filesystem;

namespace {

struct TempDir {
    fs::path path;
    TempDir() {
        path = fs::temp_directory_path() /
               ("fgc_jxl_" + std::to_string(std::chrono::steady_clock::now().time_since_epoch().count()));
        fs::create_directories(path);
    }
    ~TempDir() { fs::remove_all(path); }
    std::string file(const std::string& name) const { return (path / name).string(); }
};

// A test image: smooth gradients (what lossy coding keeps well) plus noise
// and a few hard edges, in `maxv` levels.
std::vector<uint16_t> image(int w, int h, int ch, int maxv, unsigned seed = 1) {
    std::mt19937               rng(seed);
    std::normal_distribution<> noise(0.0, maxv / 200.0);
    std::vector<uint16_t>      v(static_cast<size_t>(w) * h * ch);
    for (int y = 0; y < h; ++y)
        for (int x = 0; x < w; ++x)
            for (int c = 0; c < ch; ++c) {
                double base = maxv * (0.15 + 0.7 * (x + y * (c + 1)) / double(w + h * ch));
                if ((x / 16 + y / 16) % 7 == 0) base = maxv * 0.9;  // edges
                double val = std::round(base + noise(rng));
                v[(static_cast<size_t>(y) * w + x) * ch + c] =
                    static_cast<uint16_t>(std::min<double>(maxv, std::max(0.0, val)));
            }
    // The extremes must survive too.
    v.front() = 0;
    v.back()  = static_cast<uint16_t>(maxv);
    return v;
}

std::vector<uint8_t> to8(const std::vector<uint16_t>& v) {
    return std::vector<uint8_t>(v.begin(), v.end());
}

double meanAbsError(const std::vector<uint16_t>& a, const std::vector<uint16_t>& b) {
    double s = 0;
    for (size_t i = 0; i < a.size(); ++i) s += std::abs(double(a[i]) - double(b[i]));
    return s / a.size();
}

}  // namespace

TEST_CASE("8-bit lossless decodes exactly (known-issues #20)") {
    TempDir t;
    for (int ch : {1, 3}) {
        CAPTURE(ch);
        const auto src = image(96, 64, ch, 255);
        const auto b   = to8(src);
        REQUIRE(JpegXlEncoder::encodeToFile(t.file("ll8.jxl"), b.data(), 96, 64, ch, 0.0f, 1, 2));
        const auto d = decode(t.file("ll8.jxl"), ch);
        CHECK(d.bits == 8);
        CHECK(d.original_profile);
        CHECK(d.px == src);
    }
}

TEST_CASE("12-bit lossless decodes exactly, extremes included") {
    TempDir t;
    for (int ch : {1, 3}) {
        CAPTURE(ch);
        const auto src = image(96, 64, ch, 4095, 7);
        REQUIRE(JpegXlEncoder::encodeToFile(t.file("ll12.jxl"), src.data(), 96, 64, ch, 12, 0.0f,
                                            1, 2));
        const auto d = decode(t.file("ll12.jxl"), ch);
        CHECK(d.bits == 12);
        CHECK(d.channels == ch);
        CHECK(d.px == src);
    }
}

TEST_CASE("16-bit single-channel lossless decodes exactly (thermal counts)") {
    TempDir t;
    // Counts around room temperature, plus the full range's ends.
    auto src = image(96, 64, 1, 65535, 11);
    for (size_t i = 0; i < src.size(); i += 3) src[i] = static_cast<uint16_t>(28000 + i % 2500);
    src.front() = 0;
    src.back()  = 65535;
    REQUIRE(JpegXlEncoder::encodeToFile(t.file("ll16.jxl"), src.data(), 96, 64, 1, 16, 0.0f, 1, 2));
    const auto d = decode(t.file("ll16.jxl"), 1);
    CHECK(d.bits == 16);
    CHECK(d.channels == 1);
    CHECK(d.original_profile);
    CHECK(d.px == src);
}

TEST_CASE("lossy keeps the bit depth and stays close") {
    TempDir t;
    SUBCASE("12-bit colour") {
        const auto src = image(128, 96, 3, 4095, 3);
        REQUIRE(JpegXlEncoder::encodeToFile(t.file("l12.jxl"), src.data(), 128, 96, 3, 12, 1.0f,
                                            3, 2));
        const auto d = decode(t.file("l12.jxl"), 3);
        CHECK(d.bits == 12);
        CHECK_FALSE(d.original_profile);  // lossy goes through XYB, as before
        REQUIRE(d.px.size() == src.size());
        CHECK(meanAbsError(d.px, src) < 4095 * 0.02);
    }
    SUBCASE("8-bit colour, as the program has always saved") {
        const auto src = image(128, 96, 3, 255, 3);
        const auto b   = to8(src);
        REQUIRE(JpegXlEncoder::encodeToFile(t.file("l8.jxl"), b.data(), 128, 96, 3, 0.8f, 4));
        const auto d = decode(t.file("l8.jxl"), 3);
        CHECK(d.bits == 8);
        REQUIRE(d.px.size() == src.size());
        CHECK(meanAbsError(d.px, src) < 255 * 0.02);
    }
}

TEST_CASE("the thread count changes speed, not the file") {
    TempDir    t;
    const auto src = image(160, 120, 3, 4095, 5);
    REQUIRE(JpegXlEncoder::encodeToFile(t.file("t0.jxl"), src.data(), 160, 120, 3, 12, 1.0f, 3, 0));
    REQUIRE(JpegXlEncoder::encodeToFile(t.file("t4.jxl"), src.data(), 160, 120, 3, 12, 1.0f, 3, 4));
    std::ifstream a(t.file("t0.jxl"), std::ios::binary), b(t.file("t4.jxl"), std::ios::binary);
    std::vector<char> fa{std::istreambuf_iterator<char>(a), {}}, fb{std::istreambuf_iterator<char>(b), {}};
    CHECK(!fa.empty());
    CHECK(fa == fb);
}

TEST_CASE("bad input is refused and leaves no file") {
    TempDir    t;
    const auto src = image(16, 16, 3, 255);
    const auto b   = to8(src);
    const std::string p = t.file("bad.jxl");
    CHECK_FALSE(JpegXlEncoder::encodeToFile(p, b.data(), 16, 16, 2, 8, 1.0f, 3, 1));    // 2 channels
    CHECK_FALSE(JpegXlEncoder::encodeToFile(p, b.data(), 16, 16, 3, 7, 1.0f, 3, 1));    // 7 bits
    CHECK_FALSE(JpegXlEncoder::encodeToFile(p, b.data(), 16, 16, 3, 17, 1.0f, 3, 1));   // 17 bits
    CHECK_FALSE(JpegXlEncoder::encodeToFile(p, b.data(), 0, 16, 3, 8, 1.0f, 3, 1));     // no width
    CHECK_FALSE(JpegXlEncoder::encodeToFile(p, b.data(), 16, 16, 3, 8, -1.0f, 3, 1));   // distance
    CHECK_FALSE(JpegXlEncoder::encodeToFile(p, nullptr, 16, 16, 3, 8, 1.0f, 3, 1));     // no image
    CHECK_FALSE(fs::exists(p));
}