#include "gpu_offset.h"
#include "gpu_sdf.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>

// CUDA Error Checking Macro
#define CUDA_CHECK(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 device helper functions
__device__ __inline__ int pack_seed(short x, short y) {
    return ((int)(unsigned short)y << 16) | (int)(unsigned short)x;
}
__device__ __inline__ short get_seed_x(int packed) {
    return (short)(packed & 0xFFFF);
}
__device__ __inline__ short get_seed_y(int packed) {
    return (short)((packed >> 16) & 0xFFFF);
}

// -----------------------------------------------------------------------------
// CUDA Kernels: Batched 3D Jump Flooding Algorithm
// -----------------------------------------------------------------------------
__global__ void batched_jfa_init_kernel(
    const uint8_t* __restrict__ masks,
    int* __restrict__ seeds,
    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 layer_stride = (size_t)width * height;
    size_t idx = (size_t)z * layer_stride + y * width + x;

    // Background pixel (mask == 0) is initialized as seed point
    if (masks[idx] == 0) {
        seeds[idx] = pack_seed((short)x, (short)y);
    } else {
        seeds[idx] = -1;
    }
}

__global__ void batched_jfa_step_kernel(
    const int* __restrict__ in_seeds,
    int* __restrict__ out_seeds,
    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_stride = (size_t)width * height;
    size_t layer_offset = (size_t)z * layer_stride;
    size_t idx = layer_offset + y * width + x;

    int best_seed = in_seeds[idx];
    int best_dist2 = 0x7FFFFFFF;
    if (best_seed != -1) {
        int dx = (int)get_seed_x(best_seed) - x;
        int dy = (int)get_seed_y(best_seed) - y;
        best_dist2 = 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) {
                int n_seed = in_seeds[layer_offset + ny * width + nx];
                if (n_seed != -1) {
                    int dist_x = (int)get_seed_x(n_seed) - x;
                    int dist_y = (int)get_seed_y(n_seed) - y;
                    int dist2 = dist_x * dist_x + dist_y * dist_y;
                    if (dist2 < best_dist2) {
                        best_dist2 = dist2;
                        best_seed = n_seed;
                    }
                }
            }
        }
    }
    out_seeds[idx] = best_seed;
}

__global__ void batched_jfa_threshold_kernel(
    const int* __restrict__ seeds,
    const uint8_t* __restrict__ masks,
    uint8_t* __restrict__ out_offsets,
    int width, int height, int layers,
    float pitch_mm,
    const float* __restrict__ thresholds_mm,
    int num_offsets
) {
    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_pixels = (size_t)width * height;
    size_t total_volume = (size_t)layers * layer_pixels;
    size_t in_idx = (size_t)z * layer_pixels + y * width + x;

    uint8_t mask_val = masks[in_idx];
    int s = seeds[in_idx];

    float dist_mm = 0.0f;
    if (mask_val > 0 && s != -1) {
        float dx = (float)(get_seed_x(s) - x);
        float dy = (float)(get_seed_y(s) - y);
        dist_mm = sqrtf(dx * dx + dy * dy) * pitch_mm;
    }

    for (int o = 0; o < num_offsets; ++o) {
        size_t out_idx = (size_t)o * total_volume + in_idx;
        out_offsets[out_idx] = (dist_mm >= thresholds_mm[o]) ? 1 : 0;
    }
}

