package ai.nomyo;

import lombok.Getter;
import lombok.Setter;

import java.lang.foreign.*;
import java.lang.invoke.MethodHandle;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.Map;

/**
 * Cross-platform memory locking and secure zeroing for sensitive cryptographic buffers. Fails gracefully if unavailable.
 */
@SuppressWarnings("SameReturnValue")
public final class SecureMemory {

    @Getter
    private static final boolean HAS_MEMORY_LOCKING;
    @Getter
    private static final boolean HAS_SECURE_ZEROING;
    @Getter
    @Setter
    private static volatile boolean secureMemoryEnabled = true;

    private static final MethodHandle MLOCK_HANDLE;
    private static final MethodHandle MUNLOCK_HANDLE;

    static {
        boolean locking = false;
        boolean zeroing = false;
        MethodHandle mlockHandle = null;
        MethodHandle munlockHandle = null;

        try {
            Linker linker = Linker.nativeLinker();
            SymbolLookup stdLib = linker.defaultLookup();

            var mlockOpt = stdLib.find("mlock");
            var munlockOpt = stdLib.find("munlock");

            if (mlockOpt.isPresent() && munlockOpt.isPresent()) {
                MemorySegment mlockAddr = mlockOpt.get();
                MemorySegment munlockAddr = munlockOpt.get();
                FunctionDescriptor mlockDesc = FunctionDescriptor.of(
                        ValueLayout.JAVA_INT,
                        ValueLayout.ADDRESS,
                        ValueLayout.JAVA_LONG);

                FunctionDescriptor munlockDesc = FunctionDescriptor.of(
                        ValueLayout.JAVA_INT,
                        ValueLayout.ADDRESS,
                        ValueLayout.JAVA_LONG);

                mlockHandle = linker.downcallHandle(mlockAddr, mlockDesc);
                munlockHandle = linker.downcallHandle(munlockAddr, munlockDesc);

                locking = true;
            }
        } catch (Throwable t) {
            // Degrade gracefully
        }

        MLOCK_HANDLE = mlockHandle;
        MUNLOCK_HANDLE = munlockHandle;

        try {
            locking = initMemoryLocking(locking);
            zeroing = true; // Secure zeroing is always available at the JVM level
        } catch (Throwable t) {
            // Degrade gracefully
        }

        HAS_MEMORY_LOCKING = locking;
        HAS_SECURE_ZEROING = zeroing;
    }

    private static boolean initMemoryLocking(boolean preCheck) {
        if (MLOCK_HANDLE == null || MUNLOCK_HANDLE == null || !preCheck) {
            return false;
        }

        try (Arena arena = Arena.ofConfined()) {
            MemorySegment testSegment = arena.allocate(1);
            long result = (long) MLOCK_HANDLE.invokeExact(testSegment, (long) 1);
            return result == 0;
        } catch (Throwable t) {
            return false;
        }
    }

    /**
     * Recommended way to handle sensitive data — use within try-with-resources for secure zeroing.
     */
    public static SecureBuffer secureByteArray(byte[] data, boolean lock) {
        return new SecureBuffer(data, lock);
    }

    /**
     * Always attempts locking.
     */
    public static SecureBuffer secureByteArray(byte[] data) {
        return secureByteArray(data, true);
    }

    /**
     * Returns protection capabilities: enabled, protection_level, has_memory_locking, has_secure_zeroing, supports_full_protection, page_size.
     */
    public static Map<String, Object> getMemoryProtectionInfo() {
        String protectionLevel;
        if (HAS_MEMORY_LOCKING && HAS_SECURE_ZEROING) {
            protectionLevel = "full";
        } else if (HAS_SECURE_ZEROING) {
            protectionLevel = "zeroing_only";
        } else {
            protectionLevel = "none";
        }

        boolean supportsFull = HAS_MEMORY_LOCKING && HAS_SECURE_ZEROING && secureMemoryEnabled;

        return Map.of("enabled", secureMemoryEnabled, "protection_level", protectionLevel, "has_memory_locking", HAS_MEMORY_LOCKING, "has_secure_zeroing", HAS_SECURE_ZEROING, "supports_full_protection", supportsFull, "page_size", Constants.PAGE_SIZE);
    }

    /**
     * Wraps bytes with memory locking and guaranteed zeroing on close. AutoCloseable for try-with-resources.
     */
    @SuppressWarnings("SameReturnValue")
    public static class SecureBuffer implements AutoCloseable {

        private final Arena arena;

        @Getter
        private final MemorySegment data;

        @Getter
        private final long size;

        @Getter
        private final long address;
        private boolean locked;
        private boolean closed;

        /**
         * @param data byte array to wrap (zeroed after copy to off-heap memory)
         * @param lock whether to attempt memory locking
         */
        public SecureBuffer(byte[] data, boolean lock) {
            this.arena = Arena.ofConfined();
            this.data = data != null ? this.arena.allocate(data.length) : MemorySegment.NULL;

            if (data != null) {
                this.data.asByteBuffer().put(data);
                Arrays.fill(data, (byte) 0);
            }

            this.size = this.data.byteSize();
            this.address = this.data.address();

            this.locked = false;
            this.closed = false;

            if (lock && SecureMemory.secureMemoryEnabled) {
                this.locked = lock();
            }
        }

        /**
         * Locks buffer in memory (prevents disk swapping). Returns false if unavailable.
         */
        public boolean lock() {
            if (this.locked || this.address == 0) {
                return this.locked;
            }

            try {
                long result = (long) MLOCK_HANDLE.invokeExact(
                        MemorySegment.ofAddress(this.address),
                        this.size);
                this.locked = result == 0;
            } catch (Throwable t) {
                this.locked = false;
            }

            return this.locked;
        }

        /**
         * Unlocks buffer (allows disk swapping).
         */
        public boolean unlock() {
            if (!locked || this.address == 0) {
                return false;
            }

            try {
                long result = (long) MUNLOCK_HANDLE.invokeExact(
                        MemorySegment.ofAddress(this.address),
                        this.size);
                locked = result != 0;
                return result == 0;
            } catch (Throwable t) {
                locked = true;
                return false;
            }
        }

        /**
         * Securely zeros buffer contents.
         */
        public void zero() {
            if (data != null) {
                data.fill((byte) 0);
            }
        }

        @Override
        public void close() {
            if (closed) return;

            zero();
            unlock();

            arena.close();
            closed = true;
        }
    }
}