package com.atlassian.bitbucket.mesh.transaction;

import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamChunk;
import com.atlassian.bitbucket.mesh.util.SatiableConsumer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.time.Duration;
import java.time.Instant;
import java.util.*;

import static java.util.Objects.requireNonNull;

public class TeePackWriter implements SatiableConsumer<RpcStreamChunk> {

    private static final Logger log = LoggerFactory.getLogger(TeePackWriter.class);

    private final Map<ReplicaTransaction, SatiableConsumer<RpcStreamChunk>> delegateByReplicaTx;
    private final Instant startTimestamp;
    private volatile boolean closed;
    private volatile long size;

    public TeePackWriter(@Nonnull Map<ReplicaTransaction, SatiableConsumer<RpcStreamChunk>> delegateByReplicaTx) {
        this.delegateByReplicaTx = new HashMap<>(requireNonNull(delegateByReplicaTx, "delegates"));

        startTimestamp = Instant.now();
    }

    public boolean isClosed() {
        return closed;
    }

    @Override
    public boolean accept(@Nonnull RpcStreamChunk chunk) {
        if (delegateByReplicaTx.isEmpty()) {
            return true;
        }
        size += chunk.getData().size();
        Iterator<Map.Entry<ReplicaTransaction, SatiableConsumer<RpcStreamChunk>>> it =
                delegateByReplicaTx.entrySet().iterator();
        while (it.hasNext()) {
            Map.Entry<ReplicaTransaction, SatiableConsumer<RpcStreamChunk>> entry = it.next();
            SatiableConsumer<RpcStreamChunk> delegate = entry.getValue();
            try {
                if (delegate.accept(chunk)) {
                    log.debug("[{}] Removing pack writer because it signalled it does not want more data",
                            entry.getKey());
                    it.remove();
                }
            } catch (RuntimeException e) {
                log.warn("[{}] Dropping pack writer after write failure.", entry.getKey(), e);
                it.remove();
            }
        }
        closed |= chunk.getClosed();
        if (closed) {
            log.debug("Forwarded {} pack bytes to {} followers in {}ms", size, delegateByReplicaTx.size(),
                    Duration.between(startTimestamp, Instant.now()).toMillis());
        }
        return delegateByReplicaTx.isEmpty();
    }
}
