// pixmask nanobind Python bindings.
// See architecture/CPP_REFERENCE.md section 4 for patterns.
#include <cstring>
#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
#include <nanobind/stl/string.h>

#include <cstddef>
#include <cstdint>
#include <stdexcept>

#include "pixmask/pipeline.h"
#include "pixmask/types.h"

namespace nb = nanobind;
using namespace nb::literals;

// C-contiguous uint8 HWC array on CPU — zero-copy input from numpy.
using InputArray = nb::ndarray<uint8_t,
                               nb::shape<-1, -1, -1>,
                               nb::device::cpu,
                               nb::c_contig>;

// ---------------------------------------------------------------------------
// sanitize_bytes: accepts raw image bytes (JPEG/PNG), returns HWC uint8 array
// ---------------------------------------------------------------------------
static nb::ndarray<nb::numpy, uint8_t, nb::ndim<3>>
sanitize_bytes(nb::bytes input,
               uint8_t bit_depth,
               uint8_t median_radius,
               uint8_t jpeg_quality_lo,
               uint8_t jpeg_quality_hi) {
    const auto* data = reinterpret_cast<const uint8_t*>(input.c_str());
    const size_t len = input.size();

    pixmask::SanitizeOptions opts;
    opts.bit_depth       = bit_depth;
    opts.median_radius   = median_radius;
    opts.jpeg_quality_lo = jpeg_quality_lo;
    opts.jpeg_quality_hi = jpeg_quality_hi;

    // Pipeline must outlive the memcpy below — result.image.data points
    // into the Pipeline's arena, which is freed when Pipeline is destroyed.
    pixmask::Pipeline pipeline(opts);
    pixmask::SanitizeResult result;
    {
        nb::gil_scoped_release release;
        result = pipeline.sanitize(data, len);
    }

    if (!result.success) {
        throw std::runtime_error(
            result.error_message ? result.error_message : "sanitize failed");
    }

    const auto& img = result.image;
    const size_t H = img.height;
    const size_t W = img.width;
    const size_t C = img.channels;
    const size_t n = H * W * C;

    // Copy out of arena into heap-owned buffer before pipeline/arena is destroyed.
    auto* out_buf = new uint8_t[n];
    for (size_t y = 0; y < H; ++y) {
        const uint8_t* src_row = img.data + y * img.stride;
        uint8_t* dst_row = out_buf + y * W * C;
        std::memcpy(dst_row, src_row, W * C);
    }

    nb::capsule owner(out_buf, [](void* p) noexcept {
        delete[] static_cast<uint8_t*>(p);
    });

    return nb::ndarray<nb::numpy, uint8_t, nb::ndim<3>>(
        out_buf, {H, W, C}, owner);
}

// ---------------------------------------------------------------------------
// sanitize_array: accepts HWC uint8 numpy array, returns HWC uint8 array
// ---------------------------------------------------------------------------
static nb::ndarray<nb::numpy, uint8_t, nb::ndim<3>>
sanitize_array(InputArray img,
               uint8_t bit_depth,
               uint8_t median_radius,
               uint8_t jpeg_quality_lo,
               uint8_t jpeg_quality_hi) {
    const size_t H = img.shape(0);
    const size_t W = img.shape(1);
    const size_t C = img.shape(2);

    if (C != 1 && C != 3 && C != 4) {
        throw nb::value_error("channels must be 1, 3, or 4");
    }

    pixmask::SanitizeOptions opts;
    opts.bit_depth       = bit_depth;
    opts.median_radius   = median_radius;
    opts.jpeg_quality_lo = jpeg_quality_lo;
    opts.jpeg_quality_hi = jpeg_quality_hi;

    // Build ImageView over the numpy buffer — zero-copy.
    pixmask::ImageView view;
    view.data     = img.data();
    view.width    = static_cast<uint32_t>(W);
    view.height   = static_cast<uint32_t>(H);
    view.channels = static_cast<uint32_t>(C);
    view.stride   = static_cast<uint32_t>(W * C);

    // We need to pass raw bytes to the C++ pipeline.  The pipeline expects
    // encoded bytes (JPEG/PNG), so for array input we encode to PNG first,
    // then let the pipeline decode+sanitize.  However, the architecture also
    // supports an in-memory pixel path via Pipeline::sanitize on raw bytes.
    //
    // For v0.1 the canonical path is bytes-in → bytes-out through the
    // full pipeline.  Array input is handled at the Python layer by encoding
    // to PNG bytes first (see python/pixmask/__init__.py).  This C++ binding
    // directly sanitizes raw encoded bytes.
    //
    // TODO(v0.2): Add a pixel-buffer entry point that skips decode/validate.

    // For now, this binding is not exposed to Python — the Python layer
    // handles array→bytes conversion.  Keep the implementation for future use.
    (void)view;

    throw std::runtime_error(
        "Direct array sanitization not yet implemented in C++ pipeline. "
        "Use sanitize_bytes with encoded image data.");
}

// ---------------------------------------------------------------------------
// Module definition
// ---------------------------------------------------------------------------
NB_MODULE(pixmask_ext, m) {
    m.doc() = "pixmask C++ extension: fast image sanitization for multimodal LLMs";

    m.def(
        "sanitize_bytes",
        &sanitize_bytes,
        "data"_a,
        "bit_depth"_a       = 5,
        "median_radius"_a   = 1,
        "jpeg_quality_lo"_a = 70,
        "jpeg_quality_hi"_a = 85,
        "Sanitize raw JPEG/PNG bytes. Returns uint8 HxWxC numpy array.\n"
        "GIL is released during C++ processing."
    );

    m.def(
        "sanitize_array",
        &sanitize_array,
        "img"_a,
        "bit_depth"_a       = 5,
        "median_radius"_a   = 1,
        "jpeg_quality_lo"_a = 70,
        "jpeg_quality_hi"_a = 85,
        "Sanitize a uint8 HxWxC numpy array. Returns new array.\n"
        "GIL is released during C++ processing."
    );
}