__global__ void batched_jfa_threshold_bitpack_kernel(
    const int* __restrict__ seeds,
    const uint8_t* __restrict__ masks,
    unsigned int* __restrict__ out_bit_words,
    int width, int height, int layers,
    float pitch_mm,
    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 layer_pixels = (size_t)width * height;
    size_t in_layer_offset = (size_t)z * layer_pixels + 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];
        float thresh_sq_px = (thresh / pitch_mm) * (thresh / pitch_mm);

        unsigned int packed = 0;

        #pragma unroll
        for (int b = 0; b < 32; ++b) {
            int x = x_base + b;
            if (x < width) {
                size_t in_idx = in_layer_offset + x;
                uint8_t m = masks[in_idx];
                int s = seeds[in_idx];

                if (m > 0 && s != -1) {
                    int sx = (int)get_seed_x(s);
                    int sy = (int)get_seed_y(s);
                    int dx = sx - x;
                    int dy = sy - y;
                    int dist_sq = dx * dx + dy * dy;
                    if ((float)dist_sq >= thresh_sq_px) {
                        packed |= (1u << b);
                    }
                }
            }
        }

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

// -----------------------------------------------------------------------------
// Direct GPU Mesh Slicing & Scanline Parity Kernels
// -----------------------------------------------------------------------------
__global__ void gpu_mesh_slice_raster_kernel(
    const GpuTriangle3D* __restrict__ triangles,
    size_t num_triangles,
    float z_min_base,
    float first_layer_height,
    float layer_height,
    int num_layers,
    float bounds_min_x,
    float bounds_min_y,
    float pitch_mm,
    int cols,
    int rows,
    int row_stride_words,
    unsigned int* __restrict__ parity_words
) {
    size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= num_triangles) return;

    GpuTriangle3D tri = triangles[i];

    float tz_min = fminf(tri.v0[2], fminf(tri.v1[2], tri.v2[2]));
    float tz_max = fmaxf(tri.v0[2], fmaxf(tri.v1[2], tri.v2[2]));

    float z_first = z_min_base + first_layer_height;
    int l_start = (tz_min <= z_first) ? 0 : (int)ceilf((tz_min - z_first) / layer_height);
    int l_end   = (int)floorf((tz_max - z_first) / layer_height);

    if (l_start < 0) l_start = 0;
    if (l_end >= num_layers) l_end = num_layers - 1;
    if (l_start > l_end) return;

    const float* v[3] = { tri.v0, tri.v1, tri.v2 };

    for (int l = l_start; l <= l_end; ++l) {
        float z0 = z_first + l * layer_height;
        if (z0 < tz_min || z0 > tz_max) continue;

        float pts_x[3];
        float pts_y[3];
        int count = 0;

        #pragma unroll
        for (int e = 0; e < 3; ++e) {
            int next = (e + 1) % 3;
            float za = v[e][2];
            float zb = v[next][2];

            if ((za <= z0 && zb >= z0) || (zb <= z0 && za >= z0)) {
                if (fabsf(zb - za) > 1e-7f) {
                    float t = (z0 - za) / (zb - za);
                    pts_x[count] = v[e][0] + t * (v[next][0] - v[e][0]);
                    pts_y[count] = v[e][1] + t * (v[next][1] - v[e][1]);
                    count++;
                    if (count == 2) break;
                }
            }
        }

        if (count != 2) continue;

        float gx0 = (pts_x[0] - bounds_min_x) / pitch_mm - 0.5f;
        float gy0 = (pts_y[0] - bounds_min_y) / pitch_mm - 0.5f;
        float gx1 = (pts_x[1] - bounds_min_x) / pitch_mm - 0.5f;
        float gy1 = (pts_y[1] - bounds_min_y) / pitch_mm - 0.5f;

        if (gy0 > gy1) {
            float tx = gx0; gx0 = gx1; gx1 = tx;
            float ty = gy0; gy0 = gy1; gy1 = ty;
        }

        float dy = gy1 - gy0;
        if (dy < 1e-6f) continue;

        int y_start = (int)ceilf(gy0);
        int y_end   = (int)floorf(gy1 - 1e-6f);

        if (y_start < 0) y_start = 0;
        if (y_end >= rows) y_end = rows - 1;

        size_t layer_word_offset = (size_t)l * rows * row_stride_words;

        for (int y = y_start; y <= y_end; ++y) {
            float t = ((float)y - gy0) / dy;
            float gx = gx0 + t * (gx1 - gx0);
            int x_int = (int)roundf(gx + 0.5f);
            if (x_int < 0) x_int = 0;
            if (x_int >= cols) x_int = cols - 1;

            int w = x_int / 32;
            int b = x_int % 32;
            size_t w_idx = layer_word_offset + (size_t)y * row_stride_words + w;
            atomicXor(&parity_words[w_idx], (1u << b));
        }
    }
}

