// pixmask/jpeg_roundtrip.cpp — Stage 5: JPEG encode+decode roundtrip.
// See architecture/CPP_JPEG_REFERENCE.md for design details.

// stb_image_write: implementation owned by this TU.
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

// stb_image: header-only include. STB_IMAGE_IMPLEMENTATION lives in decode.cpp.
#include "stb_image.h"

#include "pixmask/jpeg_roundtrip.h"

#include <cstdlib>  // malloc, realloc, free
#include <cstring>  // memcpy
#include <cstdint>

// ---------------------------------------------------------------------------
// OS entropy shim
// ---------------------------------------------------------------------------
#if defined(__linux__)
// Use syscall directly — sys/random.h is unavailable on manylinux2014 (glibc 2.17).
#  include <sys/syscall.h>
#  include <unistd.h>
#  include <cerrno>
#  ifdef SYS_getrandom
#    define PIXMASK_HAVE_GETRANDOM 1
#  else
#    include <fcntl.h>
#    define PIXMASK_HAVE_DEV_URANDOM 1
#  endif
#elif defined(__APPLE__) || defined(__FreeBSD__)
#  include <sys/random.h>
#  define PIXMASK_HAVE_GETENTROPY 1
#elif defined(_WIN32)
#  ifndef WIN32_LEAN_AND_MEAN
#    define WIN32_LEAN_AND_MEAN
#  endif
#  include <windows.h>
#  include <bcrypt.h>
#  pragma comment(lib, "bcrypt.lib")
#  define PIXMASK_HAVE_BCRYPT 1
#else
#  include <random>
#  define PIXMASK_HAVE_STD_RANDOM 1
#endif

