/*
 * 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.arrow.flight;

import com.google.common.util.concurrent.SettableFuture;
import io.grpc.stub.StreamObserver;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.arrow.flight.ArrowMessage.HeaderType;
import org.apache.arrow.flight.grpc.StatusUtils;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.util.VisibleForTesting;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.VectorLoader;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.dictionary.Dictionary;
import org.apache.arrow.vector.dictionary.DictionaryProvider;
import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch;
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
import org.apache.arrow.vector.types.MetadataVersion;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.DictionaryUtility;
import org.apache.arrow.vector.validate.MetadataV4UnionChecker;

/** An adaptor between protobuf streams and flight data streams. */
public class FlightStream implements AutoCloseable {
  // Use AutoCloseable sentinel objects to simplify logic in #close
  private final AutoCloseable DONE =
      new AutoCloseable() {
        @Override
        public void close() throws Exception {}
      };
  private final AutoCloseable DONE_EX =
      new AutoCloseable() {
        @Override
        public void close() throws Exception {}
      };

  private final BufferAllocator allocator;
  private final Cancellable cancellable;
  private final LinkedBlockingQueue<AutoCloseable> queue = new LinkedBlockingQueue<>();
  private final SettableFuture<VectorSchemaRoot> root = SettableFuture.create();
  private final SettableFuture<FlightDescriptor> descriptor = SettableFuture.create();
  private final int pendingTarget;
  private final Requestor requestor;
  // The completion flags.
  // This flag is only updated as the user iterates through the data, i.e. it tracks whether the
  // user has read all the
  // data and closed the stream
  final CompletableFuture<Void> completed;
  // This flag is immediately updated when gRPC signals that the server has ended the call. This is
  // used to make sure
  // we don't block forever trying to write to a server that has rejected a call.
  final CompletableFuture<Void> cancelled;

  private final AtomicInteger pending = new AtomicInteger();
  private volatile VectorSchemaRoot fulfilledRoot;
  private DictionaryProvider.MapDictionaryProvider dictionaries;
  private volatile VectorLoader loader;
  private volatile Throwable ex;
  private volatile ArrowBuf applicationMetadata = null;
  @VisibleForTesting volatile MetadataVersion metadataVersion = null;

  /**
   * Constructs a new instance.
   *
   * @param allocator The allocator to use for creating/reallocating buffers for Vectors.
   * @param pendingTarget Target number of messages to receive.
   * @param cancellable Used to cancel mid-stream requests.
   * @param requestor A callback to determine how many pending items there are.
   */
  public FlightStream(
      BufferAllocator allocator, int pendingTarget, Cancellable cancellable, Requestor requestor) {
    Objects.requireNonNull(allocator);
    Objects.requireNonNull(requestor);
    this.allocator = allocator;
    this.pendingTarget = pendingTarget;
    this.cancellable = cancellable;
    this.requestor = requestor;
    this.dictionaries = new DictionaryProvider.MapDictionaryProvider();
    this.completed = new CompletableFuture<>();
    this.cancelled = new CompletableFuture<>();
  }

  /** Get the schema for this stream. Blocks until the schema is available. */
  public Schema getSchema() {
    return getRoot().getSchema();
  }

  /**
   * Get the provider for dictionaries in this stream.
   *
   * <p>Does NOT retain a reference to the underlying dictionaries. Dictionaries may be updated as
   * the stream is read. This method is intended for stream processing, where the application code
   * will not retain references to values after the stream is closed.
   *
   * @throws IllegalStateException if {@link #takeDictionaryOwnership()} was called
   * @see #takeDictionaryOwnership()
   */
  public DictionaryProvider getDictionaryProvider() {
    if (dictionaries == null) {
      throw new IllegalStateException("Dictionary ownership was claimed by the application.");
    }
    return dictionaries;
  }