__global__ void gpu_parity_scan_to_mask_kernel(
    const unsigned int* __restrict__ parity_words,
    uint8_t* __restrict__ masks,
    int cols,
    int rows,
    int num_layers,
    int row_stride_words
) {
    int y = blockIdx.x;
    int l = blockIdx.y;
    if (y >= rows || l >= num_layers) return;

    size_t layer_word_offset = (size_t)l * rows * row_stride_words;
    size_t layer_pixel_offset = (size_t)l * rows * cols;
    size_t row_pixel_offset = layer_pixel_offset + (size_t)y * cols;

    uint8_t inside = 0;
    int x_out = 0;

    for (int w = 0; w < row_stride_words; ++w) {
        unsigned int word = parity_words[layer_word_offset + (size_t)y * row_stride_words + w];
        for (int b = 0; b < 32 && x_out < cols; ++b, ++x_out) {
            inside ^= (uint8_t)((word >> b) & 1u);
            masks[row_pixel_offset + x_out] = inside;
        }
    }
}

// Helper function for direct mesh slicing to device mask
extern "C" 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
) {
    int row_stride_words = (cols + 31) / 32;
    size_t total_words = (size_t)num_layers * rows * row_stride_words;

    unsigned int* d_parity = nullptr;
    CUDA_CHECK(cudaMalloc(&d_parity, total_words * sizeof(unsigned int)));
    CUDA_CHECK(cudaMemsetAsync(d_parity, 0, total_words * sizeof(unsigned int), stream));

    int block_sz = 256;
    int grid_sz = (num_triangles + block_sz - 1) / block_sz;
    gpu_mesh_slice_raster_kernel<<<grid_sz, block_sz, 0, stream>>>(
        triangles, num_triangles, min_z, h_first, layer_height, num_layers,
        bounds_min_x, bounds_min_y, pitch, cols, rows, row_stride_words, d_parity
    );

    dim3 scan_grid(rows, num_layers);
    gpu_parity_scan_to_mask_kernel<<<scan_grid, 1, 0, stream>>>(
        d_parity, d_masks, cols, rows, num_layers, row_stride_words
    );

    cudaFree(d_parity);
    return GPU_OFFSET_SUCCESS;
}

