/*
 * Copyright 2004-2023 H2 Group. Multiple-Licensed under the MPL 2.0,
 * and the EPL 1.0 (https://h2database.com/html/license.html).
 * Initial Developer: H2 Group
 */
package org.h2.mvstore;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.IntConsumer;

/**
 * A facility to store streams in a map. Streams are split into blocks, which
 * are stored in a map. Very small streams are inlined in the stream id.
 * <p>
 * The key of the map is a long (incremented for each stored block). The default
 * initial value is 0. Before storing blocks into the map, the stream store
 * checks if there is already a block with the next key, and if necessary
 * searches the next free entry using a binary search (0 to Long.MAX_VALUE).
 * <p>
 * Before storing
 * <p>
 * The format of the binary id is: An empty id represents 0 bytes of data.
 * In-place data is encoded as 0, the size (a variable size int), then the data.
 * A stored block is encoded as 1, the length of the block (a variable size
 * int), then the key (a variable size long). Multiple ids can be concatenated
 * to concatenate the data. If the id is large, it is stored itself, which is
 * encoded as 2, the total length (a variable size long), and the key of the
 * block that contains the id (a variable size long).
 */
public final class StreamStore
{
    private final Map<Long, byte[]> map;
    private final int minBlockSize;
    private final int maxBlockSize;
    private final AtomicLong nextKey = new AtomicLong();
    private final AtomicReference<byte[]> nextBuffer = new AtomicReference<>();
    private final IntConsumer onStoreCallback;

    /**
     * Create a stream store instance.
     *
     * @param map the map to store blocks of data
     */
    public StreamStore(Map<Long, byte[]> map) {
        this(map, 256, 256 * 1024, null);
    }

    public StreamStore(Map<Long, byte[]> map, int minBlockSize, int maxBlockSize) {
        this(map, minBlockSize, maxBlockSize, null);
    }

    public StreamStore(Map<Long, byte[]> map, IntConsumer onStoreCallback) {
        this(map, 256, 256 * 1024, onStoreCallback);
    }

    public StreamStore(Map<Long, byte[]> map, int minBlockSize, int maxBlockSize, IntConsumer onStoreCallback) {
        this.map = map;
        this.minBlockSize = minBlockSize;
        this.maxBlockSize = maxBlockSize;
        this.onStoreCallback = onStoreCallback;
    }

    public Map<Long, byte[]> getMap() {
        return map;
    }

    public void setNextKey(long nextKey) {
        this.nextKey.set(nextKey);
    }

    public long getNextKey() {
        return nextKey.get();
    }

    public int getMinBlockSize() {
        return minBlockSize;
    }

    public long getMaxBlockSize() {
        return maxBlockSize;
    }

    /**
     * Store the stream, and return the id. The stream is not closed.
     *
     * @param in the stream
     * @return the id (potentially an empty array)
     * @throws IOException If an I/O error occurs
     */
    public byte[] put(InputStream in) throws IOException {
        ByteArrayOutputStream id = new ByteArrayOutputStream();
        int level = 0;
        try {
            while (!put(id, in, level)) {
                if (id.size() > maxBlockSize / 2) {
                    id = putIndirectId(id);
                    level++;
                }
            }
        } catch (IOException e) {
            remove(id.toByteArray());
            throw e;
        }
        if (id.size() > minBlockSize * 2) {
            id = putIndirectId(id);
        }
        return id.toByteArray();
    }

