/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.flink.streaming.connectors.influxdb.sink2.writer;

import com.influxdb.client.InfluxDBClient;
import com.influxdb.client.WriteApi;
import com.influxdb.client.domain.WritePrecision;
import com.influxdb.client.write.Point;
import org.apache.flink.api.common.operators.ProcessingTimeService;
import org.apache.flink.api.connector.sink2.SinkWriter;
import org.apache.flink.configuration.Configuration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;

import static org.apache.flink.streaming.connectors.influxdb.sink2.InfluxDBSinkOptions.*;

public class InfluxDBWriter<IN> implements SinkWriter<IN> {

    private static final Logger LOG = LoggerFactory.getLogger(InfluxDBWriter.class);

    private final int bufferSize;
    private final boolean writeCheckpoint;
    private long lastTimestamp = 0;
    private final List<Point> elements;
    private ProcessingTimeService processingTimerService;
    private final InfluxDBSchemaSerializer<IN> schemaSerializer;
    private final InfluxDBClient influxDBClient;

    public InfluxDBWriter(
            final InfluxDBSchemaSerializer<IN> schemaSerializer,
            final Configuration configuration) {
        this.schemaSerializer = schemaSerializer;
        this.bufferSize = configuration.getInteger(WRITE_BUFFER_SIZE);
        this.elements = new ArrayList<>(this.bufferSize);
        this.writeCheckpoint = configuration.getBoolean(WRITE_DATA_POINT_CHECKPOINT);
        this.influxDBClient = getInfluxDBClient(configuration);
    }

    public void setProcessingTimerService(final ProcessingTimeService processingTimerService) {
        this.processingTimerService = processingTimerService;
    }

    /**
     * This method calls the InfluxDB write API whenever the element list reaches the {@link
     * #bufferSize}. It keeps track of the latest timestamp of each element. It compares the latest
     * timestamp with the context.timestamp() and takes the bigger (latest) timestamp.
     *
     * @param in incoming data
     * @param context current Flink context
     * @see org.apache.flink.api.connector.sink2.SinkWriter.Context
     */
    @Override
    public void write(IN in, Context context) throws IOException, InterruptedException {
        LOG.trace("Adding elements to buffer. Buffer size: {}", this.elements.size());
        this.elements.add(this.schemaSerializer.serialize(in, context));

        if (this.elements.size() == this.bufferSize) {
            LOG.debug("Buffer size reached preparing to write the elements.");
            this.writeCurrentElements();
        }
        if (context.timestamp() != null) {
            this.lastTimestamp = Math.max(this.lastTimestamp, context.timestamp());
        }

    }

    @Override
    public void flush(boolean flush) throws IOException, InterruptedException {
        if (this.lastTimestamp == 0) return;

        commit(Collections.singletonList(this.lastTimestamp));
    }

    public void commit(final List<Long> committables) {
        if (this.writeCheckpoint) {
            LOG.debug("A checkpoint is set.");
            Optional<Long> lastTimestamp = Optional.empty();
            if (committables.size() >= 1) {
                lastTimestamp = Optional.ofNullable(committables.get(committables.size() - 1));
            }
            lastTimestamp.ifPresent(this::writeCheckpointDataPoint);
        }
    }

    private void writeCheckpointDataPoint(final Long timestamp) {
        final Point point = new Point("checkpoint")
                .addField("checkpoint", "flink")
                .time(timestamp, WritePrecision.MS);

        writeElementsOf(Collections.singletonList(point));
        LOG.debug("Checkpoint data point write at {}", point.toLineProtocol());
    }

    @Override
    public void close() throws Exception {
        LOG.debug("Preparing to write the elements in InfluxDB.");
        this.writeCurrentElements();

        LOG.debug("Closing the writer.");
        this.influxDBClient.close();
    }

    private void writeCurrentElements() {
        writeElementsOf(this.elements);
        this.elements.clear();
    }

    private void writeElementsOf(List<Point> toWrite) {
        if (toWrite.isEmpty()) return;

        try (final WriteApi writeApi = this.influxDBClient.getWriteApi()) {
            writeApi.writePoints(toWrite);
            LOG.debug("Wrote {} data points", toWrite.size());
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
