/*
 * 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 org.apache.beam.runners.dataflow.worker.logging;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.PrintStream;
import java.io.UnsupportedEncodingException;
import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.charset.Charset;
import java.nio.charset.CharsetDecoder;
import java.nio.charset.CoderMalfunctionError;
import java.nio.charset.CodingErrorAction;
import java.util.Formatter;
import java.util.Locale;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.logging.Handler;
import java.util.logging.Level;
import java.util.logging.LogRecord;
import java.util.logging.Logger;
import javax.annotation.concurrent.GuardedBy;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;

/**
 * A {@link PrintStream} factory that creates {@link PrintStream}s which output to the specified JUL
 * {@link Handler} at the specified {@link Level}.
 */
@SuppressWarnings({
  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
class JulHandlerPrintStreamAdapterFactory {
  private static final AtomicBoolean OUTPUT_WARNING = new AtomicBoolean(false);

  @VisibleForTesting
  static final String LOGGING_DISCLAIMER =
      String.format(
          "Please use a logger instead of System.out or System.err.%n"
              + "Please switch to using org.slf4j.Logger.%n"
              + "See: https://cloud.google.com/dataflow/pipelines/logging");

  private static class JulHandlerPrintStream extends PrintStream {
    // This limits the number of bytes which we buffer in case we don't have a flush.
    private static final int BUFFER_LIMIT = 1 << 10; // 1024 chars

    /** Hold reference of named logger to check configured {@link Level}. */
    private final Logger logger;

    private final Handler handler;
    private final String loggerName;
    private final Level messageLevel;

    @GuardedBy("this")
    private final StringBuilder buffer;

    @GuardedBy("this")
    private final CharsetDecoder decoder;

    @GuardedBy("this")
    private final CharBuffer decoded;

    @GuardedBy("this")
    private ByteArrayOutputStream carryOverBytes;

    private JulHandlerPrintStream(
        Handler handler, String loggerName, Level logLevel, Charset charset)
        throws UnsupportedEncodingException {
      super(
          new OutputStream() {
            @Override
            public void write(int i) throws IOException {
              throw new RuntimeException("All methods should be overwritten so this is unused");
            }
          },
          false,
          charset.name());
      this.handler = handler;
      this.loggerName = loggerName;
      this.messageLevel = logLevel;
      this.logger = Logger.getLogger(loggerName);
      this.buffer = new StringBuilder();
      this.decoder =
          charset
              .newDecoder()
              .onMalformedInput(CodingErrorAction.REPLACE)
              .onUnmappableCharacter(CodingErrorAction.REPLACE);
      this.carryOverBytes = new ByteArrayOutputStream();
      this.decoded = CharBuffer.allocate(BUFFER_LIMIT);
    }

    @Override
    public void flush() {
      publishIfNonEmpty(flushBufferToString());
    }

    private synchronized String flushBufferToString() {
      if (buffer.length() > 0 && buffer.charAt(buffer.length() - 1) == '\n') {
        buffer.setLength(buffer.length() - 1);
      }
      if (buffer.length() == 0) {
        return null;
      }
      String result = buffer.toString();
      buffer.setLength(0);
      return result;
    }

    @Override
    public void close() {
      flush();
    }

    @Override
    public boolean checkError() {
      return false;
    }

    @Override
    public synchronized void write(int i) {
      buffer.append(i);
    }

    @Override
    public void write(byte[] a, int offset, int length) {
      if (length == 0) {
        return;
      }

      ByteBuffer incoming = ByteBuffer.wrap(a, offset, length);
      assert incoming.hasArray();

      String msg = null;
      // Consume the added bytes, flushing on decoded newlines or if we hit
      // the buffer limit.
      synchronized (this) {
        int startLength = buffer.length();

        try {
          // Process any remaining bytes from last time by adding a byte at a time.
          while (carryOverBytes.size() > 0 && incoming.hasRemaining()) {
            carryOverBytes.write(incoming.get());
            ByteBuffer wrapped =
                ByteBuffer.wrap(carryOverBytes.toByteArray(), 0, carryOverBytes.size());
            decoder.decode(wrapped, decoded, false);
            if (!wrapped.hasRemaining()) {
              carryOverBytes.reset();
            }
          }

          // Append chunks while we are hitting the decoded buffer limit
          while (decoder.decode(incoming, decoded, false).isOverflow()) {
            decoded.flip();
            buffer.append(decoded);
            decoded.clear();
          }

          // Append the partial chunk
          decoded.flip();
          buffer.append(decoded);
          decoded.clear();

          // Check to see if we should output this message
          if (buffer.length() > BUFFER_LIMIT || buffer.indexOf("\n", startLength) >= 0) {
            msg = flushBufferToString();
          }

          // Keep all unread bytes.
          carryOverBytes.write(
              incoming.array(), incoming.arrayOffset() + incoming.position(), incoming.remaining());
        } catch (CoderMalfunctionError error) {
          decoder.reset();
          carryOverBytes.reset();
          error.printStackTrace();
        }
      }
      publishIfNonEmpty(msg);
    }

    @Override
    public synchronized void print(boolean b) {
      buffer.append(b ? "true" : "false");
    }

    @Override
    public synchronized void print(char c) {
      buffer.append(c);
    }

    @Override
    public synchronized void print(int i) {
      buffer.append(i);
    }

    @Override
    public synchronized void print(long l) {
      buffer.append(l);
    }

    @Override
    public synchronized void print(float f) {
      buffer.append(f);
    }

    @Override
    public synchronized void print(double d) {
      buffer.append(d);
    }

    @Override
    public void print(char[] a) {
      boolean flush = false;
      for (char c : a) {
        if (c == '\n') {
          flush = true;
        }
      }
      String msg;
      synchronized (this) {
        buffer.append(a);
        if (!flush) {
          return;
        }
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
    }

    @Override
    public void print(String s) {
      boolean flush = s.indexOf('\n') >= 0;
      String msg;
      synchronized (this) {
        buffer.append(s);
        if (!flush) {
          return;
        }
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
    }

    @Override
    public void print(Object o) {
      print(o.toString());
    }

    @Override
    public void println() {
      flush();
    }

    @Override
    public void println(boolean b) {
      String msg;
      synchronized (this) {
        buffer.append(b);
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
    }

    @Override
    public void println(char c) {
      String msg;
      synchronized (this) {
        buffer.append(c);
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
    }

    @Override
    public void println(int i) {
      String msg;
      synchronized (this) {
        buffer.append(i);
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
    }

    @Override
    public void println(long l) {
      String msg;
      synchronized (this) {
        buffer.append(l);
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
    }

    @Override
    public void println(float f) {
      String msg;
      synchronized (this) {
        buffer.append(f);
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
    }

    @Override
    public void println(double d) {
      String msg;
      synchronized (this) {
        buffer.append(d);
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
    }

    @Override
    public void println(char[] a) {
      String msg;
      synchronized (this) {
        buffer.append(a);
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
    }

    @Override
    public void println(String s) {
      String msg;
      synchronized (this) {
        buffer.append(s);
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
    }

    @Override
    public void println(Object o) {
      String msg;
      synchronized (this) {
        buffer.append(o);
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
    }

    @Override
    public PrintStream format(String format, Object... args) {
      return format(Locale.getDefault(), format, args);
    }

    @Override
    public PrintStream format(Locale locale, String format, Object... args) {
      String msg;
      synchronized (this) {
        int startLength = buffer.length();
        Formatter formatter = new Formatter(buffer, locale);
        formatter.format(format, args);
        if (buffer.indexOf("\n", startLength) < 0) {
          return this;
        }
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
      return this;
    }

    @Override
    public PrintStream append(CharSequence cs, int start, int limit) {
      CharSequence subsequence = cs.subSequence(start, limit);
      boolean flush = false;
      for (int i = 0; i < subsequence.length(); ++i) {
        if (subsequence.charAt(i) == '\n') {
          flush = true;
          break;
        }
      }
      String msg;
      synchronized (this) {
        buffer.append(cs.subSequence(start, limit));
        if (!flush) {
          return this;
        }
        msg = flushBufferToString();
      }
      publishIfNonEmpty(msg);
      return this;
    }

    private void publishIfNonEmpty(String message) {
      if (message == null || message.isEmpty()) {
        return;
      }
      if (logger.isLoggable(messageLevel)) {
        if (OUTPUT_WARNING.compareAndSet(false, true)) {
          LogRecord log = new LogRecord(Level.WARNING, LOGGING_DISCLAIMER);
          log.setLoggerName(loggerName);
          handler.publish(log);
        }
        LogRecord log = new LogRecord(messageLevel, message);
        log.setLoggerName(loggerName);
        handler.publish(log);
      }
    }
  }

  /**
   * Creates a {@link PrintStream} which redirects all output to the JUL {@link Handler} with the
   * specified {@code loggerName} and {@code level}.
   */
  static PrintStream create(
      Handler handler, String loggerName, Level messageLevel, Charset charset) {
    try {
      return new JulHandlerPrintStream(handler, loggerName, messageLevel, charset);
    } catch (UnsupportedEncodingException exc) {
      throw new RuntimeException("Encoding not supported: " + charset.name(), exc);
    }
  }

  static void reset() {
    OUTPUT_WARNING.set(false);
  }
}
