package org.apache.flink.streaming.connectors.http;

import org.apache.flink.connector.base.sink.AsyncSinkBase;
import org.apache.flink.connector.base.sink.writer.BufferedRequestState;
import org.apache.flink.connector.base.sink.writer.ElementConverter;
import org.apache.flink.connector.base.sink.writer.RequestEntryWrapper;
import org.apache.flink.core.io.SimpleVersionedSerializer;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

/** @param <InputT> */
public class HttpSink<InputT> extends AsyncSinkBase<InputT, byte []> {

    private final String endpointURL;
    private final String method;
    private static final ElementConverter<Object, byte []> ELEMENT_CONVERTER =
            ((element, context) -> (byte[]) element);
    private static final int MAX_BATCH_SIZE = 10000;
    private static final int MAX_IN_FLIGHT_REQUESTS = 100000; // must be > max_batch_size
    private static final int MAX_BUFFERED_REQUESTS = 100000; // must be > max_batch_size
    private static final int MAX_BATCH_SIZE_IN_BYTES = 10000000;
    private static final int MAX_TIME_IN_BUFFER_MS = 1000;
    private static final int MAX_RECORD_SIZE_IN_BYTES = 10000000;

    public HttpSink(String endpointURL, String method) {
        super(
                (ElementConverter<InputT, byte []>) ELEMENT_CONVERTER,
                MAX_BATCH_SIZE,
                MAX_IN_FLIGHT_REQUESTS,
                MAX_BUFFERED_REQUESTS,
                MAX_BATCH_SIZE_IN_BYTES,
                MAX_TIME_IN_BUFFER_MS,
                MAX_RECORD_SIZE_IN_BYTES);
        this.endpointURL = endpointURL;
        this.method = method;
    }

    @Override
    public StatefulSinkWriter<InputT, BufferedRequestState<byte []>> createWriter(
            InitContext initContext) throws IOException {
        System.out.println("creating writer...");
        return new HttpSinkWriter(
                ELEMENT_CONVERTER,
                initContext,
                MAX_BATCH_SIZE,
                MAX_IN_FLIGHT_REQUESTS,
                MAX_BUFFERED_REQUESTS,
                MAX_BATCH_SIZE_IN_BYTES,
                MAX_TIME_IN_BUFFER_MS,
                MAX_RECORD_SIZE_IN_BYTES,
                endpointURL,
                method);
    }

    @Override
    public StatefulSinkWriter<InputT, BufferedRequestState<byte []>> restoreWriter(
            InitContext initContext, Collection<BufferedRequestState<byte []>> collection)
            throws IOException {
        return new HttpSinkWriter(
                ELEMENT_CONVERTER,
                initContext,
                MAX_BATCH_SIZE,
                MAX_IN_FLIGHT_REQUESTS,
                MAX_BUFFERED_REQUESTS,
                MAX_BATCH_SIZE_IN_BYTES,
                MAX_TIME_IN_BUFFER_MS,
                MAX_RECORD_SIZE_IN_BYTES,
                endpointURL,
                method);
    }

    @Override
    public SimpleVersionedSerializer<BufferedRequestState<byte []>> getWriterStateSerializer() {

        return new SimpleVersionedSerializer<BufferedRequestState<byte []>>() {
            @Override
            public int getVersion() {
                return 0;
            }

            /**
             * Serializes state in form of
             * [DATA_IDENTIFIER,NUM_OF_ELEMENTS,SIZE1,REQUEST1,SIZE2,REQUEST2....].
             */
            @Override
            public byte[] serialize(BufferedRequestState<byte []> obj) throws IOException {
                Collection<RequestEntryWrapper<byte []>> bufferState =
                        obj.getBufferedRequestEntries();

                try (final ByteArrayOutputStream baos = new ByteArrayOutputStream();
                        final DataOutputStream out = new DataOutputStream(baos)) {

                    out.writeLong(-1);
                    out.writeInt(bufferState.size());

                    for (RequestEntryWrapper<byte []> wrapper : bufferState) {
                        out.writeLong(wrapper.getSize());
                        serializeRequestToStream(wrapper.getRequestEntry(), out);
                    }

                    return baos.toByteArray();
                }
            }

            @Override
            public BufferedRequestState<byte []> deserialize(int version, byte[] serialized)
                    throws IOException {
                try (final ByteArrayInputStream bais = new ByteArrayInputStream(serialized);
                        final DataInputStream in = new DataInputStream(bais)) {

                    validateIdentifier(in);

                    int size = in.readInt();
                    List<RequestEntryWrapper<String>> serializedState = new ArrayList<>();

                    for (int i = 0; i < size; i++) {
                        long requestSize = in.readLong();
                        String request = deserializeRequestFromStream(requestSize, in);
                        serializedState.add(new RequestEntryWrapper<>(request, requestSize));
                    }

                    return new BufferedRequestState(serializedState);
                }
            }

            protected void serializeRequestToStream(byte [] request, DataOutputStream out)
                    throws IOException {
                out.flush();
            }

            protected String deserializeRequestFromStream(long requestSize, DataInputStream in)
                    throws IOException {
                return in.toString();
            }

            private void validateIdentifier(DataInputStream in) throws IOException {
                if (in.readLong() != -1) {
                    throw new IllegalStateException("Corrupted data to deserialize");
                }
            }
        };
    }
}