// Copyright 2017 The Bazel Authors. All rights reserved.
//
// Licensed 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 build.buildfarm.worker;

import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import static java.util.concurrent.TimeUnit.SECONDS;

import build.buildfarm.common.Write;
import com.google.common.base.Preconditions;
import com.google.protobuf.ByteString;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.OutputStreamWriter;

public class ByteStringWriteReader implements Runnable {
  private final InputStream input;
  private final Write write;
  private ByteString.Output data = ByteString.newOutput();
  private boolean completed;
  private IOException exception = null;

  private final int dataLimit;

  public ByteStringWriteReader(InputStream input, Write write, int dataLimit) {
    this.input = input;
    this.write = write;
    this.dataLimit = dataLimit;
    Preconditions.checkState(dataLimit > 0);
    completed = false;
  }

  public void run() {
    boolean closed = false;
    byte[] buffer = new byte[1024 * 16];
    int len;
    write.addListener(this::complete, directExecutor());
    // may want to buffer this if not ready
    try (OutputStream writeOut = write.getOutput(1, SECONDS, () -> {})) {
      while (!isComplete() && (len = input.read(buffer)) != -1) {
        if (len != 0) {
          int dataLen = len;
          int size = data.size();
          if (size < dataLimit) {
            if (size + dataLen > dataLimit) {
              dataLen = dataLimit - size;
            }
            data.write(buffer, 0, len);
            if (size + dataLen > dataLimit) {
              new OutputStreamWriter(data).write("\nOutput truncated after exceeding limit.\n");
            }
          }
          writeOut.write(buffer, 0, len);
        }
      }
    } catch (IOException e) {
      // ignore asynchronous stream closure, this fulfills our objective
      // openjdk throws with the following copy when the process stderr/out
      // streams are asynchronously closed on process termination.
      if (!e.getMessage().equals("Stream closed")) {
        exception = e;
      }
    } finally {
      try {
        input.close();
      } catch (IOException e) {
        if (exception == null) {
          exception = e;
        } else {
          exception.addSuppressed(e);
        }
      }
    }
  }

  private synchronized void waitForComplete() throws InterruptedException {
    while (!completed) {
      wait();
    }
  }

  private synchronized void complete() {
    completed = true;
    notify();
  }

  public synchronized boolean isComplete() {
    return completed;
  }

  public ByteString getData() throws IOException, InterruptedException {
    if (exception != null) {
      throw exception;
    }
    waitForComplete();
    return data.toByteString();
  }
}