#include "gpu_sdf.h"
#include "gpu_offset.h"
#include "cpu_offset.h"
#include "vectorizer.h"
#include <cuda_runtime.h>
#include <iostream>
#include <vector>
#include <cmath>
#include <chrono>
#include <cstring>
#include <algorithm>
#include <thread>
#include <map>
#include <set>

#define CUDA_CHECK(call) do { \
    cudaError_t err = call; \
    if (err != cudaSuccess) { \
        std::cerr << "[CUDA Error] " << cudaGetErrorString(err) \
                  << " at " << __FILE__ << ":" << __LINE__ << std::endl; \
        return nullptr; \
    } \
} while (0)

#define CUDA_CHECK_INT(call) do { \
    cudaError_t err = call; \
    if (err != cudaSuccess) { \
        std::cerr << "[CUDA Error] " << cudaGetErrorString(err) \
                  << " at " << __FILE__ << ":" << __LINE__ << std::endl; \
        return GPU_OFFSET_ERR_GENERIC; \
    } \
} while (0)

// Seed packing helpers
__device__ __inline__ int pack_seed_sdf(short x, short y) {
    return ((int)(unsigned short)y << 16) | (int)(unsigned short)x;
}
__device__ __inline__ short get_seed_x_sdf(int packed) {
    return (short)(packed & 0xFFFF);
}
__device__ __inline__ short get_seed_y_sdf(int packed) {
    return (short)((packed >> 16) & 0xFFFF);
}

// -----------------------------------------------------------------------------
// Internal Implementation of GpuSdfStack
// -----------------------------------------------------------------------------
struct GpuSdfStack_T {
    GpuSdfMetadata meta;
    int32_t* d_z_um;     // GPU pointer to array of layer heights [num_layers] in micrometers
    float* d_sdf;        // GPU pointer to 3D continuous distance field [num_layers * height * width] in mm
    cudaStream_t stream;
    bool owns_stream;
};

struct LayerSamplingInfo {
    int l0_a;
    int l1_a;
    float t_a; // < 0 if not present in A
    int l0_b;
    int l1_b;
    float t_b; // < 0 if not present in B
};

// -----------------------------------------------------------------------------
// CUDA Kernels for JFA Full Signed Distance Generation
// -----------------------------------------------------------------------------
__global__ void jfa_init_interior_exterior_kernel(
    const uint8_t* __restrict__ masks,
    int* __restrict__ seeds_ext,
    int* __restrict__ seeds_int,
    int width, int height, int layers
) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    int z = blockIdx.z;
    if (x >= width || y >= height || z >= layers) return;

    size_t idx = ((size_t)z * height + y) * width + x;
    uint8_t m = masks[idx];

    // seeds_ext: background pixels (0) are seeds to propagate into foreground (1)
    if (m == 0) {
        seeds_ext[idx] = pack_seed_sdf((short)x, (short)y);
    } else {
        seeds_ext[idx] = -1;
    }

    // seeds_int: foreground pixels (1) are seeds to propagate into background (0)
    if (m > 0) {
        seeds_int[idx] = pack_seed_sdf((short)x, (short)y);
    } else {
        seeds_int[idx] = -1;
    }
}

__global__ void jfa_step_dual_kernel(
    const int* __restrict__ in_seeds_ext,
    int* __restrict__ out_seeds_ext,
    const int* __restrict__ in_seeds_int,
    int* __restrict__ out_seeds_int,
    int width, int height, int layers,
    int step
) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    int z = blockIdx.z;
    if (x >= width || y >= height || z >= layers) return;

    size_t layer_offset = (size_t)z * height * width;
    size_t idx = layer_offset + y * width + x;

    // 1. Exterior seeds
    int best_seed_ext = in_seeds_ext[idx];
    int best_dist2_ext = 0x7FFFFFFF;
    if (best_seed_ext != -1) {
        int dx = (int)get_seed_x_sdf(best_seed_ext) - x;
        int dy = (int)get_seed_y_sdf(best_seed_ext) - y;
        best_dist2_ext = dx * dx + dy * dy;
    }

    // 2. Interior seeds
    int best_seed_int = in_seeds_int[idx];
    int best_dist2_int = 0x7FFFFFFF;
    if (best_seed_int != -1) {
        int dx = (int)get_seed_x_sdf(best_seed_int) - x;
        int dy = (int)get_seed_y_sdf(best_seed_int) - y;
        best_dist2_int = dx * dx + dy * dy;
    }

    #pragma unroll
    for (int dy = -1; dy <= 1; ++dy) {
        #pragma unroll
        for (int dx = -1; dx <= 1; ++dx) {
            if (dx == 0 && dy == 0) continue;
            int nx = x + dx * step;
            int ny = y + dy * step;
            if ((unsigned int)nx < (unsigned int)width && (unsigned int)ny < (unsigned int)height) {
                size_t n_idx = layer_offset + ny * width + nx;

                int s_ext = in_seeds_ext[n_idx];
                if (s_ext != -1) {
                    int dist_x = (int)get_seed_x_sdf(s_ext) - x;
                    int dist_y = (int)get_seed_y_sdf(s_ext) - y;
                    int dist2 = dist_x * dist_x + dist_y * dist_y;
                    if (dist2 < best_dist2_ext) {
                        best_dist2_ext = dist2;
                        best_seed_ext = s_ext;
                    }
                }

                int s_int = in_seeds_int[n_idx];
                if (s_int != -1) {
                    int dist_x = (int)get_seed_x_sdf(s_int) - x;
                    int dist_y = (int)get_seed_y_sdf(s_int) - y;
                    int dist2 = dist_x * dist_x + dist_y * dist_y;
                    if (dist2 < best_dist2_int) {
                        best_dist2_int = dist2;
                        best_seed_int = s_int;
                    }
                }
            }
        }
    }

    out_seeds_ext[idx] = best_seed_ext;
    out_seeds_int[idx] = best_seed_int;
}