  /**
   * Get an owned reference to the dictionaries in this stream. Should be called after finishing
   * reading the stream, but before closing.
   *
   * <p>If called, the client is responsible for closing the dictionaries in this provider. Can only
   * be called once.
   *
   * @return The dictionary provider for the stream.
   * @throws IllegalStateException if called more than once.
   */
  public DictionaryProvider takeDictionaryOwnership() {
    if (dictionaries == null) {
      throw new IllegalStateException("Dictionary ownership was claimed by the application.");
    }
    // Swap out the provider so it is not closed
    final DictionaryProvider provider = dictionaries;
    dictionaries = null;
    return provider;
  }

  /**
   * Get the descriptor for this stream. Only applicable on the server side of a DoPut operation.
   * Will block until the client sends the descriptor.
   */
  public FlightDescriptor getDescriptor() {
    // This blocks until the first message from the client is received.
    try {
      return descriptor.get();
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      throw CallStatus.INTERNAL.withCause(e).withDescription("Interrupted").toRuntimeException();
    } catch (ExecutionException e) {
      throw CallStatus.INTERNAL
          .withCause(e)
          .withDescription("Error getting descriptor")
          .toRuntimeException();
    }
  }

  /**
   * Closes the stream (freeing any existing resources).
   *
   * <p>If the stream isn't complete and is cancellable, this method will cancel and drain the
   * stream first.
   */
  @Override
  public void close() throws Exception {
    final List<AutoCloseable> closeables = new ArrayList<>();
    Throwable suppressor = null;
    if (cancellable != null) {
      // Client-side stream. Cancel the call, to help ensure gRPC doesn't deliver a message after
      // close() ends.
      // On the server side, we can't rely on draining the stream , because this gRPC bug means the
      // completion callback
      // may never run https://github.com/grpc/grpc-java/issues/5882
      try {
        synchronized (cancellable) {
          if (!cancelled.isDone()) {
            // Only cancel if the call is not done on the gRPC side
            cancellable.cancel("Stream closed before end", /* no exception to report */ null);
          }
        }
        // Drain the stream without the lock (as next() implicitly needs the lock)
        while (next()) {}
      } catch (FlightRuntimeException e) {
        suppressor = e;
      }
    }
    // Perform these operations under a lock. This way the observer can't enqueue new messages while
    // we're in the
    // middle of cleanup. This should only be a concern for server-side streams since client-side
    // streams are drained
    // by the lambda above.
    synchronized (completed) {
      try {
        if (fulfilledRoot != null) {
          closeables.add(fulfilledRoot);
        }
        closeables.add(applicationMetadata);
        closeables.addAll(queue);
        if (dictionaries != null) {
          dictionaries
              .getDictionaryIds()
              .forEach(id -> closeables.add(dictionaries.lookup(id).getVector()));
        }
        if (suppressor != null) {
          AutoCloseables.close(suppressor, closeables);
        } else {
          AutoCloseables.close(closeables);
        }
        // Remove any metadata after closing to prevent negative refcnt
        applicationMetadata = null;
      } finally {
        // The value of this CompletableFuture is meaningless, only whether it's completed (or has
        // an exception)
        // No-op if already complete
        completed.complete(null);
      }
    }
  }

