package stroom.query.common.v2;

import stroom.lmdb.LmdbEnv.BatchingWriteTxn;
import stroom.query.api.v2.QueryKey;
import stroom.query.common.v2.SearchProgressLog.SearchPhase;
import stroom.util.concurrent.CompleteException;
import stroom.util.logging.LambdaLogger;
import stroom.util.logging.LambdaLoggerFactory;
import stroom.util.logging.Metrics;

import com.esotericsoftware.kryo.KryoException;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.lmdbjava.CursorIterable;
import org.lmdbjava.CursorIterable.KeyVal;
import org.lmdbjava.Dbi;
import org.lmdbjava.Txn;

import java.io.ByteArrayInputStream;
import java.nio.ByteBuffer;

public class LmdbPayloadCreator {

    private static final LambdaLogger LOGGER = LambdaLoggerFactory.getLogger(LmdbPayloadCreator.class);

    private final int maxPayloadSize;
    private final QueryKey queryKey;
    private final LmdbDataStore lmdbDataStore;
    private final LmdbPayloadQueue currentPayload = new LmdbPayloadQueue(1);
    private final LmdbRowKeyFactory lmdbRowKeyFactory;
    private final int minPayloadSize;

    LmdbPayloadCreator(final QueryKey queryKey,
                       final LmdbDataStore lmdbDataStore,
                       final AbstractResultStoreConfig resultStoreConfig,
                       final LmdbRowKeyFactory lmdbRowKeyFactory) {
        this.queryKey = queryKey;
        this.lmdbDataStore = lmdbDataStore;
        maxPayloadSize = (int) resultStoreConfig.getMaxPayloadSize().getBytes();
        this.lmdbRowKeyFactory = lmdbRowKeyFactory;
        this.minPayloadSize = (int) resultStoreConfig.getMinPayloadSize().getBytes();
    }

    /**
     * Read items from the supplied input and transfer them to the data store.
     *
     * @param input The input to read.
     */
    void readPayload(final Input input) {
        SearchProgressLog.increment(queryKey, SearchPhase.LMDB_DATA_STORE_READ_PAYLOAD);
        Metrics.measure("readPayload", () -> {
            // Determine how many bytes the payload contains.
            final int length = input.readInt();
            if (length > 0) {
                // If there are some bytes to read then read them.
                final byte[] bytes = input.readBytes(length);
                try (final Input in = new Input(new ByteArrayInputStream(bytes))) {
                    while (!in.end()) {
                        final int rowKeyLength = in.readInt();
                        final byte[] key = in.readBytes(rowKeyLength);
                        ByteBuffer keyBuffer = ByteBuffer.allocateDirect(key.length);
                        keyBuffer.put(key, 0, key.length);
                        keyBuffer.flip();

                        final int valueLength = in.readInt();
                        final byte[] value = in.readBytes(valueLength);
                        ByteBuffer valueBuffer = ByteBuffer.allocateDirect(value.length);
                        valueBuffer.put(value, 0, value.length);
                        valueBuffer.flip();

                        // Create a new unique key if this isn't a group key.
                        LmdbKV lmdbKV = new LmdbKV(null, keyBuffer, valueBuffer);
                        lmdbKV = lmdbRowKeyFactory.makeUnique(lmdbKV);

                        lmdbDataStore.put(lmdbKV);
                    }
                }
            }
        });
    }

    /**
     * Write data from the data store to an output removing them from the datastore as we go as they will be transferred
     * to another store.
     *
     * @param output The output to write to.
     */
    void writePayload(final Output output) {
        try {
            SearchProgressLog.increment(queryKey, SearchPhase.LMDB_DATA_STORE_WRITE_PAYLOAD);
            Metrics.measure("writePayload", () -> {
                try {
                    // Get the current payload.
                    final LmdbPayload payload = currentPayload.poll();
                    if (payload == null) {
                        // There will be 0 bytes to read so just write 0.
                        output.writeInt(0);

                    } else {
                        final byte[] data = payload.getData();
                        output.writeInt(data.length);
                        output.writeBytes(data);
                    }
                } catch (final KryoException e) {
                    // Expected as sometimes the output stream is closed by the receiving node.
                    LOGGER.debug(e::getMessage, e);
                } catch (final InterruptedException e) {
                    // There will be 0 bytes to read so just write 0.
                    output.writeInt(0);
                    LOGGER.trace(e::getMessage, e);
                    // Keep interrupting this thread.
                    Thread.currentThread().interrupt();
                } catch (final CompleteException e) {
                    // There will be 0 bytes to read so just write 0.
                    output.writeInt(0);
                    LOGGER.debug(() -> "Complete");
                    LOGGER.trace(e::getMessage, e);
                }
            });
        } catch (final RuntimeException e) {
            LOGGER.debug(e::getMessage, e);
        }
    }

