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 utilities.
 *
 * <p>This module provides optional memory protection for sensitive
 * cryptographic buffers. It fails gracefully if memory locking is
 * unavailable on the current platform (e.g., Windows on some JVM
 * configurations).</p>
 *
 * <h3>Protection Levels</h3>
 * <ul>
 *   <li><b>"full"</b> — Memory locking and secure zeroing both available</li>
 *   <li><b>"zeroing_only"</b> — Only secure zeroing available</li>
 *   <li><b>"none"</b> — No memory protection available</li>
 * </ul>
 *
 * <h3>Usage</h3>
 * <pre>{@code
 * try (SecureBuffer buf = SecureMemory.secureBytearray(sensitiveData)) {
 *     // buf is locked and ready to use
 *     process(buf.getData());
 * }
 * // buf is automatically zeroed and unlocked on exit
 * }</pre>
 */
public final class SecureMemory {

    @Getter
    @Setter
    private static volatile boolean secureMemoryEnabled = true;

    @Getter
    private static final boolean HAS_MEMORY_LOCKING;
    @Getter
    private static final boolean HAS_SECURE_ZEROING;

    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
        // TODO: Bypass this with native libraries
        return false;
    }

    /**
     * Wraps a byte array with memory locking and guaranteed zeroing on exit.
     *
     * <p>Implements {@link AutoCloseable} for use with try-with-resources.
     * The buffer is automatically zeroed and unlocked when closed, even
     * if an exception occurs.</p>
     */
    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;

        /**
         * Creates a new SecureBuffer wrapping the given data.
         *
         * @param data     the 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();
            }
        }

        /**
         * Attempts to lock the buffer in memory, preventing swapping to disk.
         *
         * @return {@code true} if locking succeeded, {@code false} otherwise
         */
        public boolean lock() {
            //data = data.asReadOnly();
            return false;
        }

        /**
         * Unlocks the buffer, allowing it to be swapped to disk.
         *
         * @return {@code true} if unlocking succeeded, {@code false} otherwise
         */
        public boolean unlock() {
            if (!locked) return false;
            locked = false;
            return false;
        }

        /**
         * Securely zeros the 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;
        }
    }

    /**
     * Creates a SecureBuffer for the given data with memory locking.
     *
     * <p>This is the recommended way to handle sensitive data. The returned
     * buffer should be used within a try-with-resources block to ensure
     * secure zeroing on exit.</p>
     *
     * @param data   the sensitive data bytes
     * @param lock   whether to attempt memory locking
     * @return a new SecureBuffer
     */
    public static SecureBuffer secureByteArray(byte[] data, boolean lock) {
        return new SecureBuffer(data, lock);
    }

    /**
     * Creates a SecureBuffer for the given data with memory locking.
     * Convenience variant that always attempts locking.
     *
     * @param data the sensitive data bytes
     * @return a new SecureBuffer
     */
    public static SecureBuffer secureByteArray(byte[] data) {
        return secureByteArray(data, true);
    }

    /**
     * Creates a SecureBuffer for the given data with memory locking.
     *
     * <p><b>Deprecated:</b> Use {@link #secureByteArray(byte[])} instead.
     * This method exists for compatibility with the Python reference.</p>
     *
     * @param data   the sensitive data bytes
     * @param lock   whether to attempt memory locking
     * @return a new SecureBuffer
     * @deprecated Use {@link #secureByteArray(byte[])} instead
     */
    @Deprecated
    public static SecureBuffer secureBytes(byte[] data, boolean lock) {
        return new SecureBuffer(data, lock);
    }

    /**
     * Creates a SecureBuffer for the given data with memory locking.
     *
     * <p><b>Deprecated:</b> Use {@link #secureByteArray(byte[])} instead.
     * This method exists for compatibility with the Python reference.</p>
     *
     * @param data the sensitive data bytes
     * @return a new SecureBuffer
     * @deprecated Use {@link #secureByteArray(byte[])} instead
     */
    @Deprecated
    public static SecureBuffer secureBytes(byte[] data) {
        return secureBytes(data, true);
    }

    /**
     * Returns information about the current memory protection capabilities.
     *
     * @return a map of protection capabilities
     */
    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
        );
    }
}