    private boolean put(ByteArrayOutputStream id, InputStream in, int level)
            throws IOException {
        if (level > 0) {
            ByteArrayOutputStream id2 = new ByteArrayOutputStream();
            while (true) {
                boolean eof = put(id2, in, level - 1);
                if (id2.size() > maxBlockSize / 2) {
                    id2 = putIndirectId(id2);
                    id2.writeTo(id);
                    return eof;
                } else if (eof) {
                    id2.writeTo(id);
                    return true;
                }
            }
        }
        byte[] readBuffer = nextBuffer.getAndSet(null);
        if (readBuffer == null) {
            readBuffer = new byte[maxBlockSize];
        }
        byte[] buff = read(in, readBuffer);
        if (buff != readBuffer) {
            // re-use the buffer if the result was shorter
            nextBuffer.set(readBuffer);
        }
        int len = buff.length;
        if (len == 0) {
            return true;
        }
        boolean eof = len < maxBlockSize;
        if (len < minBlockSize) {
            // in-place: 0, len (int), data
            id.write(0);
            DataUtils.writeVarInt(id, len);
            id.write(buff);
        } else {
            // block: 1, len (int), blockId (long)
            id.write(1);
            DataUtils.writeVarInt(id, len);
            DataUtils.writeVarLong(id, writeBlock(buff));
        }
        return eof;
    }

    private static byte[] read(InputStream in, byte[] target)
            throws IOException {
        int copied = 0;
        int remaining = target.length;
        while (remaining > 0) {
            try {
                int len = in.read(target, copied, remaining);
                if (len < 0) {
                    return Arrays.copyOf(target, copied);
                }
                copied += len;
                remaining -= len;
            } catch (RuntimeException e) {
                throw new IOException(e);
            }
        }
        return target;
    }

    private ByteArrayOutputStream putIndirectId(ByteArrayOutputStream id)
            throws IOException {
        byte[] data = id.toByteArray();
        id = new ByteArrayOutputStream();
        // indirect: 2, total len (long), blockId (long)
        id.write(2);
        DataUtils.writeVarLong(id, length(data));
        DataUtils.writeVarLong(id, writeBlock(data));
        return id;
    }

    private long writeBlock(byte[] data) {
        long key = getAndIncrementNextKey();
        map.put(key, data);
        if (onStoreCallback != null) {
            onStoreCallback.accept(data.length);
        }
        return key;
    }

    /**
     * Generate a new key.
     *
     * @return the new key
     */
    private long getAndIncrementNextKey() {
        long key = nextKey.getAndIncrement();
        if (!map.containsKey(key)) {
            return key;
        }
        // search the next free id using binary search
        synchronized (this) {
            long low = key, high = Long.MAX_VALUE;
            while (low < high) {
                long x = (low + high) >>> 1;
                if (map.containsKey(x)) {
                    low = x + 1;
                } else {
                    high = x;
                }
            }
            key = low;
            nextKey.set(key + 1);
            return key;
        }
    }

    /**
     * Get the key of the biggest block, of -1 for inline data.
     * This method is used to garbage collect orphaned blocks.
     *
     * @param id the id
     * @return the key, or -1
     */
    public long getMaxBlockKey(byte[] id) {
        long maxKey = -1;
        ByteBuffer idBuffer = ByteBuffer.wrap(id);
        while (idBuffer.hasRemaining()) {
            switch (idBuffer.get()) {
            case 0:
                // in-place: 0, len (int), data
                int len = DataUtils.readVarInt(idBuffer);
                idBuffer.position(idBuffer.position() + len);
                break;
            case 1:
                // block: 1, len (int), blockId (long)
                DataUtils.readVarInt(idBuffer);
                long k = DataUtils.readVarLong(idBuffer);
                maxKey = Math.max(maxKey, k);
                break;
            case 2:
                // indirect: 2, total len (long), blockId (long)
                DataUtils.readVarLong(idBuffer);
                long k2 = DataUtils.readVarLong(idBuffer);
                maxKey = k2;
                byte[] r = map.get(k2);
                // recurse
                long m = getMaxBlockKey(r);
                if (m >= 0) {
                    maxKey = Math.max(maxKey, m);
                }
                break;
            default:
                throw DataUtils.newIllegalArgumentException(
                        "Unsupported id {0}", Arrays.toString(id));
            }
        }
        return maxKey;
    }