__global__ void fuse_signed_sdf_kernel(
    const uint8_t* __restrict__ masks,
    const int* __restrict__ seeds_ext,
    const int* __restrict__ seeds_int,
    float* __restrict__ out_sdf,
    int width, int height, int layers,
    float pitch_mm
) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    int z = blockIdx.z;
    if (x >= width || y >= height || z >= layers) return;

    size_t idx = ((size_t)z * height + y) * width + x;
    uint8_t m = masks[idx];

    if (m > 0) {
        int s = seeds_ext[idx];
        if (s != -1) {
            float dx = (float)(get_seed_x_sdf(s) - x);
            float dy = (float)(get_seed_y_sdf(s) - y);
            out_sdf[idx] = sqrtf(dx * dx + dy * dy) * pitch_mm;
        } else {
            out_sdf[idx] = 1e5f;
        }
    } else {
        int s = seeds_int[idx];
        if (s != -1) {
            float dx = (float)(get_seed_x_sdf(s) - x);
            float dy = (float)(get_seed_y_sdf(s) - y);
            out_sdf[idx] = -sqrtf(dx * dx + dy * dy) * pitch_mm;
        } else {
            out_sdf[idx] = -1e5f;
        }
    }
}

// -----------------------------------------------------------------------------
// CSG Mathematical Device Functions
// -----------------------------------------------------------------------------
__device__ __inline__ float csg_eval_op(float a, float b, int op, float blend_k) {
    switch (op) {
        case GPU_SDF_OP_UNION:
            return fmaxf(a, b);
        case GPU_SDF_OP_INTERSECTION:
            return fminf(a, b);
        case GPU_SDF_OP_DIFFERENCE:
            return fminf(a, -b);
        case GPU_SDF_OP_SMOOTH_UNION: {
            if (blend_k <= 1e-6f) return fmaxf(a, b);
            float h = fmaxf(blend_k - fabsf(a - b), 0.0f) / blend_k;
            return fmaxf(a, b) + h * h * blend_k * 0.25f;
        }
        case GPU_SDF_OP_SMOOTH_INTERSECT: {
            if (blend_k <= 1e-6f) return fminf(a, b);
            float h = fmaxf(blend_k - fabsf(a - b), 0.0f) / blend_k;
            return fminf(a, b) - h * h * blend_k * 0.25f;
        }
        case GPU_SDF_OP_SMOOTH_DIFF: {
            float neg_b = -b;
            if (blend_k <= 1e-6f) return fminf(a, neg_b);
            float h = fmaxf(blend_k - fabsf(a - neg_b), 0.0f) / blend_k;
            return fminf(a, neg_b) - h * h * blend_k * 0.25f;
        }
        default:
            return a;
    }
}

// -----------------------------------------------------------------------------
// Identical Grid Fast CSG Kernel
// -----------------------------------------------------------------------------
__global__ void gpu_sdf_csg_identical_kernel(
    const float* __restrict__ sdf_a,
    const float* __restrict__ sdf_b,
    float* __restrict__ sdf_out,
    size_t total_elements,
    int op,
    float blend_k
) {
    size_t idx = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_elements) return;

    float a = sdf_a[idx];
    float b = sdf_b[idx];
    sdf_out[idx] = csg_eval_op(a, b, op, blend_k);
}

// -----------------------------------------------------------------------------
// Harmonized Resampling + CSG Kernel
// -----------------------------------------------------------------------------
__device__ __inline__ float sample_stack_sdf(
    const float* __restrict__ sdf,
    int w, int h,
    int x, int y,
    int l0, int l1, float t
) {
    if (t < -0.5f) return -1e5f; // Not present in stack (outside)
    if (x < 0 || x >= w || y < 0 || y >= h) return -1e5f;

    size_t layer_stride = (size_t)w * h;
    if (l0 == l1 || t <= 1e-6f) {
        return sdf[(size_t)l0 * layer_stride + y * w + x];
    } else {
        float v0 = sdf[(size_t)l0 * layer_stride + y * w + x];
        float v1 = sdf[(size_t)l1 * layer_stride + y * w + x];
        return (1.0f - t) * v0 + t * v1;
    }
}

__global__ void gpu_sdf_csg_resample_kernel(
    const float* __restrict__ sdf_a,
    GpuSdfMetadata meta_a,
    const float* __restrict__ sdf_b,
    GpuSdfMetadata meta_b,
    float* __restrict__ sdf_out,
    GpuSdfMetadata meta_out,
    const LayerSamplingInfo* __restrict__ layer_map,
    int op,
    float blend_k
) {
    int x_out = blockIdx.x * blockDim.x + threadIdx.x;
    int y_out = blockIdx.y * blockDim.y + threadIdx.y;
    int l_out = blockIdx.z;

    if (x_out >= meta_out.width || y_out >= meta_out.height || l_out >= meta_out.num_layers) return;

    LayerSamplingInfo info = layer_map[l_out];

    // Compute world X, Y in micrometers
    int32_t world_x_um = meta_out.min_x_um + x_out * meta_out.pitch_um;
    int32_t world_y_um = meta_out.min_y_um + y_out * meta_out.pitch_um;

    // Sample from Stack A
    int x_a = (world_x_um - meta_a.min_x_um) / meta_a.pitch_um;
    int y_a = (world_y_um - meta_a.min_y_um) / meta_a.pitch_um;
    float val_a = sample_stack_sdf(sdf_a, meta_a.width, meta_a.height, x_a, y_a, info.l0_a, info.l1_a, info.t_a);

    // Sample from Stack B
    int x_b = (world_x_um - meta_b.min_x_um) / meta_b.pitch_um;
    int y_b = (world_y_um - meta_b.min_y_um) / meta_b.pitch_um;
    float val_b = sample_stack_sdf(sdf_b, meta_b.width, meta_b.height, x_b, y_b, info.l0_b, info.l1_b, info.t_b);

    float res = csg_eval_op(val_a, val_b, op, blend_k);

    size_t out_idx = ((size_t)l_out * meta_out.height + y_out) * meta_out.width + x_out;
    sdf_out[out_idx] = res;
}

// -----------------------------------------------------------------------------
// Unary GPU Operations (Offset, Shell, Infill)
// -----------------------------------------------------------------------------
__global__ void gpu_sdf_offset_kernel(
    const float* __restrict__ in_sdf,
    float* __restrict__ out_sdf,
    size_t total_elements,
    float delta_mm
) {
    size_t idx = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_elements) return;
    out_sdf[idx] = in_sdf[idx] + delta_mm;
}

