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.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.lang.foreign.*;
import java.lang.invoke.MethodHandle;
import java.nio.file.*;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;

import static java.lang.foreign.ValueLayout.*;

/**
 * FFM backend: calls the AcceleratedRecoiling C++ shared library via java.lang.foreign.
 * Requires Java 21+ and the native library placed in /natives/{os}-{arch}/ on classpath.
 */
public class FFMCollisionBackend implements INativeCollisionBackend {

    private static final Logger LOGGER = LoggerFactory.getLogger("NekoCore/CollisionFFM");
    private static final AtomicLong maxSizeSeen = new AtomicLong(-1);

    private static Linker linker;
    private static Arena nativeArena;
    private static MethodHandle hPush, hCreateCtx, hDestroyCtx, hCreateCfg, hUpdateCfg, hDestroyCfg;
    private static volatile boolean initialized = false;

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

    private static class ThreadState {
        final Arena arena = Arena.ofConfined();
        long ctx;
        long cfg;
        MemorySegment aabbBuf;
        MemorySegment pairBufA;
        MemorySegment pairBufB;
        MemorySegment densityBuf;
        int bufCapacity = 0;
    }

    private static final ThreadLocal<ThreadState> THREAD_STATE = ThreadLocal.withInitial(ThreadState::new);

    @Override
    public void initialize() throws Exception {
        String osName = System.getProperty("os.name", "").toLowerCase();
        String arch    = System.getProperty("os.arch", "").toLowerCase();
        String[] platforms;
        String[] libraryNames;

        if (osName.contains("win")) {
            if (arch.contains("aarch")) {
                platforms = new String[]{"windows-aarch64"};
            } else {
                // Compatible with both x86_64 and x64 package layouts.
                platforms = new String[]{"windows-x86_64", "windows-x64"};
            }
            libraryNames = new String[]{"AcceleratedRecoiling.dll", "acceleratedRecoilingLib.dll"};
        } else if (osName.contains("linux")) {
            if (arch.contains("aarch")) {
                platforms = new String[]{"linux-aarch64"};
            } else {
                // Compatible with both x86_64 and x64 package layouts.
                platforms = new String[]{"linux-x86_64", "linux-x64"};
            }
            libraryNames = new String[]{"AcceleratedRecoiling.so", "libAcceleratedRecoiling.so", "acceleratedRecoilingLib.so"};
        } else if (osName.contains("mac")) {
            platforms = arch.contains("aarch") ? new String[]{"macos-aarch64"} : new String[]{"macos-x86_64"};
            libraryNames = new String[]{"AcceleratedRecoiling.dylib", "libAcceleratedRecoiling.dylib"};
        } else {
            throw new UnsupportedOperationException("Unsupported OS: " + osName + ", arch: " + arch);
        }

        InputStream is = null;
        String selectedResourcePath = null;
        String selectedLibName = null;
        for (String platform : platforms) {
            for (String libName : libraryNames) {
                String candidate = "/natives/" + platform + "/" + libName;
                InputStream candidateStream = FFMCollisionBackend.class.getResourceAsStream(candidate);
                if (candidateStream != null) {
                    is = candidateStream;
                    selectedResourcePath = candidate;
                    selectedLibName = libName;
                    break;
                }
            }
            if (is != null) {
                break;
            }
        }

        if (is == null) {
            throw new FileNotFoundException(
                "Native library not found in classpath for OS=" + osName + ", arch=" + arch +
                ", tried platforms=" + String.join(",", platforms) +
                ", names=" + String.join(",", libraryNames)
            );
        }

        String tempSuffix = selectedLibName.endsWith(".dll") ? ".dll"
            : (selectedLibName.endsWith(".dylib") ? ".dylib" : ".so");
        Path tempLib = Files.createTempFile("AccelRecoil_" + UUID.randomUUID(), tempSuffix);
        tempLib.toFile().deleteOnExit();
        final InputStream selectedStream = is;
        try (selectedStream) {
            Files.copy(selectedStream, tempLib, StandardCopyOption.REPLACE_EXISTING);
        }

        linker     = Linker.nativeLinker();
        nativeArena = Arena.ofAuto();
        SymbolLookup lib = SymbolLookup.libraryLookup(tempLib, nativeArena);

        FunctionDescriptor pushDesc = FunctionDescriptor.of(JAVA_INT, ADDRESS, ADDRESS, ADDRESS, JAVA_INT, ADDRESS, JAVA_LONG, JAVA_LONG);
        FunctionDescriptor ctxDesc  = FunctionDescriptor.of(JAVA_LONG);
        FunctionDescriptor cfgDesc  = FunctionDescriptor.of(JAVA_LONG, JAVA_INT, JAVA_INT, JAVA_INT, JAVA_INT);
        FunctionDescriptor updDesc  = FunctionDescriptor.ofVoid(JAVA_LONG, JAVA_INT, JAVA_INT, JAVA_INT, JAVA_INT);
        FunctionDescriptor dCtxDesc = FunctionDescriptor.ofVoid(JAVA_LONG);
        FunctionDescriptor dCfgDesc = FunctionDescriptor.ofVoid(JAVA_LONG);

        hPush       = linker.downcallHandle(findSymbol(lib, "push", "push_entities"), pushDesc);
        hCreateCtx  = linker.downcallHandle(findSymbol(lib, "createCtx", "create_context"), ctxDesc);
        hCreateCfg  = linker.downcallHandle(findSymbol(lib, "createCfg", "create_config"), cfgDesc);
        hUpdateCfg  = linker.downcallHandle(findSymbol(lib, "updateCfg", "update_config"), updDesc);
        hDestroyCtx = linker.downcallHandle(findSymbol(lib, "destroyCtx", "destroy_context"), dCtxDesc);
        hDestroyCfg = linker.downcallHandle(findSymbol(lib, "destroyCfg", "destroy_config"), dCfgDesc);

        initialized = true;
        LOGGER.info("FFM backend initialized from {} (resource: {})", tempLib, selectedResourcePath);
    }

