package be.bagofwords.db.remote;

import be.bagofwords.db.methods.DataStream;
import be.bagofwords.db.methods.DataStreamUtils;
import be.bagofwords.db.methods.ObjectSerializer;
import be.bagofwords.iterator.CloseableIterator;
import be.bagofwords.util.KeyValue;
import org.xerial.snappy.Snappy;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import static be.bagofwords.db.remote.Protocol.LONG_END;
import static be.bagofwords.db.remote.Protocol.LONG_ERROR;

/**
 * Created by koen on 21/05/17.
 */
public class KeyValueSocketIterator<T> extends CloseableIterator<KeyValue<T>> {

    private RemoteDataInterface<T> remoteDataInterface;
    private final Connection connection;
    private Iterator<KeyValue<T>> nextValues;
    private boolean readAllValuesFromConnection;
    private final ObjectSerializer<T> objectSerializer;

    public KeyValueSocketIterator(RemoteDataInterface<T> remoteDataInterface, Connection connection) {
        this.remoteDataInterface = remoteDataInterface;
        this.connection = connection;
        this.readAllValuesFromConnection = false;
        this.objectSerializer = remoteDataInterface.getObjectSerializer();
        //Constructor
        findNextValues();
    }

    private synchronized void findNextValues() {
        if (!wasClosed()) {
            try {
                long numOfValues = connection.readLong();
                if (numOfValues == LONG_END) {
                    nextValues = null;
                    readAllValuesFromConnection = true;
                } else if (numOfValues != LONG_ERROR) {
                    byte[] keys = connection.readByteArray();
                    byte[] compressedValues = connection.readByteArray();
                    byte[] uncompressedValues = Snappy.uncompress(compressedValues);
                    DataStream keyIS = new DataStream(keys);
                    DataStream valueIS = new DataStream(uncompressedValues);
                    List<KeyValue<T>> nextValuesList = new ArrayList<>();
                    while (nextValuesList.size() < numOfValues) {
                        long key = keyIS.readLong();
                        int objectSize = DataStreamUtils.getObjectSize(valueIS, objectSerializer);
                        T value = objectSerializer.readValue(valueIS, objectSize);
                        nextValuesList.add(new KeyValue<>(key, value));
                    }
                    if (nextValuesList.isEmpty()) {
                        throw new RuntimeException("Received zero values! numOfValues=" + numOfValues);
                    }
                    nextValues = nextValuesList.iterator();
                } else {
                    throw new RuntimeException("Unexpected response " + connection.readString());

                }
            } catch (Exception e) {
                remoteDataInterface.dropConnection(connection);
                throw new RuntimeException(e);
            }
        } else {
            nextValues = null;
        }
    }

    @Override
    public void closeInt() {
        if (readAllValuesFromConnection) {
            remoteDataInterface.releaseConnection(connection);
        } else {
            //server will still be sending data through this connection, so it can not be reused.
            remoteDataInterface.dropConnection(connection);
        }
    }

    @Override
    public boolean hasNext() {
        return nextValues != null;
    }

    @Override
    public KeyValue<T> next() {
        KeyValue<T> result = nextValues.next();
        if (!nextValues.hasNext()) {
            findNextValues();
        }
        return result;
    }

    @Override
    public void remove() {
        throw new RuntimeException("Not implemented");
    }

}