__global__ void gpu_sdf_shell_kernel(
    const float* __restrict__ in_sdf,
    float* __restrict__ out_sdf,
    size_t total_elements,
    float thickness_mm
) {
    size_t idx = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_elements) return;

    float d = in_sdf[idx];
    float half_t = thickness_mm * 0.5f;
    // Positive inside shell, negative outside / hollow center
    out_sdf[idx] = half_t - fabsf(d - half_t);
}

__global__ void gpu_sdf_infill_kernel(
    const float* __restrict__ in_sdf,
    float* __restrict__ out_sdf,
    int width, int height, int num_layers,
    int32_t min_x_um, int32_t min_y_um, int32_t pitch_um,
    const int32_t* __restrict__ z_heights_um,
    int pattern,
    float cell_size_mm,
    float density_ratio,
    float wall_offset_mm
) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    int l = blockIdx.z;

    if (x >= width || y >= height || l >= num_layers) return;

    size_t idx = ((size_t)l * height + y) * width + x;
    float dist = in_sdf[idx];

    // Outside boundary or within shell wall offset -> retain wall distance
    if (dist <= wall_offset_mm) {
        out_sdf[idx] = dist;
        return;
    }

    float world_x = (float)(min_x_um + x * pitch_um) * 0.001f;
    float world_y = (float)(min_y_um + y * pitch_um) * 0.001f;
    float world_z = (float)(z_heights_um[l]) * 0.001f;

    float cell_sz = (cell_size_mm > 0.1f) ? cell_size_mm : 5.0f;
    float freq = 6.28318530718f / cell_sz;

    float px = world_x * freq;
    float py = world_y * freq;
    float pz = world_z * freq;

    float field = 0.0f;
    if (pattern == GPU_SDF_INFILL_GYROID) {
        field = sinf(px) * cosf(py) + sinf(py) * cosf(pz) + sinf(pz) * cosf(px);
    } else if (pattern == GPU_SDF_INFILL_SCHWARZ_P) {
        field = cosf(px) + cosf(py) + cosf(pz);
    } else if (pattern == GPU_SDF_INFILL_SCHWARZ_D) {
        field = cosf(px)*cosf(py)*cosf(pz) - sinf(px)*sinf(py)*sinf(pz);
    } else if (pattern == GPU_SDF_INFILL_GRID) {
        float fx = fabsf(fmodf(world_x, cell_sz) - cell_sz * 0.5f);
        float fy = fabsf(fmodf(world_y, cell_sz) - cell_sz * 0.5f);
        field = fminf(fx, fy) / (cell_sz * 0.5f) - 0.5f;
    }

    float threshold = (density_ratio - 0.5f) * 1.5f;
    // Combine lattice with the infill bounding region
    float lattice_sdf = (field >= threshold) ? 1.0f : -1.0f;
    out_sdf[idx] = fminf(dist - wall_offset_mm, lattice_sdf);
}

// -----------------------------------------------------------------------------
// Direct Bitpack Vectorizer Kernel from SDF
// -----------------------------------------------------------------------------
__global__ void gpu_sdf_threshold_bitpack_kernel(
    const float* __restrict__ sdf,
    unsigned int* __restrict__ out_bit_words,
    int width, int height, int layers,
    const float* __restrict__ thresholds_mm,
    int num_offsets,
    int row_stride_words
) {
    int x_word = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y;
    int z = blockIdx.z;

    if (x_word >= row_stride_words || y >= height || z >= layers) return;

    size_t in_layer_offset = ((size_t)z * height + y) * width;
    size_t total_word_layer = (size_t)layers * height * row_stride_words;
    size_t out_layer_offset = (size_t)z * height * row_stride_words + y * row_stride_words;

    int x_base = x_word * 32;

    for (int o = 0; o < num_offsets; ++o) {
        float thresh = thresholds_mm[o];
        unsigned int packed = 0;

        #pragma unroll
        for (int b = 0; b < 32; ++b) {
            int x = x_base + b;
            if (x < width) {
                float d = sdf[in_layer_offset + x];
                if (d >= thresh) {
                    packed |= (1u << b);
                }
            }
        }

        size_t out_idx = (size_t)o * total_word_layer + out_layer_offset + x_word;
        out_bit_words[out_idx] = packed;
    }
}

