/*
 * Copyright (c) 2018 Spotify AB.
 *
 * 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 com.spotify.heroic.consumer.pubsub;

import static io.opencensus.trace.AttributeValue.stringAttributeValue;

import com.google.cloud.pubsub.v1.AckReplyConsumer;
import com.google.cloud.pubsub.v1.MessageReceiver;
import com.google.protobuf.ByteString;
import com.google.pubsub.v1.PubsubMessage;
import com.spotify.heroic.consumer.ConsumerSchema;
import com.spotify.heroic.consumer.ConsumerSchemaValidationException;
import com.spotify.heroic.statistics.ConsumerReporter;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.LongAdder;
import io.opencensus.common.Scope;
import io.opencensus.trace.Span;
import io.opencensus.trace.Status;
import io.opencensus.trace.Tracer;
import io.opencensus.trace.Tracing;
import lombok.extern.slf4j.Slf4j;

@Slf4j
class Receiver implements MessageReceiver {
    private final ConsumerSchema.Consumer consumer;
    private final ConsumerReporter reporter;
    private final AtomicLong errors;
    private final LongAdder consumed;
    private final Tracer tracer = Tracing.getTracer();

    Receiver(
        final ConsumerSchema.Consumer consumer,
        final ConsumerReporter reporter,
        final AtomicLong errors,
        final LongAdder consumed
    ) {
        this.consumer = consumer;
        this.reporter = reporter;
        this.errors = errors;
        this.consumed = consumed;
    }

    @Override
    public void receiveMessage(PubsubMessage message, AckReplyConsumer replyConsumer) {
        Scope ss = tracer.spanBuilder("PubSub.receiveMessage").startScopedSpan();

        // handle incoming message, then ack/nack the received message
        final ByteString data = message.getData();
        final String messageId = message.getMessageId();
        log.debug("Received ID:{} with content: {}", messageId, data.toStringUtf8());
        final byte[] bytes = data.toByteArray();

        Span span = tracer.getCurrentSpan();
        span.putAttribute("id", stringAttributeValue(messageId));

        // process the data
        try {
            consumer.consume(bytes);
            reporter.reportMessageSize(bytes.length);
            consumed.increment();
            replyConsumer.ack();
        } catch (ConsumerSchemaValidationException e) {
            reporter.reportConsumerSchemaError();
            log.error("ID:{} - {}", messageId, e.getMessage(), e);

            // The message will never be processable, ack it to make it go away
            replyConsumer.ack();
        } catch (Exception e) {
            errors.incrementAndGet();
            log.error("ID:{} - Failed to consume", messageId, e);
            span.setStatus(Status.INTERNAL.withDescription(e.toString()));
            reporter.reportMessageError();
            replyConsumer.nack();
        } finally {
            span.end();
            ss.close();
        }
  }

}

@@ -1313,24 +1313,77 @@
 erReporter;%0A
+import com.spotify.heroic.statistics.FutureReporter;%0A
 import java.
@@ -2875,24 +2875,113 @@
 ssageId));%0A%0A
+        final FutureReporter.Context consumptionContext = reporter.reportConsumption();%0A%0A
         // p
@@ -3045,18 +3045,67 @@
 e(bytes)
-;%0A
+.onDone(consumptionContext).onFinished(() -%3E %7B%0A    
         
@@ -3166,26 +3166,29 @@
     
-consumed.increment
+    replyConsumer.ack
 ();%0A
@@ -3199,34 +3199,45 @@
         
-replyConsumer.ack(
+    span.end();%0A            %7D
 );%0A     
@@ -3510,32 +3510,56 @@
 Consumer.ack();%0A
+            span.end();%0A
         %7D catch 
@@ -3826,24 +3826,48 @@
 mer.nack();%0A
+            span.end();%0A
         %7D fi
@@ -3882,32 +3882,42 @@
             
-span.end
+consumed.increment
 ();%0A        

Move ack and span end to wait for the future to finish.