    /**
     * Remove all stored blocks for the given id.
     *
     * @param id the id
     */
    public void remove(byte[] id) {
        ByteBuffer idBuffer = ByteBuffer.wrap(id);
        while (idBuffer.hasRemaining()) {
            switch (idBuffer.get()) {
            case 0:
                // in-place: 0, len (int), data
                int len = DataUtils.readVarInt(idBuffer);
                idBuffer.position(idBuffer.position() + len);
                break;
            case 1:
                // block: 1, len (int), blockId (long)
                DataUtils.readVarInt(idBuffer);
                long k = DataUtils.readVarLong(idBuffer);
                map.remove(k);
                break;
            case 2:
                // indirect: 2, total len (long), blockId (long)
                DataUtils.readVarLong(idBuffer);
                long k2 = DataUtils.readVarLong(idBuffer);
                // recurse
                remove(map.get(k2));
                map.remove(k2);
                break;
            default:
                throw DataUtils.newIllegalArgumentException(
                        "Unsupported id {0}", Arrays.toString(id));
            }
        }
    }

    /**
     * Convert the id to a human readable string.
     *
     * @param id the stream id
     * @return the string
     */
    public static String toString(byte[] id) {
        StringBuilder buff = new StringBuilder();
        ByteBuffer idBuffer = ByteBuffer.wrap(id);
        long length = 0;
        while (idBuffer.hasRemaining()) {
            long block;
            int len;
            switch (idBuffer.get()) {
            case 0:
                // in-place: 0, len (int), data
                len = DataUtils.readVarInt(idBuffer);
                idBuffer.position(idBuffer.position() + len);
                buff.append("data len=").append(len);
                length += len;
                break;
            case 1:
                // block: 1, len (int), blockId (long)
                len = DataUtils.readVarInt(idBuffer);
                length += len;
                block = DataUtils.readVarLong(idBuffer);
                buff.append("block ").append(block).append(" len=").append(len);
                break;
            case 2:
                // indirect: 2, total len (long), blockId (long)
                len = DataUtils.readVarInt(idBuffer);
                length += DataUtils.readVarLong(idBuffer);
                block = DataUtils.readVarLong(idBuffer);
                buff.append("indirect block ").append(block).append(" len=").append(len);
                break;
            default:
                buff.append("error");
            }
            buff.append(", ");
        }
        buff.append("length=").append(length);
        return buff.toString();
    }

    /**
     * Calculate the number of data bytes for the given id. As the length is
     * encoded in the id, this operation does not cause any reads in the map.
     *
     * @param id the id
     * @return the length
     */
    public long length(byte[] id) {
        ByteBuffer idBuffer = ByteBuffer.wrap(id);
        long length = 0;
        while (idBuffer.hasRemaining()) {
            switch (idBuffer.get()) {
            case 0:
                // in-place: 0, len (int), data
                int len = DataUtils.readVarInt(idBuffer);
                idBuffer.position(idBuffer.position() + len);
                length += len;
                break;
            case 1:
                // block: 1, len (int), blockId (long)
                length += DataUtils.readVarInt(idBuffer);
                DataUtils.readVarLong(idBuffer);
                break;
            case 2:
                // indirect: 2, total len (long), blockId (long)
                length += DataUtils.readVarLong(idBuffer);
                DataUtils.readVarLong(idBuffer);
                break;
            default:
                throw DataUtils.newIllegalArgumentException(
                        "Unsupported id {0}", Arrays.toString(id));
            }
        }
        return length;
    }

    /**
     * Check whether the id itself contains all the data. This operation does
     * not cause any reads in the map.
     *
     * @param id the id
     * @return if the id contains the data
     */
    public boolean isInPlace(byte[] id) {
        ByteBuffer idBuffer = ByteBuffer.wrap(id);
        while (idBuffer.hasRemaining()) {
            if (idBuffer.get() != 0) {
                return false;
            }
            int len = DataUtils.readVarInt(idBuffer);
            idBuffer.position(idBuffer.position() + len);
        }
        return true;
    }

    /**
     * Open an input stream to read data.
     *
     * @param id the id
     * @return the stream
     */
    public InputStream get(byte[] id) {
        return new Stream(this, id);
    }

