package ai.nomyo;

import lombok.Getter;
import lombok.Setter;

import java.lang.foreign.Arena;
import java.lang.foreign.MemorySegment;
import java.util.Map;

/**
 * Cross-platform memory locking and secure zeroing for sensitive cryptographic buffers. Fails gracefully if unavailable.
 */
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;

    static {
        boolean locking = false;
        boolean zeroing = false;

        try {
            locking = initMemoryLocking();
            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() {
        // FFM doesn't support memory locking at this point in time
        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.
     */
    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
         * @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);
            }

            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() {
            return false;
        }

        /**
         * Unlocks buffer (allows disk swapping).
         */
        public boolean unlock() {
            if (!locked) return false;
            locked = false;
            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;
        }
    }
}