    boolean isEmpty() {
        return currentPayload.isEmpty();
    }

    boolean addPayload(final BatchingWriteTxn writeTxn,
                       final Dbi<ByteBuffer> dbi,
                       final boolean complete) {
        final LmdbPayload payload = createPayload(writeTxn, dbi, complete);
        doPut(payload);
        return payload.isFinalPayload();
    }

    void finalPayload() {
        currentPayload.complete();
    }

    private void doPut(final LmdbPayload payload) {
        try {
            currentPayload.put(payload);
        } catch (final InterruptedException e) {
            LOGGER.trace(e.getMessage(), e);
            // Keep interrupting this thread.
            Thread.currentThread().interrupt();
        }
    }

    private LmdbPayload createPayload(final BatchingWriteTxn writeTxn,
                                      final Dbi<ByteBuffer> dbi,
                                      final boolean complete) {
        SearchProgressLog.increment(queryKey, SearchPhase.LMDB_DATA_STORE_CREATE_PAYLOAD);
        final Txn<ByteBuffer> txn = writeTxn.getTxn();

        return Metrics.measure("createPayload", () -> {
            if (maxPayloadSize > 0) {
                final PayloadOutput payloadOutput = new PayloadOutput(minPayloadSize);
                boolean finalPayload = complete;
                long size = 0;
                long count = 0;
                try (final CursorIterable<ByteBuffer> cursorIterable = dbi.iterate(txn)) {
                    for (final KeyVal<ByteBuffer> kv : cursorIterable) {
                        final ByteBuffer keyBuffer = kv.key();
                        final ByteBuffer valBuffer = kv.val();

                        // Make sure we don't add a state key to the payload.
                        if (LmdbRowKeyFactoryFactory.isNotStateKey(keyBuffer)) {
                            // Add to the size of the current payload.
                            size += 4;
                            size += keyBuffer.remaining();
                            size += 4;
                            size += valBuffer.remaining();
                            count++;

                            if (size < maxPayloadSize || count == 1) {
                                payloadOutput.writeInt(keyBuffer.remaining());
                                payloadOutput.writeByteBuffer(keyBuffer);
                                payloadOutput.writeInt(valBuffer.remaining());
                                payloadOutput.writeByteBuffer(valBuffer);

                                dbi.delete(txn, keyBuffer.flip());

                            } else {
                                // We have reached the maximum payload size so stop adding.
                                // We also can't be complete as more payloads will be needed.
                                finalPayload = false;
                                break;
                            }
                        }
                    }
                }

                writeTxn.commit();
                payloadOutput.close();
                return new LmdbPayload(finalPayload, payloadOutput.toBytes());

            } else {
                final PayloadOutput payloadOutput = new PayloadOutput(minPayloadSize);
                try (final CursorIterable<ByteBuffer> cursorIterable = dbi.iterate(txn)) {
                    for (final KeyVal<ByteBuffer> kv : cursorIterable) {
                        final ByteBuffer keyBuffer = kv.key();
                        final ByteBuffer valBuffer = kv.val();

                        payloadOutput.writeInt(keyBuffer.remaining());
                        payloadOutput.writeByteBuffer(keyBuffer);
                        payloadOutput.writeInt(valBuffer.remaining());
                        payloadOutput.writeByteBuffer(valBuffer);
                    }
                }

                dbi.drop(txn);
                writeTxn.commit();
                payloadOutput.close();
                return new LmdbPayload(complete, payloadOutput.toBytes());
            }
        });
    }
}
