package com.atlassian.bitbucket.scm.cache.web;

import com.google.common.base.Charsets;

import javax.annotation.Nonnull;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;

public class ServiceAdvertisementStrippingOutputStream extends OutputStream {

    private final OutputStream delegate;
    private final ByteArrayOutputStream serviceHeader = new ByteArrayOutputStream(256);

    private volatile int advertisementLength = -1;
    private volatile boolean stripped;

    public ServiceAdvertisementStrippingOutputStream(OutputStream delegate) {
        this.delegate = delegate;
    }

    @Override
    public void write(int b) throws IOException {
        if (stripped) {
            delegate.write(b);
        } else {
            serviceHeader.write(b);
            stripAdvertisement();
        }
    }

    @Override
    public void write(@Nonnull byte[] b) throws IOException {
        if (stripped) {
            delegate.write(b);
        } else {
            serviceHeader.write(b);
            stripAdvertisement();
        }
    }

    @Override
    public void write(@Nonnull byte[] b, int off, int len) throws IOException {
        if (stripped) {
            delegate.write(b, off, len);
        } else {
            serviceHeader.write(b, off, len);
            stripAdvertisement();
        }
    }

    @Override
    public void flush() throws IOException {
        delegate.flush();
    }

    @Override
    public void close() throws IOException {
        delegate.close();
    }

    private void stripAdvertisement() throws IOException {
        // the service advertisement consists of 2 packets:
        // 001e# service=git-upload-pack
        // 0000
        if (advertisementLength == -1 && serviceHeader.size() >= 4) {
            try {
                advertisementLength =
                        Integer.parseInt(new String(serviceHeader.toByteArray(), 0, 4, Charsets.UTF_8), 16) + 4;
            } catch (NumberFormatException e) {
                // not a packet after all - nothing to do but get out of the way
                advertisementLength = 0;
            }
        }
        if (advertisementLength >= 0) {
            byte[] packetBytes = serviceHeader.toByteArray();
            stripped = packetBytes.length >= advertisementLength;
            if (packetBytes.length > advertisementLength) {
                delegate.write(packetBytes, advertisementLength, packetBytes.length - advertisementLength);
                delegate.flush();
            }
        }
    }
}