    /**
     * Get the block.
     *
     * @param key the key
     * @return the block
     */
    byte[] getBlock(long key) {
        byte[] data = map.get(key);
        if (data == null) {
            throw DataUtils.newMVStoreException(
                    DataUtils.ERROR_BLOCK_NOT_FOUND,
                    "Block {0} not found",  key);
        }
        return data;
    }

    /**
     * A stream backed by a map.
     */
    static class Stream extends InputStream {

        private final StreamStore store;
        private byte[] oneByteBuffer;
        private ByteBuffer idBuffer;
        private ByteArrayInputStream buffer;
        private long skip;
        private final long length;
        private long pos;

        Stream(StreamStore store, byte[] id) {
            this.store = store;
            this.length = store.length(id);
            this.idBuffer = ByteBuffer.wrap(id);
        }

        @Override
        public int read() throws IOException {
            byte[] buffer = oneByteBuffer;
            if (buffer == null) {
                buffer = oneByteBuffer = new byte[1];
            }
            int len = read(buffer, 0, 1);
            return len == -1 ? -1 : (buffer[0] & 255);
        }

        @Override
        public long skip(long n) {
            n = Math.min(length - pos, n);
            if (n == 0) {
                return 0;
            }
            if (buffer != null) {
                long s = buffer.skip(n);
                if (s > 0) {
                    n = s;
                } else {
                    buffer = null;
                    skip += n;
                }
            } else {
                skip += n;
            }
            pos += n;
            return n;
        }

        @Override
        public void close() {
            buffer = null;
            idBuffer.position(idBuffer.limit());
            pos = length;
        }

        @Override
        public int read(byte[] b, int off, int len) throws IOException {
            if (len <= 0) {
                return 0;
            }
            while (true) {
                if (buffer == null) {
                    try {
                        buffer = nextBuffer();
                    } catch (MVStoreException e) {
                        String msg = DataUtils.formatMessage(
                                DataUtils.ERROR_BLOCK_NOT_FOUND,
                                "Block not found in id {0}",
                                Arrays.toString(idBuffer.array()));
                        throw new IOException(msg, e);
                    }
                    if (buffer == null) {
                        return -1;
                    }
                }
                int result = buffer.read(b, off, len);
                if (result > 0) {
                    pos += result;
                    return result;
                }
                buffer = null;
            }
        }

        private ByteArrayInputStream nextBuffer() {
            while (idBuffer.hasRemaining()) {
                switch (idBuffer.get()) {
                case 0: {
                    int len = DataUtils.readVarInt(idBuffer);
                    if (skip >= len) {
                        skip -= len;
                        idBuffer.position(idBuffer.position() + len);
                        continue;
                    }
                    int p = (int) (idBuffer.position() + skip);
                    int l = (int) (len - skip);
                    idBuffer.position(p + l);
                    return new ByteArrayInputStream(idBuffer.array(), p, l);
                }
                case 1: {
                    int len = DataUtils.readVarInt(idBuffer);
                    long key = DataUtils.readVarLong(idBuffer);
                    if (skip >= len) {
                        skip -= len;
                        continue;
                    }
                    byte[] data = store.getBlock(key);
                    int s = (int) skip;
                    skip = 0;
                    return new ByteArrayInputStream(data, s, data.length - s);
                }
                case 2: {
                    long len = DataUtils.readVarLong(idBuffer);
                    long key = DataUtils.readVarLong(idBuffer);
                    if (skip >= len) {
                        skip -= len;
                        continue;
                    }
                    byte[] k = store.getBlock(key);
                    ByteBuffer newBuffer = ByteBuffer.allocate(k.length
                            + idBuffer.limit() - idBuffer.position());
                    newBuffer.put(k);
                    newBuffer.put(idBuffer);
                    newBuffer.flip();
                    idBuffer = newBuffer;
                    return nextBuffer();
                }
                default:
                    throw DataUtils.newIllegalArgumentException(
                            "Unsupported id {0}",
                            Arrays.toString(idBuffer.array()));
                }
            }
            return null;
        }

    }

}
