package ch.threema.app.webrtc;

import org.saltyrtc.tasks.webrtc.exceptions.IllegalStateError;
import org.slf4j.Logger;
import org.webrtc.DataChannel;
import java.util.concurrent.CompletableFuture;
import androidx.annotation.AnyThread;
import androidx.annotation.NonNull;
import ch.threema.logging.ThreemaLogger;

import static ch.threema.base.utils.LoggingKt.getThreemaLogger;


/**
 * A flow-controlled (sender side) data channel.
 * <p>
 * When using this, make sure to properly call `bufferedAmountChange` when the corresponding
 * event on the data channel is received.
 */
@AnyThread
public class FlowControlledDataChannel {
    @NonNull
    final private Logger logger = getThreemaLogger("FlowControlledDataChannel");
    @NonNull
    public final DataChannel dc;
    private final long lowWaterMark;
    private final long highWaterMark;
    @NonNull
    private CompletableFuture<Void> readyFuture = CompletableFuture.completedFuture(null);

    /**
     * Create a flow-controlled (sender side) data channel.
     *
     * @param dc The data channel to be flow-controlled
     */
    public FlowControlledDataChannel(@NonNull final String logPrefix, @NonNull final DataChannel dc) {
        this(logPrefix, dc, 256 * 1024, 1024 * 1024);
    }

    /**
     * Create a flow-controlled (sender side) data channel.
     *
     * @param dc            The data channel to be flow-controlled
     * @param lowWaterMark  The low water mark unpauses the data channel once
     *                      the buffered amount of bytes becomes less or equal to it.
     * @param highWaterMark The high water mark pauses the data channel once
     *                      the buffered amount of bytes becomes greater or equal to it.
     */
    public FlowControlledDataChannel(
        @NonNull final String logPrefix,
        @NonNull final DataChannel dc,
        final long lowWaterMark,
        final long highWaterMark
    ) {
        // Set logger prefix
        if (logger instanceof ThreemaLogger) {
            ((ThreemaLogger) logger).setPrefix(logPrefix + "." + dc.label() + "/" + dc.id());
        }

        this.dc = dc;
        this.lowWaterMark = lowWaterMark;
        this.highWaterMark = highWaterMark;
    }

    /**
     * Return the low water mark.
     */
    public long getLowWaterMark() {
        return this.lowWaterMark;
    }

    /**
     * Return the high water mark.
     */
    public long getHighWaterMark() {
        return this.highWaterMark;
    }

    /**
     * A future whether the data channel is ready to be written on.
     */
    @NonNull
    public synchronized CompletableFuture<Void> ready() {
        return this.readyFuture;
    }

    /**
     * Write a message to the data channel's internal buffer for delivery to
     * the remote side.
     * <p>
     * Important: Before calling this, the `ready` Promise must be awaited.
     *
     * @param message The message to be sent.
     * @throws IllegalStateError in case the data channel is currently paused.
     */
    public synchronized void write(@NonNull final DataChannel.Buffer message) {
        // Note: Locked since the "onBufferedAmountChange" event must run in parallel to the send
        //       calls.

        // Throw if paused
        if (!this.ready().isDone()) {
            throw new IllegalStateError("Unable to write, data channel is paused!");
        }

        // Try sending
        // Note: Technically we should be able to catch an Exception in case the
        //       underlying buffer is full. However, webrtc.org is utterly
        //       outdated and just closes when its buffer would overflow. Thus,
        //       we use a well-tested high water mark instead and try to never
        //       fill the buffer completely.
        if (!this.dc.send(message)) {
            // This can happen when the data channel is closing.
            throw new IllegalStateError("Unable to send in state " + this.dc.state());
        }

        // Pause once high water mark has been reached
        final long bufferedAmount = this.dc.bufferedAmount();
        if (bufferedAmount >= this.highWaterMark) {
            this.readyFuture = new CompletableFuture<>();
            if (logger.isDebugEnabled()) {
                logger.debug("{} paused (buffered={})", this.dc.label(), bufferedAmount);
            }
        }
    }

    /**
     * Must be called when the data channel's buffered amount changed.
     * <p>
     * Important: You MUST ensure that you're not calling this from the send thread of the data
     * channel! When in doubt, post it to some other thread!
     */
    public synchronized void bufferedAmountChange() {
        final long bufferedAmount;
        try {
            bufferedAmount = this.dc.bufferedAmount();
        } catch (IllegalStateException e) {
            logger.warn("IllegalStateException when calling `dc.bufferedAmount`, data channel already disposed?");
            return;
        }

        // Unpause once low water mark has been reached
        if (bufferedAmount <= this.lowWaterMark && !this.readyFuture.isDone()) {
            if (logger.isDebugEnabled()) {
                logger.debug("{} resumed (buffered={})", this.dc.label(), bufferedAmount);
            }
            this.readyFuture.complete(null);
        }
    }
}