    @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; }

        ThreadState ts = THREAD_STATE.get();
        ensureBuffers(ts, count);

        // Bulk-copy AABB data (much cheaper than per-element setAtIndex on the hot path).
        MemorySegment.copy(aabb, 0, ts.aabbBuf, JAVA_DOUBLE, 0L, aabb.length);

        try {
            int pairCount = (int) hPush.invoke(
                ts.aabbBuf,
                ts.pairBufA,
                ts.pairBufB,
                count,
                ts.densityBuf,
                ts.ctx,
                ts.cfg
            );
            resultSizeOut[0] = pairCount;
            if (pairCount == 0) return null;

            int[] arrA = new int[pairCount], arrB = new int[pairCount];
            float[] dens = new float[count];
            // Bulk-copy native results back into Java arrays.
            MemorySegment.copy(ts.pairBufA, JAVA_INT, 0L, arrA, 0, pairCount);
            MemorySegment.copy(ts.pairBufB, JAVA_INT, 0L, arrB, 0, pairCount);
            MemorySegment.copy(ts.densityBuf, JAVA_FLOAT, 0L, dens, 0, count);

            return new me.earthme.luminol.collision.CollisionPushResult() {
                @Override public int getA(int idx) { return arrA[idx]; }
                @Override public int getB(int idx) { return arrB[idx]; }
                @Override public float getDensity(int idx) { return dens[idx]; }
                @Override public void copyATo(int[] dest, int len) { System.arraycopy(arrA,0,dest,0,len); }
                @Override public void copyBTo(int[] dest, int len) { System.arraycopy(arrB,0,dest,0,len); }
                @Override public void copyDensityTo(float[] dest, int len) { System.arraycopy(dens,0,dest,0,len); }
            };
        } catch (Throwable e) {
            LOGGER.error("FFM push failed", e);
            resultSizeOut[0] = 0;
            return null;
        }
    }

    @Override
    public int pushInto(double[] aabb, int count, int[] outA, int[] outB) {
        if (!initialized || count == 0) return -1;

        ThreadState ts = THREAD_STATE.get();
        ensureBuffers(ts, count);

        // Bulk-copy AABBs into the native input buffer.
        MemorySegment.copy(aabb, 0, ts.aabbBuf, JAVA_DOUBLE, 0L, count * 6);

        try {
            int pairCount = (int) hPush.invoke(
                ts.aabbBuf,
                ts.pairBufA,
                ts.pairBufB,
                count,
                ts.densityBuf,   // native writes density here; we never read it back
                ts.ctx,
                ts.cfg
            );
            if (pairCount <= 0) return 0;
            int n = Math.min(pairCount, Math.min(outA.length, outB.length));
            // Bulk-copy pair indices straight into the caller's reusable arrays. No density,
            // no anonymous result object, no per-tick int[] allocation.
            MemorySegment.copy(ts.pairBufA, JAVA_INT, 0L, outA, 0, n);
            MemorySegment.copy(ts.pairBufB, JAVA_INT, 0L, outB, 0, n);
            return n;
        } catch (Throwable e) {
            LOGGER.error("FFM pushInto failed", e);
            return -1;
        }
    }

    private void ensureBuffers(ThreadState ts, int count) {
        if (count > ts.bufCapacity) {
            ts.aabbBuf   = ts.arena.allocate(JAVA_DOUBLE, (long)count * 6);
            int pairBufSize = Math.toIntExact((long) count * AcceleratedCollisionConfig.maxCollision);
            ts.pairBufA  = ts.arena.allocate(JAVA_INT, pairBufSize);
            ts.pairBufB  = ts.arena.allocate(JAVA_INT, pairBufSize);
            ts.densityBuf= ts.arena.allocate(JAVA_FLOAT,  count);
            if (ts.ctx != 0) { try { hDestroyCtx.invoke(ts.ctx); } catch (Throwable ignored) {} }
            if (ts.cfg != 0) { try { hDestroyCfg.invoke(ts.cfg); } catch (Throwable ignored) {} }
            try {
                ts.ctx = (long) hCreateCtx.invoke();
                ts.cfg = (long) hCreateCfg.invoke(
                    AcceleratedCollisionConfig.maxCollision,
                    AcceleratedCollisionConfig.gridSize,
                    AcceleratedCollisionConfig.densityWindow,
                    AcceleratedCollisionConfig.maxThreads
                );
            } catch (Throwable e) { LOGGER.error("Failed to create FFM context", e); }
            ts.bufCapacity = count;
        }
    }

    @Override
    public void applyConfig() {
        // Per-thread config update handled in push()
    }

    @Override
    public void destroy() {
        initialized = false;
        if (nativeArena != null) { try { nativeArena.close(); } catch (Exception ignored) {} }
    }

    private static MemorySegment findSymbol(SymbolLookup lib, String... candidates) {
        for (String name : candidates) {
            var found = lib.find(name);
            if (found.isPresent()) {
                return found.get();
            }
        }
        throw new IllegalStateException("Missing native symbol, tried: " + String.join(", ", candidates));
    }
}
