package me.earthme.luminol.collision.backend;

import me.earthme.luminol.collision.CollisionPushResult;
import me.earthme.luminol.collision.INativeCollisionBackend;
import me.earthme.luminol.config.modules.optimizations.AcceleratedCollisionConfig;
import org.jocl.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Arrays;

import static org.jocl.CL.*;

/**
 * GPU backend using OpenCL via JOCL. Provides highest throughput for dense entity clusters.
 * Requires org.jocl:jocl dependency and a compatible GPU driver.
 */
public class GPUCollisionBackend implements INativeCollisionBackend {

    private static final Logger LOGGER = LoggerFactory.getLogger("NekoCore/CollisionGPU");

    // Single-precision broadphase. Consumer GPUs run float many times faster than double,
    // so AABBs are uploaded as float (entity-coordinate precision is ample for overlap tests).
    // Each thread owns entity i and, to avoid reporting a pair twice, only scans j > i.
    // Overlapping pairs are written to two separate index arrays (pairsA/pairsB) so the host
    // can copy them straight into the caller's outA/outB with no de-interleaving.
    private static final String KERNEL_SOURCE =
        "__kernel void detect_collisions(__global const float* aabb,\n" +
        "                                __global int* pairsA, __global int* pairsB,\n" +
        "                                __global volatile int* pairCount,\n" +
        "                                int count, int maxPairs) {\n" +
        "    int i = get_global_id(0);\n" +
        "    if (i >= count) return;\n" +
        "    int b = i*6;\n" +
        "    float minX=aabb[b], minY=aabb[b+1], minZ=aabb[b+2];\n" +
        "    float maxX=aabb[b+3], maxY=aabb[b+4], maxZ=aabb[b+5];\n" +
        "    for (int j = i+1; j < count; j++) {\n" +
        "        int c = j*6;\n" +
        "        if (minX > aabb[c+3] || maxX < aabb[c])   continue;\n" +
        "        if (minY > aabb[c+4] || maxY < aabb[c+1]) continue;\n" +
        "        if (minZ > aabb[c+5] || maxZ < aabb[c+2]) continue;\n" +
        "        int idx = atomic_inc(pairCount);\n" +
        "        if (idx < maxPairs) { pairsA[idx]=i; pairsB[idx]=j; }\n" +
        "    }\n" +
        "}\n";

    private static volatile cl_platform_id platform;
    private static volatile cl_device_id device;
    private static volatile cl_context context;
    private static volatile cl_program program;
    private static volatile boolean initialized = false;

    /**
     * Per-thread reusable OpenCL resources. Folia ticks regions on multiple threads, so a
     * command queue / kernel / device buffers are kept thread-local and only reallocated
     * when capacity must grow, instead of being created and released on every tick.
     */
    private static final class GpuLocal {
        cl_command_queue queue;
        cl_kernel kernel;
        cl_mem aabbMem, pairsAMem, pairsBMem, pairCountMem;
        float[] aabbF = new float[0];   // host-side double->float staging
        int capCount = 0;   // entity capacity of the buffers
        int capPairs = 0;
    }
    private static final ThreadLocal<GpuLocal> LOCAL = ThreadLocal.withInitial(GpuLocal::new);

    @Override public String getName() { return "GPU"; }