// -----------------------------------------------------------------------------
// Public C API Implementation & Backend Management
// -----------------------------------------------------------------------------
extern "C" {

static cudaStream_t g_cuda_stream = nullptr;
static int g_initialized = 0;
static int g_global_backend = GPU_OFFSET_BACKEND_AUTO;

GPU_OFFSET_API int gpu_offset_has_cuda(void) {
    int count = 0;
    cudaError_t err = cudaGetDeviceCount(&count);
    if (err != cudaSuccess || count <= 0) {
        cudaGetLastError(); // Clear error
        return 0;
    }
    return 1;
}

GPU_OFFSET_API int gpu_offset_set_backend(int backend) {
    if (backend < 0 || backend > 2) return GPU_OFFSET_ERR_INVALID_ARG;
    g_global_backend = backend;
    return GPU_OFFSET_SUCCESS;
}

GPU_OFFSET_API int gpu_offset_get_backend(void) {
    return g_global_backend;
}

GPU_OFFSET_API int gpu_offset_init(int device_id) {
    if (g_initialized) return GPU_OFFSET_SUCCESS;

    if (!gpu_offset_has_cuda()) {
        if (g_global_backend == GPU_OFFSET_BACKEND_CUDA) {
            return GPU_OFFSET_ERR_NO_CUDA;
        }
        return GPU_OFFSET_SUCCESS; // CPU fallback ready
    }

    CUDA_CHECK(cudaSetDevice(device_id));
    CUDA_CHECK(cudaStreamCreate(&g_cuda_stream));

    // Warm up CUDA runtime
    void* d_dummy;
    CUDA_CHECK(cudaMalloc(&d_dummy, 1024));
    CUDA_CHECK(cudaFree(d_dummy));

    g_initialized = 1;
    return GPU_OFFSET_SUCCESS;
}

GPU_OFFSET_API void gpu_offset_cleanup() {
    if (g_cuda_stream) {
        cudaStreamDestroy(g_cuda_stream);
        g_cuda_stream = nullptr;
    }
    g_initialized = 0;
}

// Core GPU Distance Transform execution on masks
static int gpu_compute_masks_internal(
    const uint8_t* h_masks,
    int width,
    int height,
    int num_layers,
    float pitch_mm,
    const float* offset_distances,
    int num_offsets,
    uint8_t* h_out_masks,
    GpuOffsetStats* stats
) {
    if (!g_initialized) {
        if (gpu_offset_init(0) != GPU_OFFSET_SUCCESS) return GPU_OFFSET_ERR_GENERIC;
    }

    size_t layer_pixels = (size_t)width * height;
    size_t total_volume = (size_t)num_layers * layer_pixels;
    size_t total_out_bytes = (size_t)num_offsets * total_volume * sizeof(uint8_t);

    cudaEvent_t ev_start, ev_h2d_done, ev_jfa_done, ev_thresh_done, ev_d2h_done;
    CUDA_CHECK(cudaEventCreate(&ev_start));
    CUDA_CHECK(cudaEventCreate(&ev_h2d_done));
    CUDA_CHECK(cudaEventCreate(&ev_jfa_done));
    CUDA_CHECK(cudaEventCreate(&ev_thresh_done));
    CUDA_CHECK(cudaEventCreate(&ev_d2h_done));

    // Allocate GPU buffers
    uint8_t* d_masks;
    int* d_seeds_a;
    int* d_seeds_b;
    uint8_t* d_out_offsets;
    float* d_thresholds;

    CUDA_CHECK(cudaMalloc(&d_masks, total_volume * sizeof(uint8_t)));
    CUDA_CHECK(cudaMalloc(&d_seeds_a, total_volume * sizeof(int)));
    CUDA_CHECK(cudaMalloc(&d_seeds_b, total_volume * sizeof(int)));
    CUDA_CHECK(cudaMalloc(&d_out_offsets, total_out_bytes));
    CUDA_CHECK(cudaMalloc(&d_thresholds, num_offsets * sizeof(float)));

    // 1. H2D Transfer
    CUDA_CHECK(cudaEventRecord(ev_start, g_cuda_stream));
    CUDA_CHECK(cudaMemcpyAsync(d_masks, h_masks, total_volume * sizeof(uint8_t),
                               cudaMemcpyHostToDevice, g_cuda_stream));
    CUDA_CHECK(cudaMemcpyAsync(d_thresholds, offset_distances, num_offsets * sizeof(float),
                               cudaMemcpyHostToDevice, g_cuda_stream));
    CUDA_CHECK(cudaEventRecord(ev_h2d_done, g_cuda_stream));

    // 2. Batched 3D JFA Computation
    dim3 block(16, 16, 1);
    dim3 grid((width + 15) / 16, (height + 15) / 16, num_layers);

    batched_jfa_init_kernel<<<grid, block, 0, g_cuda_stream>>>(
        d_masks, d_seeds_a, width, height, num_layers
    );

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

    int* in_seeds = d_seeds_a;
    int* out_seeds = d_seeds_b;

    while (step >= 1) {
        batched_jfa_step_kernel<<<grid, block, 0, g_cuda_stream>>>(
            in_seeds, out_seeds, width, height, num_layers, step
        );
        std::swap(in_seeds, out_seeds);
        step /= 2;
    }

    // JFA+2 and JFA+1 correction passes
    batched_jfa_step_kernel<<<grid, block, 0, g_cuda_stream>>>(
        in_seeds, out_seeds, width, height, num_layers, 2
    );
    std::swap(in_seeds, out_seeds);

    batched_jfa_step_kernel<<<grid, block, 0, g_cuda_stream>>>(
        in_seeds, out_seeds, width, height, num_layers, 1
    );
    std::swap(in_seeds, out_seeds);

    CUDA_CHECK(cudaEventRecord(ev_jfa_done, g_cuda_stream));

    // 3. Multi-Offset Thresholding
    batched_jfa_threshold_kernel<<<grid, block, 0, g_cuda_stream>>>(
        in_seeds, d_masks, d_out_offsets,
        width, height, num_layers,
        pitch_mm, d_thresholds, num_offsets
    );
    CUDA_CHECK(cudaEventRecord(ev_thresh_done, g_cuda_stream));

    // 4. D2H Transfer
    CUDA_CHECK(cudaMemcpyAsync(h_out_masks, d_out_offsets, total_out_bytes,
                               cudaMemcpyDeviceToHost, g_cuda_stream));
    CUDA_CHECK(cudaEventRecord(ev_d2h_done, g_cuda_stream));
    CUDA_CHECK(cudaStreamSynchronize(g_cuda_stream));

    // Timing stats
    if (stats) {
        float h2d_ms = 0, jfa_ms = 0, thresh_ms = 0, d2h_ms = 0;
        cudaEventElapsedTime(&h2d_ms, ev_start, ev_h2d_done);
        cudaEventElapsedTime(&jfa_ms, ev_h2d_done, ev_jfa_done);
        cudaEventElapsedTime(&thresh_ms, ev_jfa_done, ev_thresh_done);
        cudaEventElapsedTime(&d2h_ms, ev_thresh_done, ev_d2h_done);

        stats->t_h2d_ms = h2d_ms;
        stats->t_gpu_jfa_ms = jfa_ms;
        stats->t_gpu_thresh_ms = thresh_ms;
        stats->t_d2h_ms = d2h_ms;
    }

    // Cleanup
    cudaFree(d_masks);
    cudaFree(d_seeds_a);
    cudaFree(d_seeds_b);
    cudaFree(d_out_offsets);
    cudaFree(d_thresholds);

    cudaEventDestroy(ev_start);
    cudaEventDestroy(ev_h2d_done);
    cudaEventDestroy(ev_jfa_done);
    cudaEventDestroy(ev_thresh_done);
    cudaEventDestroy(ev_d2h_done);

    return GPU_OFFSET_SUCCESS;
}

GPU_OFFSET_API int gpu_offset_compute_masks(
    const uint8_t* h_masks,
    int width,
    int height,
    int num_layers,
    float pitch_mm,
    const float* offset_distances,
    int num_offsets,
    uint8_t* h_out_masks,
    GpuOffsetStats* stats
) {
    int eff_backend = g_global_backend;
    if (eff_backend == GPU_OFFSET_BACKEND_AUTO) {
        eff_backend = gpu_offset_has_cuda() ? GPU_OFFSET_BACKEND_CUDA : GPU_OFFSET_BACKEND_CPU;
    }

    if (eff_backend == GPU_OFFSET_BACKEND_CUDA) {
        if (!gpu_offset_has_cuda()) {
            return GPU_OFFSET_ERR_NO_CUDA;
        }
        return gpu_compute_masks_internal(h_masks, width, height, num_layers, pitch_mm,
                                         offset_distances, num_offsets, h_out_masks, stats);
    } else {
        return cpu_offset_compute_masks(h_masks, width, height, num_layers, pitch_mm,
                                        offset_distances, num_offsets, h_out_masks, stats);
    }
}

GPU_OFFSET_API int gpu_offset_polygon_layers(
    const GpuSliceLayer* in_layers,
    size_t num_layers,
    float min_x, float min_y, float max_x, float max_y,
    const GpuOffsetConfig* config,
    GpuLayerResult** out_results,
    GpuOffsetStats* stats
) {
    if (!in_layers || num_layers == 0 || !config || !out_results) return GPU_OFFSET_ERR_INVALID_ARG;

    int eff_backend = (config->backend != GPU_OFFSET_BACKEND_AUTO) ? config->backend : g_global_backend;
    if (eff_backend == GPU_OFFSET_BACKEND_AUTO) {
        eff_backend = gpu_offset_has_cuda() ? GPU_OFFSET_BACKEND_CUDA : GPU_OFFSET_BACKEND_CPU;
    }

    if (eff_backend == GPU_OFFSET_BACKEND_CUDA) {
        if (!gpu_offset_has_cuda()) {
            return GPU_OFFSET_ERR_NO_CUDA;
        }

        // GPU Offsetting pipeline
        auto t_total_start = std::chrono::high_resolution_clock::now();

        float pad = (config->pad_mm > 0.0f) ? config->pad_mm : 1.5f;
        float pitch = (config->pitch_mm > 0.0f) ? config->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);
        size_t total_voxels = (size_t)num_layers * cols * rows;

        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();

        // 1. Rasterization
        auto t_rast_0 = std::chrono::high_resolution_clock::now();
        std::vector<uint8_t> h_masks(total_voxels, 0);

        // Uses CPU polygon rasterizer to populate h_masks
        extern void cpu_rasterize_polygon_layers_internal(
            const GpuSliceLayer* in_layers, size_t num_layers, uint8_t* out_masks,
            int width, int height, float min_x, float min_y, float pitch_mm);

        cpu_rasterize_polygon_layers_internal(in_layers, num_layers, h_masks.data(),
                                              cols, rows, bounds_min_x, bounds_min_y, pitch);
        auto t_rast_1 = std::chrono::high_resolution_clock::now();
        double t_raster_ms = std::chrono::duration<double, std::milli>(t_rast_1 - t_rast_0).count();

        // 2. GPU Distance Transform & Offsetting
        std::vector<uint8_t> h_offset_masks(num_offsets * total_voxels);
        if (gpu_compute_masks_internal(
                h_masks.data(), cols, rows, (int)num_layers, pitch,
                resolved_offsets.data(), (int)num_offsets,
                h_offset_masks.data(), stats) != GPU_OFFSET_SUCCESS) {
            return GPU_OFFSET_ERR_GENERIC;
        }

        if (stats) stats->t_raster_ms = t_raster_ms;

        // 3. Multi-Threaded C++ 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 = in_layers[l].z_height;
                res_array[l].num_perimeters = P;
                res_array[l].perimeters = (GpuOffsetPerimeter*)malloc(P * sizeof(GpuOffsetPerimeter));

                for (int o = 0; o < P; ++o) {
                    size_t offset_layer_idx = (o * num_layers + l) * ((size_t)cols * rows);
                    const uint8_t* mask_ptr = h_offset_masks.data() + offset_layer_idx;

                    std::vector<GpuContour> contours;
                    vectorize_mask_layer(mask_ptr, cols, rows, bounds_min_x, bounds_min_y, pitch,
                                         next_edge_buf, active_edges_buf, visited_buf, contours);

                    res_array[l].perimeters[o].offset_distance_mm = resolved_offsets[o];
                    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
                size_t infill_offset_layer_idx = (P * num_layers + l) * ((size_t)cols * rows);
                const uint8_t* infill_mask_ptr = h_offset_masks.data() + infill_offset_layer_idx;

                std::vector<GpuContour> infill_contours;
                vectorize_mask_layer(infill_mask_ptr, cols, rows, bounds_min_x, bounds_min_y, pitch,
                                     next_edge_buf, active_edges_buf, visited_buf, infill_contours);

                res_array[l].infill.offset_distance_mm = resolved_offsets[P];
                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);
                    memcpy(infill_mask, infill_mask_ptr, (size_t)cols * rows);
                    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;
        }

        *out_results = res_array;
        return GPU_OFFSET_SUCCESS;
    } else {
        return cpu_offset_polygon_layers(in_layers, num_layers, min_x, min_y, max_x, max_y, config, out_results, stats);
    }
}