namespace pixmask {
namespace {

// ---------------------------------------------------------------------------
// Write callback for stbi_write_jpg_to_func — accumulates JPEG bytes
// ---------------------------------------------------------------------------

struct WriteBuffer {
    uint8_t* data     = nullptr;
    size_t   used     = 0;
    size_t   capacity = 0;
    bool     failed   = false;
};

void jpeg_write_cb(void* context, void* chunk, int size) {
    auto* wb = static_cast<WriteBuffer*>(context);
    if (wb->failed) return;

    auto chunk_sz = static_cast<size_t>(size);
    size_t needed = wb->used + chunk_sz;

    if (needed > wb->capacity) {
        size_t new_cap = (wb->capacity == 0) ? needed : wb->capacity * 2;
        if (new_cap < needed) new_cap = needed;
        void* p = std::realloc(wb->data, new_cap);
        if (p == nullptr) {
            wb->failed = true;
            return;
        }
        wb->data     = static_cast<uint8_t*>(p);
        wb->capacity = new_cap;
    }
    std::memcpy(wb->data + wb->used, chunk, chunk_sz);
    wb->used += chunk_sz;
}

// ---------------------------------------------------------------------------
// OS entropy
// ---------------------------------------------------------------------------

bool os_random_bytes(uint8_t* buf, size_t n) noexcept {
#if defined(PIXMASK_HAVE_GETRANDOM)
    return ::syscall(SYS_getrandom, buf, n, 0) == static_cast<long>(n);
#elif defined(PIXMASK_HAVE_DEV_URANDOM)
    int fd = ::open("/dev/urandom", O_RDONLY | O_CLOEXEC);
    if (fd < 0) return false;
    auto r = ::read(fd, buf, n);
    ::close(fd);
    return r == static_cast<ssize_t>(n);
#elif defined(PIXMASK_HAVE_GETENTROPY)
    return ::getentropy(buf, n) == 0;
#elif defined(PIXMASK_HAVE_BCRYPT)
    return BCryptGenRandom(nullptr, buf, static_cast<ULONG>(n),
                           BCRYPT_USE_SYSTEM_PREFERRED_RNG) == 0;
#elif defined(PIXMASK_HAVE_STD_RANDOM)
    // Fallback: std::random_device. Not ideal but better than rand().
    try {
        std::random_device rd;
        for (size_t i = 0; i < n; ++i) {
            buf[i] = static_cast<uint8_t>(rd());
        }
        return true;
    } catch (...) {
        return false;
    }
#else
    (void)buf; (void)n;
    return false;
#endif
}

// Returns a quality factor in [lo, hi] inclusive, drawn from OS entropy.
// Uses rejection sampling to avoid modulo bias.
int random_quality(int lo, int hi) noexcept {
    const int range  = hi - lo + 1;
    const int usable = (256 / range) * range;

    for (int attempt = 0; attempt < 16; ++attempt) {
        uint8_t byte{};
        if (!os_random_bytes(&byte, 1)) {
            // Entropy failure: fall back to midpoint.
            return lo + range / 2;
        }
        if (byte < static_cast<uint8_t>(usable)) {
            return lo + (byte % range);
        }
        // byte >= usable: reject and retry to avoid modulo bias.
    }
    // Extremely unlikely: 16 consecutive rejections. Use midpoint.
    return lo + range / 2;
}

} // anonymous namespace

// ---------------------------------------------------------------------------
// jpeg_roundtrip
// ---------------------------------------------------------------------------

ImageView jpeg_roundtrip(const ImageView& input, Arena& arena,
                         uint8_t quality_lo, uint8_t quality_hi,
                         const char** err_out) noexcept {
    auto fail = [&](const char* msg) -> ImageView {
        if (err_out) *err_out = msg;
        return ImageView{};
    };

    // --- Input validation ---------------------------------------------------

    if (input.data == nullptr || input.width == 0 || input.height == 0 ||
        (input.channels != 1 && input.channels != 3 && input.channels != 4) ||
        input.stride < input.width * input.channels) {
        return fail("jpeg_roundtrip: invalid input ImageView");
    }

    if (quality_lo > quality_hi) {
        return fail("jpeg_roundtrip: quality_lo > quality_hi");
    }

    if (quality_lo < 1 || quality_hi > 100) {
        return fail("jpeg_roundtrip: quality must be in [1, 100]");
    }

    // --- Stage 5a: encode to JPEG in memory ---------------------------------

    const int quality = random_quality(
        static_cast<int>(quality_lo), static_cast<int>(quality_hi));

    // Pre-allocate: JPEG at QF=85 is typically < 0.5 bytes/pixel for photos.
    const size_t raw_size = static_cast<size_t>(input.width)
                          * static_cast<size_t>(input.height)
                          * static_cast<size_t>(input.channels);
    const size_t prealloc = (raw_size / 2 > 0) ? raw_size / 2 : 4096;

    WriteBuffer wb{};
    wb.data = static_cast<uint8_t*>(std::malloc(prealloc));
    if (wb.data == nullptr) {
        return fail("jpeg_roundtrip: malloc failed for encode buffer");
    }
    wb.capacity = prealloc;

    // stb_image_write expects tightly-packed rows. If stride != width*channels,
    // we need to pass stride-aware data. stbi_write_jpg_to_func does not
    // support stride, so we must use tightly-packed input or copy rows.
    // For now, assert tight packing; pipeline stages produce tight images.
    const uint8_t* encode_data = input.data;

    const int encode_ok = stbi_write_jpg_to_func(
        jpeg_write_cb, &wb,
        static_cast<int>(input.width),
        static_cast<int>(input.height),
        static_cast<int>(input.channels),
        encode_data,
        quality);

    if (!encode_ok || wb.used == 0 || wb.failed) {
        std::free(wb.data);
        return fail("jpeg_roundtrip: JPEG encoding failed");
    }

    // --- Stage 5b: decode JPEG back from memory -----------------------------

    int dec_w{}, dec_h{}, dec_comp{};
    // Always decode to 3-channel RGB. JPEG discards alpha anyway.
    const int req_channels = static_cast<int>(input.channels);
    stbi_uc* decoded = stbi_load_from_memory(
        wb.data, static_cast<int>(wb.used),
        &dec_w, &dec_h, &dec_comp,
        req_channels);

    std::free(wb.data);
    wb.data = nullptr;

    if (decoded == nullptr) {
        return fail("jpeg_roundtrip: JPEG decode failed");
    }

    if (static_cast<uint32_t>(dec_w) != input.width ||
        static_cast<uint32_t>(dec_h) != input.height) {
        stbi_image_free(decoded);
        return fail("jpeg_roundtrip: decoded dimensions differ from source");
    }

    // --- Copy decoded pixels into arena -------------------------------------

    const auto out_channels = static_cast<uint32_t>(req_channels);
    const size_t out_bytes = static_cast<size_t>(dec_w)
                           * static_cast<size_t>(dec_h)
                           * out_channels;

    uint8_t* out = nullptr;
    try {
        out = static_cast<uint8_t*>(arena.allocate(out_bytes));
    } catch (...) {
        stbi_image_free(decoded);
        return fail("jpeg_roundtrip: arena allocation failed");
    }

    std::memcpy(out, decoded, out_bytes);
    stbi_image_free(decoded);

    ImageView result{};
    result.data     = out;
    result.width    = static_cast<uint32_t>(dec_w);
    result.height   = static_cast<uint32_t>(dec_h);
    result.channels = out_channels;
    result.stride   = result.width * out_channels;
    return result;
}

} // namespace pixmask