  /**
   * Blocking request to load next item into list.
   *
   * @return Whether or not more data was found.
   */
  public boolean next() {
    try {
      if (completed.isDone() && queue.isEmpty()) {
        return false;
      }

      pending.decrementAndGet();
      requestOutstanding();

      Object data = queue.take();
      if (DONE == data) {
        queue.put(DONE);
        // Other code ignores the value of this CompletableFuture, only whether it's completed (or
        // has an exception)
        completed.complete(null);
        return false;
      } else if (DONE_EX == data) {
        queue.put(DONE_EX);
        if (ex instanceof Exception) {
          throw (Exception) ex;
        } else {
          throw new Exception(ex);
        }
      } else {
        try (ArrowMessage msg = ((ArrowMessage) data)) {
          if (msg.getMessageType() == HeaderType.NONE) {
            updateMetadata(msg);
            // We received a message without data, so erase any leftover data
            if (fulfilledRoot != null) {
              fulfilledRoot.clear();
            }
          } else if (msg.getMessageType() == HeaderType.RECORD_BATCH) {
            checkMetadataVersion(msg);
            // Ensure we have the root
            root.get().clear();
            try (ArrowRecordBatch arb = msg.asRecordBatch()) {
              loader.load(arb);
            }
            updateMetadata(msg);
          } else if (msg.getMessageType() == HeaderType.DICTIONARY_BATCH) {
            checkMetadataVersion(msg);
            // Ensure we have the root
            root.get().clear();
            try (ArrowDictionaryBatch arb = msg.asDictionaryBatch()) {
              final long id = arb.getDictionaryId();
              if (dictionaries == null) {
                throw new IllegalStateException(
                    "Dictionary ownership was claimed by the application.");
              }
              final Dictionary dictionary = dictionaries.lookup(id);
              if (dictionary == null) {
                throw new IllegalArgumentException("Dictionary not defined in schema: ID " + id);
              }

              final FieldVector vector = dictionary.getVector();
              final VectorSchemaRoot dictionaryRoot =
                  new VectorSchemaRoot(
                      Collections.singletonList(vector.getField()),
                      Collections.singletonList(vector),
                      0);
              final VectorLoader dictionaryLoader = new VectorLoader(dictionaryRoot);
              dictionaryLoader.load(arb.getDictionary());
            }
            return next();
          } else {
            throw new UnsupportedOperationException(
                "Message type is unsupported: " + msg.getMessageType());
          }
          return true;
        }
      }
    } catch (RuntimeException e) {
      throw e;
    } catch (ExecutionException e) {
      throw StatusUtils.fromThrowable(e.getCause());
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
  }

  /** Update our metadata reference with a new one from this message. */
  private void updateMetadata(ArrowMessage msg) {
    if (this.applicationMetadata != null) {
      this.applicationMetadata.close();
    }
    this.applicationMetadata = msg.getApplicationMetadata();
    if (this.applicationMetadata != null) {
      this.applicationMetadata.getReferenceManager().retain();
    }
  }

  /** Ensure the Arrow metadata version doesn't change mid-stream. */
  private void checkMetadataVersion(ArrowMessage msg) {
    if (msg.asSchemaMessage() == null) {
      return;
    }
    MetadataVersion receivedVersion =
        MetadataVersion.fromFlatbufID(msg.asSchemaMessage().getMessage().version());
    if (this.metadataVersion != receivedVersion) {
      throw new IllegalStateException(
          "Metadata version mismatch: stream started as "
              + this.metadataVersion
              + " but got message with version "
              + receivedVersion);
    }
  }

  /**
   * Get the current vector data from the stream.
   *
   * <p>The data in the root may change at any time. Clients should NOT modify the root, but instead
   * unload the data into their own root.
   *
   * @throws FlightRuntimeException if there was an error reading the schema from the stream.
   */
  public VectorSchemaRoot getRoot() {
    try {
      return root.get();
    } catch (InterruptedException e) {
      throw CallStatus.INTERNAL.withCause(e).toRuntimeException();
    } catch (ExecutionException e) {
      throw StatusUtils.fromThrowable(e.getCause());
    }
  }

  /**
   * Check if there is a root (i.e. whether the other end has started sending data).
   *
   * <p>Updated by calls to {@link #next()}.
   *
   * @return true if and only if the other end has started sending data.
   */
  public boolean hasRoot() {
    return root.isDone();
  }

  /**
   * Get the most recent metadata sent from the server. This may be cleared by calls to {@link
   * #next()} if the server sends a message without metadata. This does NOT take ownership of the
   * buffer - call retain() to create a reference if you need the buffer after a call to {@link
   * #next()}.
   *
   * @return the application metadata. May be null.
   */
  public ArrowBuf getLatestMetadata() {
    return applicationMetadata;
  }

  private synchronized void requestOutstanding() {
    if (pending.get() < pendingTarget) {
      requestor.request(pendingTarget - pending.get());
      pending.set(pendingTarget);
    }
  }

  private class Observer implements StreamObserver<ArrowMessage> {

    Observer() {
      super();
    }

    /** Helper to add an item to the queue under the appropriate lock. */
    private void enqueue(AutoCloseable message) {
      synchronized (completed) {
        if (completed.isDone()) {
          // The stream is already closed (RPC ended), discard the message
          AutoCloseables.closeNoChecked(message);
        } else {
          queue.add(message);
        }
      }
    }

    @Override
    public void onNext(ArrowMessage msg) {
      // Operations here have to be under a lock so that we don't add a message to the queue while
      // in the middle of
      // close().
      requestOutstanding();
      switch (msg.getMessageType()) {
        case NONE:
          {
            // No IPC message - pure metadata or descriptor
            if (msg.getDescriptor() != null) {
              descriptor.set(new FlightDescriptor(msg.getDescriptor()));
            }
            if (msg.getApplicationMetadata() != null) {
              enqueue(msg);
            }
            break;
          }
        case SCHEMA:
          {
            Schema schema = msg.asSchema();

            // if there is app metadata in the schema message, make sure
            // that we don't leak it.
            ArrowBuf meta = msg.getApplicationMetadata();
            if (meta != null) {
              meta.close();
            }

            final List<Field> fields = new ArrayList<>();
            final Map<Long, Dictionary> dictionaryMap = new HashMap<>();
            for (final Field originalField : schema.getFields()) {
              final Field updatedField =
                  DictionaryUtility.toMemoryFormat(originalField, allocator, dictionaryMap);
              fields.add(updatedField);
            }
            for (final Map.Entry<Long, Dictionary> entry : dictionaryMap.entrySet()) {
              dictionaries.put(entry.getValue());
            }
            schema = new Schema(fields, schema.getCustomMetadata());
            metadataVersion =
                MetadataVersion.fromFlatbufID(msg.asSchemaMessage().getMessage().version());
            try {
              MetadataV4UnionChecker.checkRead(schema, metadataVersion);
            } catch (IOException e) {
              ex = e;
              enqueue(DONE_EX);
              break;
            }

            synchronized (completed) {
              if (!completed.isDone()) {
                fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
                loader = new VectorLoader(fulfilledRoot);
                if (msg.getDescriptor() != null) {
                  descriptor.set(new FlightDescriptor(msg.getDescriptor()));
                }
                root.set(fulfilledRoot);
              }
            }
            break;
          }
        case RECORD_BATCH:
        case DICTIONARY_BATCH:
          enqueue(msg);
          break;
        case TENSOR:
        default:
          ex =
              new UnsupportedOperationException(
                  "Unable to handle message of type: " + msg.getMessageType());
          enqueue(DONE_EX);
      }
    }

    @Override
    public void onError(Throwable t) {
      ex = StatusUtils.fromThrowable(t);
      queue.add(DONE_EX);
      cancelled.complete(null);
      root.setException(ex);
    }

    @Override
    public void onCompleted() {
      // Depends on gRPC calling onNext and onCompleted non-concurrently
      cancelled.complete(null);
      queue.add(DONE);
    }
  }

  /**
   * Cancels sending the stream to a client.
   *
   * <p>Callers should drain the stream (with {@link #next()}) to ensure all messages sent before
   * cancellation are received and to wait for the underlying transport to acknowledge cancellation.
   */
  public void cancel(String message, Throwable exception) {
    if (cancellable == null) {
      throw new UnsupportedOperationException(
          "Streams cannot be cancelled that are produced by client. "
              + "Instead, server should reject incoming messages.");
    }
    cancellable.cancel(message, exception);
    // Do not mark the stream as completed, as gRPC may still be delivering messages.
  }

  StreamObserver<ArrowMessage> asObserver() {
    return new Observer();
  }

  /** Provides a callback to cancel a process that is in progress. */
  @FunctionalInterface
  public interface Cancellable {
    void cancel(String message, Throwable exception);
  }

  /** Provides a interface to request more items from a stream producer. */
  @FunctionalInterface
  public interface Requestor {
    /** Requests <code>count</code> more messages from the instance of this object. */
    void request(int count);
  }
}