GPU_OFFSET_API int gpu_offset_slice_mesh_and_offset(
    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,
    float first_layer_height,
    const GpuOffsetConfig* config,
    GpuLayerResult** out_results,
    size_t* out_num_layers,
    GpuOffsetStats* stats
) {
    if (!triangles || num_triangles == 0 || !config || !out_results) return GPU_OFFSET_ERR_INVALID_ARG;

    int eff_backend = (config->backend != GPU_OFFSET_BACKEND_AUTO) ? config->backend : g_global_backend;
    if (eff_backend == GPU_OFFSET_BACKEND_AUTO) {
        eff_backend = gpu_offset_has_cuda() ? GPU_OFFSET_BACKEND_CUDA : GPU_OFFSET_BACKEND_CPU;
    }

    if (eff_backend == GPU_OFFSET_BACKEND_CUDA) {
        if (!gpu_offset_has_cuda()) {
            return GPU_OFFSET_ERR_NO_CUDA;
        }

        if (!g_initialized) {
            if (gpu_offset_init(0) != GPU_OFFSET_SUCCESS) return GPU_OFFSET_ERR_GENERIC;
        }

        float pad = (config->pad_mm > 0.0f) ? config->pad_mm : 1.5f;
        float pitch = (config->pitch_mm > 0.0f) ? config->pitch_mm : 0.10f;
        float h_first = first_layer_height;
        if (h_first <= 0.0f) h_first = config->first_layer_height;
        if (h_first <= 0.0f) h_first = layer_height;

        GpuSdfStackHandle stack = gpu_sdf_from_mesh(
            triangles, num_triangles,
            min_x, min_y, min_z, max_x, max_y, max_z,
            layer_height, h_first,
            pitch, pad
        );
        if (!stack) {
            return GPU_OFFSET_ERR_GENERIC;
        }

        int ret = gpu_sdf_vectorize(stack, config, out_results, out_num_layers, stats);
        gpu_sdf_destroy(stack);
        return ret;
    } else {
        return cpu_offset_slice_mesh_and_offset(triangles, num_triangles, min_x, min_y, min_z, max_x, max_y, max_z,
                                                layer_height, first_layer_height, config, out_results, out_num_layers, stats);
    }
}

GPU_OFFSET_API void gpu_offset_free_results(GpuLayerResult* results, size_t num_layers) {
    if (!results) return;
    for (size_t l = 0; l < num_layers; ++l) {
        if (results[l].perimeters) {
            for (size_t o = 0; o < results[l].num_perimeters; ++o) {
                if (results[l].perimeters[o].contours) {
                    for (size_t c = 0; c < results[l].perimeters[o].num_contours; ++c) {
                        if (results[l].perimeters[o].contours[c].points) {
                            free(results[l].perimeters[o].contours[c].points);
                        }
                    }
                    free(results[l].perimeters[o].contours);
                }
            }
            free(results[l].perimeters);
        }
        if (results[l].infill.contours) {
            for (size_t c = 0; c < results[l].infill.num_contours; ++c) {
                if (results[l].infill.contours[c].points) {
                    free(results[l].infill.contours[c].points);
                }
            }
            free(results[l].infill.contours);
        }
        if (results[l].infill.mask) {
            free(results[l].infill.mask);
        }
    }
    free(results);
}

} // extern "C"