    @Override
    public void initialize() throws Exception {
        CL.setExceptionsEnabled(true);
        cl_platform_id[] platforms = new cl_platform_id[1];
        clGetPlatformIDs(1, platforms, null);
        platform = platforms[0];

        cl_device_id[] devices = new cl_device_id[1];
        clGetDeviceIDs(platform, CL_DEVICE_TYPE_GPU, 1, devices, null);
        device = devices[0];

        context = clCreateContext(null, 1, new cl_device_id[]{device}, null, null, null);
        program = clCreateProgramWithSource(context, 1, new String[]{KERNEL_SOURCE}, null, null);
        try {
            clBuildProgram(program, 0, null, null, null, null);
        } catch (CLException e) {
            long[] logSize = new long[1];
            clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, null, logSize);
            byte[] log = new byte[(int)logSize[0]];
            clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, log.length, Pointer.to(log), null);
            throw new RuntimeException("OpenCL build failed: " + new String(log));
        }

        initialized = true;
        LOGGER.info("GPU backend initialized (OpenCL)");
    }

    @Override
    public CollisionPushResult push(double[] aabb, int[] resultSizeOut) {
        if (!initialized) { resultSizeOut[0] = 0; return null; }
        int count = aabb.length / 6;
        if (count == 0) { resultSizeOut[0] = 0; return null; }
        final int maxPairs = count * AcceleratedCollisionConfig.maxCollision;
        final int[] a = new int[maxPairs], b = new int[maxPairs];
        int n = runKernel(aabb, count, maxPairs, a, b);
        if (n <= 0) { resultSizeOut[0] = 0; return null; }
        resultSizeOut[0] = n;
        final int pairs = n;
        return new CollisionPushResult() {
            @Override public int getA(int i) { return a[i]; }
            @Override public int getB(int i) { return b[i]; }
            @Override public float getDensity(int i) { return 0f; }
            @Override public void copyATo(int[] d, int l) { System.arraycopy(a,0,d,0,Math.min(l,pairs)); }
            @Override public void copyBTo(int[] d, int l) { System.arraycopy(b,0,d,0,Math.min(l,pairs)); }
            @Override public void copyDensityTo(float[] d, int l) { }
        };
    }

    @Override
    public int pushInto(double[] aabb, int count, int[] outA, int[] outB) {
        if (!initialized || count == 0) return -1;
        final int maxPairs = Math.min(count * AcceleratedCollisionConfig.maxCollision,
                                      Math.min(outA.length, outB.length));
        return runKernel(aabb, count, maxPairs, outA, outB);
    }

    /**
     * Core OpenCL broadphase: uploads AABBs (as float), runs the kernel, and reads at most
     * {@code maxPairs} overlapping pairs directly into {@code outA}/{@code outB}.
     * Returns the pair count written, or -1 on failure.
     */
    private int runKernel(double[] aabb, int count, int maxPairs, int[] outA, int[] outB) {
        if (maxPairs <= 0) return 0;
        final GpuLocal g = LOCAL.get();
        try {
            if (g.queue == null) {
                g.queue = clCreateCommandQueueWithProperties(context, device, null, null);
                g.kernel = clCreateKernel(program, "detect_collisions", null);
            }
            ensureGpuBuffers(g, count, maxPairs);

            // double -> float staging on host, then upload.
            if (g.aabbF.length < count * 6) g.aabbF = new float[count * 6];
            final float[] aabbF = g.aabbF;
            for (int i = 0, m = count * 6; i < m; i++) aabbF[i] = (float) aabb[i];
            clEnqueueWriteBuffer(g.queue, g.aabbMem, CL_FALSE, 0, (long)count * 6 * Sizeof.cl_float,
                                 Pointer.to(aabbF), 0, null, null);
            int[] pairCountArr = {0};
            clEnqueueWriteBuffer(g.queue, g.pairCountMem, CL_FALSE, 0, Sizeof.cl_int,
                                 Pointer.to(pairCountArr), 0, null, null);

            clSetKernelArg(g.kernel, 0, Sizeof.cl_mem, Pointer.to(g.aabbMem));
            clSetKernelArg(g.kernel, 1, Sizeof.cl_mem, Pointer.to(g.pairsAMem));
            clSetKernelArg(g.kernel, 2, Sizeof.cl_mem, Pointer.to(g.pairsBMem));
            clSetKernelArg(g.kernel, 3, Sizeof.cl_mem, Pointer.to(g.pairCountMem));
            clSetKernelArg(g.kernel, 4, Sizeof.cl_int, Pointer.to(new int[]{count}));
            clSetKernelArg(g.kernel, 5, Sizeof.cl_int, Pointer.to(new int[]{maxPairs}));

            clEnqueueNDRangeKernel(g.queue, g.kernel, 1, null, new long[]{count}, null, 0, null, null);

            clEnqueueReadBuffer(g.queue, g.pairCountMem, CL_TRUE, 0, Sizeof.cl_int, Pointer.to(pairCountArr), 0, null, null);
            int pairCount = Math.min(pairCountArr[0], maxPairs);
            if (pairCount <= 0) return 0;

            clEnqueueReadBuffer(g.queue, g.pairsAMem, CL_TRUE, 0, (long)pairCount*Sizeof.cl_int, Pointer.to(outA), 0, null, null);
            clEnqueueReadBuffer(g.queue, g.pairsBMem, CL_TRUE, 0, (long)pairCount*Sizeof.cl_int, Pointer.to(outB), 0, null, null);
            return pairCount;
        } catch (Exception e) {
            LOGGER.error("GPU kernel failed", e);
            releaseLocal(g);
            return -1;
        }
    }

    /** (Re)allocate device buffers when the live entity/pair count exceeds current capacity. */
    private void ensureGpuBuffers(GpuLocal g, int count, int maxPairs) {
        if (count > g.capCount || g.aabbMem == null) {
            if (g.aabbMem != null)    clReleaseMemObject(g.aabbMem);
            if (g.pairCountMem != null) clReleaseMemObject(g.pairCountMem);
            g.aabbMem    = clCreateBuffer(context, CL_MEM_READ_ONLY,  (long)count * 6 * Sizeof.cl_float, null, null);
            g.pairCountMem = clCreateBuffer(context, CL_MEM_READ_WRITE, Sizeof.cl_int, null, null);
            g.capCount = count;
        }
        if (maxPairs > g.capPairs || g.pairsAMem == null) {
            if (g.pairsAMem != null) clReleaseMemObject(g.pairsAMem);
            if (g.pairsBMem != null) clReleaseMemObject(g.pairsBMem);
            g.pairsAMem = clCreateBuffer(context, CL_MEM_WRITE_ONLY, (long)maxPairs * Sizeof.cl_int, null, null);
            g.pairsBMem = clCreateBuffer(context, CL_MEM_WRITE_ONLY, (long)maxPairs * Sizeof.cl_int, null, null);
            g.capPairs = maxPairs;
        }
    }

    private static void releaseLocal(GpuLocal g) {
        try { if (g.aabbMem != null) clReleaseMemObject(g.aabbMem); } catch (Throwable ignored) {}
        try { if (g.pairsAMem != null) clReleaseMemObject(g.pairsAMem); } catch (Throwable ignored) {}
        try { if (g.pairsBMem != null) clReleaseMemObject(g.pairsBMem); } catch (Throwable ignored) {}
        try { if (g.pairCountMem != null) clReleaseMemObject(g.pairCountMem); } catch (Throwable ignored) {}
        try { if (g.kernel != null) clReleaseKernel(g.kernel); } catch (Throwable ignored) {}
        try { if (g.queue != null) clReleaseCommandQueue(g.queue); } catch (Throwable ignored) {}
        g.aabbMem = g.pairsAMem = g.pairsBMem = g.pairCountMem = null;
        g.kernel = null; g.queue = null; g.capCount = 0; g.capPairs = 0;
    }

    @Override public void applyConfig() {}

    @Override
    public void destroy() {
        if (!initialized) return;
        initialized = false;
        if (program != null) clReleaseProgram(program);
        if (context != null) clReleaseContext(context);
    }
}