// -----------------------------------------------------------------------------
// Public C API Implementation
// -----------------------------------------------------------------------------
extern "C" {

GPU_SDF_API GpuSdfStackHandle gpu_sdf_create(
    const GpuSdfMetadata* meta,
    const int32_t* layer_z_heights_um
) {
    if (!meta || meta->width <= 0 || meta->height <= 0 || meta->num_layers <= 0) return nullptr;

    GpuSdfStackHandle stack = new GpuSdfStack_T();
    stack->meta = *meta;
    stack->owns_stream = true;
    if (cudaStreamCreate(&stack->stream) != cudaSuccess) {
        delete stack;
        return nullptr;
    }

    size_t total_voxels = (size_t)meta->num_layers * meta->height * meta->width;
    size_t sdf_bytes = total_voxels * sizeof(float);
    size_t z_bytes = meta->num_layers * sizeof(int32_t);

    if (cudaMalloc(&stack->d_sdf, sdf_bytes) != cudaSuccess ||
        cudaMalloc(&stack->d_z_um, z_bytes) != cudaSuccess) {
        if (stack->d_sdf) cudaFree(stack->d_sdf);
        if (stack->d_z_um) cudaFree(stack->d_z_um);
        cudaStreamDestroy(stack->stream);
        delete stack;
        return nullptr;
    }

    if (layer_z_heights_um) {
        cudaMemcpyAsync(stack->d_z_um, layer_z_heights_um, z_bytes, cudaMemcpyHostToDevice, stack->stream);
    } else {
        // Default zero heights
        cudaMemsetAsync(stack->d_z_um, 0, z_bytes, stack->stream);
    }
    cudaMemsetAsync(stack->d_sdf, 0, sdf_bytes, stack->stream);

    return stack;
}

GPU_SDF_API GpuSdfStackHandle gpu_sdf_from_masks(
    const uint8_t* h_masks,
    int width,
    int height,
    size_t num_layers,
    const int32_t* layer_z_heights_um,
    int32_t min_x_um,
    int32_t min_y_um,
    int32_t pitch_um
) {
    if (!h_masks || width <= 0 || height <= 0 || num_layers == 0) return nullptr;

    GpuSdfMetadata meta;
    meta.width = width;
    meta.height = height;
    meta.num_layers = num_layers;
    meta.min_x_um = min_x_um;
    meta.min_y_um = min_y_um;
    meta.pitch_um = pitch_um > 0 ? pitch_um : 100;

    GpuSdfStackHandle stack = gpu_sdf_create(&meta, layer_z_heights_um);
    if (!stack) return nullptr;

    size_t total_voxels = (size_t)num_layers * height * width;
    size_t mask_bytes = total_voxels * sizeof(uint8_t);
    size_t seed_bytes = total_voxels * sizeof(int);

    uint8_t* d_masks = nullptr;
    int* d_seeds_ext_a = nullptr;
    int* d_seeds_ext_b = nullptr;
    int* d_seeds_int_a = nullptr;
    int* d_seeds_int_b = nullptr;

    if (cudaMalloc(&d_masks, mask_bytes) != cudaSuccess ||
        cudaMalloc(&d_seeds_ext_a, seed_bytes) != cudaSuccess ||
        cudaMalloc(&d_seeds_ext_b, seed_bytes) != cudaSuccess ||
        cudaMalloc(&d_seeds_int_a, seed_bytes) != cudaSuccess ||
        cudaMalloc(&d_seeds_int_b, seed_bytes) != cudaSuccess) {
        if (d_masks) cudaFree(d_masks);
        if (d_seeds_ext_a) cudaFree(d_seeds_ext_a);
        if (d_seeds_ext_b) cudaFree(d_seeds_ext_b);
        if (d_seeds_int_a) cudaFree(d_seeds_int_a);
        if (d_seeds_int_b) cudaFree(d_seeds_int_b);
        gpu_sdf_destroy(stack);
        return nullptr;
    }

    cudaMemcpyAsync(d_masks, h_masks, mask_bytes, cudaMemcpyHostToDevice, stack->stream);

    dim3 block(16, 16, 1);
    dim3 grid((width + 15) / 16, (height + 15) / 16, (unsigned int)num_layers);

    jfa_init_interior_exterior_kernel<<<grid, block, 0, stack->stream>>>(
        d_masks, d_seeds_ext_a, d_seeds_int_a, width, height, (int)num_layers
    );

    int max_dim = std::max(width, height);
    int step = 1;
    while (step * 2 < max_dim) step *= 2;

    int* in_seeds_ext = d_seeds_ext_a;
    int* out_seeds_ext = d_seeds_ext_b;
    int* in_seeds_int = d_seeds_int_a;
    int* out_seeds_int = d_seeds_int_b;

    while (step >= 1) {
        jfa_step_dual_kernel<<<grid, block, 0, stack->stream>>>(
            in_seeds_ext, out_seeds_ext,
            in_seeds_int, out_seeds_int,
            width, height, (int)num_layers, step
        );
        std::swap(in_seeds_ext, out_seeds_ext);
        std::swap(in_seeds_int, out_seeds_int);
        step /= 2;
    }

    // JFA+2 and JFA+1 corrections
    jfa_step_dual_kernel<<<grid, block, 0, stack->stream>>>(
        in_seeds_ext, out_seeds_ext,
        in_seeds_int, out_seeds_int,
        width, height, (int)num_layers, 2
    );
    std::swap(in_seeds_ext, out_seeds_ext);
    std::swap(in_seeds_int, out_seeds_int);

    jfa_step_dual_kernel<<<grid, block, 0, stack->stream>>>(
        in_seeds_ext, out_seeds_ext,
        in_seeds_int, out_seeds_int,
        width, height, (int)num_layers, 1
    );
    std::swap(in_seeds_ext, out_seeds_ext);
    std::swap(in_seeds_int, out_seeds_int);

    float pitch_mm = (float)meta.pitch_um * 0.001f;
    fuse_signed_sdf_kernel<<<grid, block, 0, stack->stream>>>(
        d_masks, in_seeds_ext, in_seeds_int, stack->d_sdf,
        width, height, (int)num_layers, pitch_mm
    );

    cudaStreamSynchronize(stack->stream);

    cudaFree(d_masks);
    cudaFree(d_seeds_ext_a);
    cudaFree(d_seeds_ext_b);
    cudaFree(d_seeds_int_a);
    cudaFree(d_seeds_int_b);

    return stack;
}

// Forward declare GPU mesh rasterizer helper from gpu_offset.cu
extern int gpu_slice_mesh_to_device_mask(
    const GpuTriangle3D* triangles,
    size_t num_triangles,
    float min_z, float h_first, float layer_height, int num_layers,
    float bounds_min_x, float bounds_min_y, float pitch,
    int cols, int rows,
    uint8_t* d_masks,
    cudaStream_t stream
);

GPU_SDF_API GpuSdfStackHandle gpu_sdf_from_mesh(
    const GpuTriangle3D* triangles,
    size_t num_triangles,
    float min_x, float min_y, float min_z,
    float max_x, float max_y, float max_z,
    float layer_height_mm,
    float first_layer_height_mm,
    float pitch_mm,
    float pad_mm
) {
    if (!triangles || num_triangles == 0) return nullptr;

    float pad = (pad_mm > 0.0f) ? pad_mm : 1.5f;
    float pitch = (pitch_mm > 0.0f) ? pitch_mm : 0.10f;
    float bounds_min_x = min_x - pad;
    float bounds_min_y = min_y - pad;
    float bounds_max_x = max_x + pad;
    float bounds_max_y = max_y + pad;

    int cols = (int)std::ceil((bounds_max_x - bounds_min_x) / pitch);
    int rows = (int)std::ceil((bounds_max_y - bounds_min_y) / pitch);

    float h_first = (first_layer_height_mm > 0.0f) ? first_layer_height_mm : layer_height_mm;
    float z_first = min_z + h_first;

    int num_layers = 0;
    if (max_z >= z_first) {
        num_layers = (int)std::floor((max_z - z_first) / layer_height_mm) + 1;
    }
    if (num_layers <= 0) return nullptr;

    std::vector<int32_t> z_heights_um(num_layers);
    for (int l = 0; l < num_layers; ++l) {
        float z_val = z_first + l * layer_height_mm;
        z_heights_um[l] = (int32_t)std::round(z_val * 1000.0f);
    }

    GpuSdfMetadata meta;
    meta.width = cols;
    meta.height = rows;
    meta.num_layers = num_layers;
    meta.min_x_um = (int32_t)std::round(bounds_min_x * 1000.0f);
    meta.min_y_um = (int32_t)std::round(bounds_min_y * 1000.0f);
    meta.pitch_um = (int32_t)std::round(pitch * 1000.0f);

    GpuSdfStackHandle stack = gpu_sdf_create(&meta, z_heights_um.data());
    if (!stack) return nullptr;

    size_t total_volume = (size_t)num_layers * rows * cols;
    size_t tri_bytes = num_triangles * sizeof(GpuTriangle3D);
    size_t seed_bytes = total_volume * sizeof(int);

    GpuTriangle3D* d_triangles = nullptr;
    uint8_t* d_masks = nullptr;
    int* d_seeds_ext_a = nullptr;
    int* d_seeds_ext_b = nullptr;
    int* d_seeds_int_a = nullptr;
    int* d_seeds_int_b = nullptr;

    if (cudaMalloc(&d_triangles, tri_bytes) != cudaSuccess ||
        cudaMalloc(&d_masks, total_volume * sizeof(uint8_t)) != cudaSuccess ||
        cudaMalloc(&d_seeds_ext_a, seed_bytes) != cudaSuccess ||
        cudaMalloc(&d_seeds_ext_b, seed_bytes) != cudaSuccess ||
        cudaMalloc(&d_seeds_int_a, seed_bytes) != cudaSuccess ||
        cudaMalloc(&d_seeds_int_b, seed_bytes) != cudaSuccess) {
        if (d_triangles) cudaFree(d_triangles);
        if (d_masks) cudaFree(d_masks);
        if (d_seeds_ext_a) cudaFree(d_seeds_ext_a);
        if (d_seeds_ext_b) cudaFree(d_seeds_ext_b);
        if (d_seeds_int_a) cudaFree(d_seeds_int_a);
        if (d_seeds_int_b) cudaFree(d_seeds_int_b);
        gpu_sdf_destroy(stack);
        return nullptr;
    }

    cudaMemcpyAsync(d_triangles, triangles, tri_bytes, cudaMemcpyHostToDevice, stack->stream);

    int slice_ret = gpu_slice_mesh_to_device_mask(
        d_triangles, num_triangles, min_z, h_first, layer_height_mm, num_layers,
        bounds_min_x, bounds_min_y, pitch, cols, rows, d_masks, stack->stream
    );
    if (slice_ret != GPU_OFFSET_SUCCESS) {
        cudaFree(d_triangles);
        cudaFree(d_masks);
        cudaFree(d_seeds_ext_a);
        cudaFree(d_seeds_ext_b);
        cudaFree(d_seeds_int_a);
        cudaFree(d_seeds_int_b);
        gpu_sdf_destroy(stack);
        return nullptr;
    }

    // Run Batched 3D JFA Dual Direction Signed Distance
    dim3 block(16, 16, 1);
    dim3 grid((cols + 15) / 16, (rows + 15) / 16, num_layers);

    jfa_init_interior_exterior_kernel<<<grid, block, 0, stack->stream>>>(
        d_masks, d_seeds_ext_a, d_seeds_int_a, cols, rows, num_layers
    );

    int max_dim = std::max(cols, rows);
    int step = 1;
    while (step * 2 < max_dim) step *= 2;

    int* in_seeds_ext = d_seeds_ext_a;
    int* out_seeds_ext = d_seeds_ext_b;
    int* in_seeds_int = d_seeds_int_a;
    int* out_seeds_int = d_seeds_int_b;

    while (step >= 1) {
        jfa_step_dual_kernel<<<grid, block, 0, stack->stream>>>(
            in_seeds_ext, out_seeds_ext,
            in_seeds_int, out_seeds_int,
            cols, rows, num_layers, step
        );
        std::swap(in_seeds_ext, out_seeds_ext);
        std::swap(in_seeds_int, out_seeds_int);
        step /= 2;
    }

    jfa_step_dual_kernel<<<grid, block, 0, stack->stream>>>(
        in_seeds_ext, out_seeds_ext,
        in_seeds_int, out_seeds_int,
        cols, rows, num_layers, 2
    );
    std::swap(in_seeds_ext, out_seeds_ext);
    std::swap(in_seeds_int, out_seeds_int);

    jfa_step_dual_kernel<<<grid, block, 0, stack->stream>>>(
        in_seeds_ext, out_seeds_ext,
        in_seeds_int, out_seeds_int,
        cols, rows, num_layers, 1
    );
    std::swap(in_seeds_ext, out_seeds_ext);
    std::swap(in_seeds_int, out_seeds_int);

    fuse_signed_sdf_kernel<<<grid, block, 0, stack->stream>>>(
        d_masks, in_seeds_ext, in_seeds_int, stack->d_sdf,
        cols, rows, num_layers, pitch
    );

    cudaStreamSynchronize(stack->stream);

    cudaFree(d_triangles);
    cudaFree(d_masks);
    cudaFree(d_seeds_ext_a);
    cudaFree(d_seeds_ext_b);
    cudaFree(d_seeds_int_a);
    cudaFree(d_seeds_int_b);

    return stack;
}

GPU_SDF_API GpuSdfStackHandle gpu_sdf_clone(GpuSdfStackHandle stack) {
    if (!stack) return nullptr;

    std::vector<int32_t> z_um(stack->meta.num_layers);
    gpu_sdf_get_layer_z_um(stack, z_um.data());

    GpuSdfStackHandle clone = gpu_sdf_create(&stack->meta, z_um.data());
    if (!clone) return nullptr;

    size_t total_voxels = (size_t)stack->meta.num_layers * stack->meta.height * stack->meta.width;
    cudaMemcpyAsync(clone->d_sdf, stack->d_sdf, total_voxels * sizeof(float),
                    cudaMemcpyDeviceToDevice, clone->stream);
    cudaStreamSynchronize(clone->stream);

    return clone;
}

GPU_SDF_API void gpu_sdf_destroy(GpuSdfStackHandle stack) {
    if (!stack) return;
    if (stack->d_sdf) {
        cudaFree(stack->d_sdf);
        stack->d_sdf = nullptr;
    }
    if (stack->d_z_um) {
        cudaFree(stack->d_z_um);
        stack->d_z_um = nullptr;
    }
    if (stack->owns_stream && stack->stream) {
        cudaStreamDestroy(stack->stream);
        stack->stream = nullptr;
    }
    delete stack;
}

GPU_SDF_API int gpu_sdf_get_metadata(GpuSdfStackHandle stack, GpuSdfMetadata* out_meta) {
    if (!stack || !out_meta) return GPU_OFFSET_ERR_INVALID_ARG;
    *out_meta = stack->meta;
    return GPU_OFFSET_SUCCESS;
}

GPU_SDF_API int gpu_sdf_get_layer_z_um(GpuSdfStackHandle stack, int32_t* out_z_um) {
    if (!stack || !out_z_um) return GPU_OFFSET_ERR_INVALID_ARG;
    CUDA_CHECK_INT(cudaMemcpy(out_z_um, stack->d_z_um, stack->meta.num_layers * sizeof(int32_t), cudaMemcpyDeviceToHost));
    return GPU_OFFSET_SUCCESS;
}

GPU_SDF_API int gpu_sdf_to_host(GpuSdfStackHandle stack, float* h_out_data) {
    if (!stack || !h_out_data) return GPU_OFFSET_ERR_INVALID_ARG;
    size_t total_voxels = (size_t)stack->meta.num_layers * stack->meta.height * stack->meta.width;
    CUDA_CHECK_INT(cudaMemcpy(h_out_data, stack->d_sdf, total_voxels * sizeof(float), cudaMemcpyDeviceToHost));
    return GPU_OFFSET_SUCCESS;
}

GPU_SDF_API float* gpu_sdf_get_device_ptr(GpuSdfStackHandle stack) {
    return stack ? stack->d_sdf : nullptr;
}

GPU_SDF_API int32_t* gpu_sdf_get_device_z_ptr(GpuSdfStackHandle stack) {
    return stack ? stack->d_z_um : nullptr;
}

// -----------------------------------------------------------------------------
// CSG Execution Engine
// -----------------------------------------------------------------------------
GPU_SDF_API GpuSdfStackHandle gpu_sdf_csg_op(
    GpuSdfStackHandle a,
    GpuSdfStackHandle b,
    GpuSdfOp op,
    GpuSdfLayerMatchMode match_mode,
    float blend_radius_mm
) {
    if (!a || !b) return nullptr;

    // Check if grids and layers are 100% identical
    std::vector<int32_t> z_a(a->meta.num_layers);
    std::vector<int32_t> z_b(b->meta.num_layers);
    gpu_sdf_get_layer_z_um(a, z_a.data());
    gpu_sdf_get_layer_z_um(b, z_b.data());

    bool identical_grid = (a->meta.min_x_um == b->meta.min_x_um) &&
                         (a->meta.min_y_um == b->meta.min_y_um) &&
                         (a->meta.pitch_um == b->meta.pitch_um) &&
                         (a->meta.width == b->meta.width) &&
                         (a->meta.height == b->meta.height) &&
                         (a->meta.num_layers == b->meta.num_layers) &&
                         (z_a == z_b);

    if (identical_grid) {
        GpuSdfStackHandle out_stack = gpu_sdf_create(&a->meta, z_a.data());
        if (!out_stack) return nullptr;

        size_t total_voxels = (size_t)a->meta.num_layers * a->meta.height * a->meta.width;
        int block_sz = 256;
        int grid_sz = (total_voxels + block_sz - 1) / block_sz;

        gpu_sdf_csg_identical_kernel<<<grid_sz, block_sz, 0, out_stack->stream>>>(
            a->d_sdf, b->d_sdf, out_stack->d_sdf, total_voxels, (int)op, blend_radius_mm
        );
        cudaStreamSynchronize(out_stack->stream);
        return out_stack;
    }

    // Harmonize Spatial Grid (XY)
    int32_t pitch_um = a->meta.pitch_um;
    int32_t min_x_um = std::min(a->meta.min_x_um, b->meta.min_x_um);
    int32_t min_y_um = std::min(a->meta.min_y_um, b->meta.min_y_um);
    int32_t max_x_um = std::max(a->meta.min_x_um + a->meta.width * a->meta.pitch_um,
                                b->meta.min_x_um + b->meta.width * b->meta.pitch_um);
    int32_t max_y_um = std::max(a->meta.min_y_um + a->meta.height * a->meta.pitch_um,
                                b->meta.min_y_um + b->meta.height * b->meta.pitch_um);

    int out_w = (int)std::ceil((float)(max_x_um - min_x_um) / pitch_um);
    int out_h = (int)std::ceil((float)(max_y_um - min_y_um) / pitch_um);

    // Harmonize Layers (Z)
    std::vector<int32_t> out_z_um;
    if (match_mode == GPU_SDF_LAYER_MATCH_A_DOMINANT) {
        out_z_um = z_a;
    } else if (match_mode == GPU_SDF_LAYER_MATCH_UNION) {
        std::set<int32_t> z_set(z_a.begin(), z_a.end());
        z_set.insert(z_b.begin(), z_b.end());
        out_z_um.assign(z_set.begin(), z_set.end());
    } else { // INTERSECTION
        std::set<int32_t> set_b(z_b.begin(), z_b.end());
        for (int32_t z : z_a) {
            if (set_b.count(z)) out_z_um.push_back(z);
        }
    }

    if (out_z_um.empty()) return nullptr;

    GpuSdfMetadata meta_out;
    meta_out.width = out_w;
    meta_out.height = out_h;
    meta_out.num_layers = out_z_um.size();
    meta_out.min_x_um = min_x_um;
    meta_out.min_y_um = min_y_um;
    meta_out.pitch_um = pitch_um;

    GpuSdfStackHandle out_stack = gpu_sdf_create(&meta_out, out_z_um.data());
    if (!out_stack) return nullptr;

    // Build Layer Sampling Info table on Host & upload to Device
    std::vector<LayerSamplingInfo> layer_infos(meta_out.num_layers);

    auto find_layer_sampling = [](const std::vector<int32_t>& z_list, int32_t target_z, int& l0, int& l1, float& t) {
        if (z_list.empty()) {
            l0 = l1 = -1; t = -1.0f; return;
        }
        if (target_z < z_list.front() || target_z > z_list.back()) {
            l0 = l1 = -1; t = -1.0f; return;
        }
        auto it = std::lower_bound(z_list.begin(), z_list.end(), target_z);
        if (it != z_list.end() && *it == target_z) {
            l0 = l1 = (int)(it - z_list.begin());
            t = 0.0f;
            return;
        }
        // Interpolation between adjacent layers
        size_t idx1 = it - z_list.begin();
        size_t idx0 = idx1 - 1;
        l0 = (int)idx0;
        l1 = (int)idx1;
        t = (float)(target_z - z_list[idx0]) / (float)(z_list[idx1] - z_list[idx0]);
    };

    for (size_t l = 0; l < meta_out.num_layers; ++l) {
        int32_t tz = out_z_um[l];
        find_layer_sampling(z_a, tz, layer_infos[l].l0_a, layer_infos[l].l1_a, layer_infos[l].t_a);
        find_layer_sampling(z_b, tz, layer_infos[l].l0_b, layer_infos[l].l1_b, layer_infos[l].t_b);
    }

    LayerSamplingInfo* d_layer_map = nullptr;
    cudaMalloc(&d_layer_map, meta_out.num_layers * sizeof(LayerSamplingInfo));
    cudaMemcpyAsync(d_layer_map, layer_infos.data(), meta_out.num_layers * sizeof(LayerSamplingInfo),
                    cudaMemcpyHostToDevice, out_stack->stream);

    dim3 block(16, 16, 1);
    dim3 grid((out_w + 15) / 16, (out_h + 15) / 16, (unsigned int)meta_out.num_layers);

    gpu_sdf_csg_resample_kernel<<<grid, block, 0, out_stack->stream>>>(
        a->d_sdf, a->meta,
        b->d_sdf, b->meta,
        out_stack->d_sdf, meta_out,
        d_layer_map, (int)op, blend_radius_mm
    );

    cudaStreamSynchronize(out_stack->stream);
    cudaFree(d_layer_map);

    return out_stack;
}

GPU_SDF_API GpuSdfStackHandle gpu_sdf_offset(
    GpuSdfStackHandle stack,
    float delta_mm
) {
    if (!stack) return nullptr;

    std::vector<int32_t> z_um(stack->meta.num_layers);
    gpu_sdf_get_layer_z_um(stack, z_um.data());

    GpuSdfStackHandle out_stack = gpu_sdf_create(&stack->meta, z_um.data());
    if (!out_stack) return nullptr;

    size_t total_voxels = (size_t)stack->meta.num_layers * stack->meta.height * stack->meta.width;
    int block_sz = 256;
    int grid_sz = (total_voxels + block_sz - 1) / block_sz;

    gpu_sdf_offset_kernel<<<grid_sz, block_sz, 0, out_stack->stream>>>(
        stack->d_sdf, out_stack->d_sdf, total_voxels, delta_mm
    );
    cudaStreamSynchronize(out_stack->stream);

    return out_stack;
}

GPU_SDF_API GpuSdfStackHandle gpu_sdf_shell(
    GpuSdfStackHandle stack,
    float thickness_mm
) {
    if (!stack || thickness_mm <= 0.0f) return nullptr;

    std::vector<int32_t> z_um(stack->meta.num_layers);
    gpu_sdf_get_layer_z_um(stack, z_um.data());

    GpuSdfStackHandle out_stack = gpu_sdf_create(&stack->meta, z_um.data());
    if (!out_stack) return nullptr;

    size_t total_voxels = (size_t)stack->meta.num_layers * stack->meta.height * stack->meta.width;
    int block_sz = 256;
    int grid_sz = (total_voxels + block_sz - 1) / block_sz;

    gpu_sdf_shell_kernel<<<grid_sz, block_sz, 0, out_stack->stream>>>(
        stack->d_sdf, out_stack->d_sdf, total_voxels, thickness_mm
    );
    cudaStreamSynchronize(out_stack->stream);

    return out_stack;
}

GPU_SDF_API GpuSdfStackHandle gpu_sdf_generate_infill(
    GpuSdfStackHandle stack,
    GpuSdfInfillPattern pattern,
    float cell_size_mm,
    float density_ratio,
    float wall_offset_mm
) {
    if (!stack) return nullptr;

    std::vector<int32_t> z_um(stack->meta.num_layers);
    gpu_sdf_get_layer_z_um(stack, z_um.data());

    GpuSdfStackHandle out_stack = gpu_sdf_create(&stack->meta, z_um.data());
    if (!out_stack) return nullptr;

    dim3 block(16, 16, 1);
    dim3 grid((stack->meta.width + 15) / 16, (stack->meta.height + 15) / 16, (unsigned int)stack->meta.num_layers);

    gpu_sdf_infill_kernel<<<grid, block, 0, out_stack->stream>>>(
        stack->d_sdf, out_stack->d_sdf,
        stack->meta.width, stack->meta.height, (int)stack->meta.num_layers,
        stack->meta.min_x_um, stack->meta.min_y_um, stack->meta.pitch_um,
        stack->d_z_um, (int)pattern, cell_size_mm, density_ratio, wall_offset_mm
    );
    cudaStreamSynchronize(out_stack->stream);

    return out_stack;
}

GPU_SDF_API int gpu_sdf_vectorize(
    GpuSdfStackHandle stack,
    const GpuOffsetConfig* config,
    GpuLayerResult** out_results,
    size_t* out_num_layers,
    GpuOffsetStats* stats
) {
    if (!stack || !config || !out_results) return GPU_OFFSET_ERR_INVALID_ARG;

    auto t_total_start = std::chrono::high_resolution_clock::now();

    int cols = stack->meta.width;
    int rows = stack->meta.height;
    size_t num_layers = stack->meta.num_layers;
    if (out_num_layers) *out_num_layers = num_layers;

    float bounds_min_x = (float)stack->meta.min_x_um * 0.001f;
    float bounds_min_y = (float)stack->meta.min_y_um * 0.001f;
    float pitch_mm = (float)stack->meta.pitch_um * 0.001f;

    int P = (config->perimeters_count >= 0) ? config->perimeters_count : 2;
    float W = (config->layer_width > 0.0f) ? config->layer_width : 0.45f;
    std::vector<float> resolved_offsets;
    if (config->offset_distances_mm && config->num_offsets > 0) {
        resolved_offsets.assign(config->offset_distances_mm, config->offset_distances_mm + config->num_offsets);
        P = (config->perimeters_count > 0) ? config->perimeters_count : (int)resolved_offsets.size() - 1;
        if (P < 1) P = (int)resolved_offsets.size();
    } else {
        resolved_offsets.resize(P + 1);
        for (int i = 0; i < P; ++i) {
            resolved_offsets[i] = (i + 0.5f) * W;
        }
        resolved_offsets[P] = (float)P * W;
    }
    size_t num_offsets = resolved_offsets.size();

    float simplify_tol = config->simplify_tolerance_mm;
    if (simplify_tol <= 0.0f && config->subpixel_lerp != 0) {
        simplify_tol = 0.02f; // default 0.02 mm (20 µm) tolerance
    }

    size_t total_voxels = (size_t)num_layers * rows * cols;
    size_t sdf_bytes = total_voxels * sizeof(float);

    float* h_pinned_sdf = nullptr;
    CUDA_CHECK_INT(cudaMallocHost(&h_pinned_sdf, sdf_bytes));
    CUDA_CHECK_INT(cudaMemcpyAsync(h_pinned_sdf, stack->d_sdf, sdf_bytes,
                                   cudaMemcpyDeviceToHost, stack->stream));
    CUDA_CHECK_INT(cudaStreamSynchronize(stack->stream));

    std::vector<int32_t> z_heights_um(num_layers);
    gpu_sdf_get_layer_z_um(stack, z_heights_um.data());

    // Multi-threaded C++ Sub-Pixel LERP Vectorization
    auto t_vec_0 = std::chrono::high_resolution_clock::now();

    GpuLayerResult* res_array = (GpuLayerResult*)malloc(num_layers * sizeof(GpuLayerResult));
    memset(res_array, 0, num_layers * sizeof(GpuLayerResult));

    unsigned int num_threads = std::max(1u, std::thread::hardware_concurrency());
    std::vector<std::thread> workers;

    auto vec_worker = [&](size_t start_layer, size_t end_layer) {
        int stride_w = cols + 1;
        size_t max_edges = 2 * (size_t)(rows + 1) * stride_w;
        std::vector<int> next_edge_buf(max_edges, -1);
        std::vector<int> active_edges_buf;
        active_edges_buf.reserve(cols * rows / 16);
        std::vector<uint8_t> visited_buf(max_edges, 0);

        for (size_t l = start_layer; l < end_layer; ++l) {
            res_array[l].z_height = (float)z_heights_um[l] * 0.001f;
            res_array[l].num_perimeters = P;
            res_array[l].perimeters = (GpuOffsetPerimeter*)malloc(P * sizeof(GpuOffsetPerimeter));

            const float* layer_sdf_ptr = h_pinned_sdf + l * ((size_t)rows * cols);

            for (int o = 0; o < P; ++o) {
                float iso_thresh = resolved_offsets[o];

                std::vector<GpuContour> contours;
                vectorize_sdf_layer(layer_sdf_ptr, cols, rows, bounds_min_x, bounds_min_y, pitch_mm,
                                    iso_thresh, simplify_tol,
                                    next_edge_buf, active_edges_buf, visited_buf, contours);

                res_array[l].perimeters[o].offset_distance_mm = iso_thresh;
                res_array[l].perimeters[o].num_contours = contours.size();
                if (!contours.empty()) {
                    res_array[l].perimeters[o].contours = (GpuContour*)malloc(contours.size() * sizeof(GpuContour));
                    memcpy(res_array[l].perimeters[o].contours, contours.data(), contours.size() * sizeof(GpuContour));
                } else {
                    res_array[l].perimeters[o].contours = nullptr;
                }
            }

            // Infill
            float infill_iso_thresh = resolved_offsets[P];
            std::vector<GpuContour> infill_contours;
            vectorize_sdf_layer(layer_sdf_ptr, cols, rows, bounds_min_x, bounds_min_y, pitch_mm,
                                infill_iso_thresh, simplify_tol,
                                next_edge_buf, active_edges_buf, visited_buf, infill_contours);

            res_array[l].infill.offset_distance_mm = infill_iso_thresh;
            res_array[l].infill.width = cols;
            res_array[l].infill.height = rows;
            res_array[l].infill.num_contours = infill_contours.size();
            if (!infill_contours.empty()) {
                res_array[l].infill.contours = (GpuContour*)malloc(infill_contours.size() * sizeof(GpuContour));
                memcpy(res_array[l].infill.contours, infill_contours.data(), infill_contours.size() * sizeof(GpuContour));
            } else {
                res_array[l].infill.contours = nullptr;
            }

            if (config->export_infill_mask) {
                uint8_t* infill_mask = (uint8_t*)malloc((size_t)cols * rows);
                for (int i = 0; i < cols * rows; ++i) {
                    infill_mask[i] = (layer_sdf_ptr[i] >= infill_iso_thresh) ? 1 : 0;
                }
                res_array[l].infill.mask = infill_mask;
            } else {
                res_array[l].infill.mask = nullptr;
            }
        }
    };

    size_t chunk = (num_layers + num_threads - 1) / num_threads;
    for (unsigned int t = 0; t < num_threads; ++t) {
        size_t s = t * chunk;
        size_t e = std::min(s + chunk, num_layers);
        if (s < e) workers.emplace_back(vec_worker, s, e);
    }
    for (auto& w : workers) w.join();

    auto t_vec_1 = std::chrono::high_resolution_clock::now();
    double t_vec_ms = std::chrono::duration<double, std::milli>(t_vec_1 - t_vec_0).count();

    auto t_total_end = std::chrono::high_resolution_clock::now();
    double t_total_ms = std::chrono::duration<double, std::milli>(t_total_end - t_total_start).count();

    if (stats) {
        stats->t_vectorize_ms = t_vec_ms;
        stats->t_total_pipeline_ms = t_total_ms;
    }

    cudaFreeHost(h_pinned_sdf);

    *out_results = res_array;
    return GPU_OFFSET_SUCCESS;
}

} // extern "C"