/*
 * 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.integration.tests;

import java.io.IOException;
import java.net.URISyntaxException;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import org.apache.arrow.flight.Action;
import org.apache.arrow.flight.ActionType;
import org.apache.arrow.flight.CallStatus;
import org.apache.arrow.flight.CancelFlightInfoRequest;
import org.apache.arrow.flight.CancelFlightInfoResult;
import org.apache.arrow.flight.CancelStatus;
import org.apache.arrow.flight.FlightConstants;
import org.apache.arrow.flight.FlightDescriptor;
import org.apache.arrow.flight.FlightEndpoint;
import org.apache.arrow.flight.FlightInfo;
import org.apache.arrow.flight.Location;
import org.apache.arrow.flight.NoOpFlightProducer;
import org.apache.arrow.flight.RenewFlightEndpointRequest;
import org.apache.arrow.flight.Result;
import org.apache.arrow.flight.Ticket;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.UInt4Vector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;

/**
 * The server used for testing FlightEndpoint.expiration_time.
 *
 * <p>GetFlightInfo() returns a FlightInfo that has the following three FlightEndpoints:
 *
 * <ol>
 *   <li>No expiration time
 *   <li>5 seconds expiration time
 *   <li>6 seconds expiration time
 * </ol>
 *
 * The client can't read data from the first endpoint multiple times but can read data from the
 * second and third endpoints. The client can't re-read data from the second endpoint 5 seconds
 * later. The client can't re-read data from the third endpoint 6 seconds later.
 *
 * <p>The client can cancel a returned FlightInfo by pre-defined CancelFlightInfo action. The client
 * can't read data from endpoints even within 6 seconds after the action.
 *
 * <p>The client can extend the expiration time of a FlightEndpoint in a returned FlightInfo by
 * pre-defined RenewFlightEndpoint action. The client can read data from endpoints multiple times
 * within more 10 seconds after the action.
 */
final class ExpirationTimeProducer extends NoOpFlightProducer {
  public static final Schema SCHEMA =
      new Schema(
          Collections.singletonList(Field.notNullable("number", Types.MinorType.UINT4.getType())));

  private final BufferAllocator allocator;
  private final List<EndpointStatus> statuses;

  ExpirationTimeProducer(BufferAllocator allocator) {
    this.allocator = allocator;
    this.statuses = new ArrayList<>();
  }

  @Override
  public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) {
    statuses.clear();
    List<FlightEndpoint> endpoints = new ArrayList<>();
    Instant now = Instant.now();
    endpoints.add(addEndpoint("No expiration time", null));
    endpoints.add(addEndpoint("5 seconds", now.plus(5, ChronoUnit.SECONDS)));
    endpoints.add(addEndpoint("6 seconds", now.plus(6, ChronoUnit.SECONDS)));
    return new FlightInfo(SCHEMA, descriptor, endpoints, -1, -1);
  }

  @Override
  public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) {
    // Obviously, not safe (since we don't lock), but we assume calls are not concurrent
    int index = parseIndexFromTicket(ticket);
    EndpointStatus status = statuses.get(index);
    if (status.cancelled) {
      listener.error(
          CallStatus.NOT_FOUND
              .withDescription(
                  "Invalid flight: cancelled: "
                      + new String(ticket.getBytes(), StandardCharsets.UTF_8))
              .toRuntimeException());
      return;
    } else if (status.expirationTime != null && Instant.now().isAfter(status.expirationTime)) {
      listener.error(
          CallStatus.NOT_FOUND
              .withDescription(
                  "Invalid flight: expired: "
                      + new String(ticket.getBytes(), StandardCharsets.UTF_8))
              .toRuntimeException());
      return;
    } else if (status.expirationTime == null && status.numGets > 0) {
      listener.error(
          CallStatus.NOT_FOUND
              .withDescription(
                  "Invalid flight: can't read multiple times: "
                      + new String(ticket.getBytes(), StandardCharsets.UTF_8))
              .toRuntimeException());
      return;
    }
    status.numGets++;

    try (final VectorSchemaRoot root = VectorSchemaRoot.create(SCHEMA, allocator)) {
      listener.start(root);
      UInt4Vector vector = (UInt4Vector) root.getVector(0);
      vector.setSafe(0, index);
      root.setRowCount(1);
      listener.putNext();
    }
    listener.completed();
  }

  @Override
  public void doAction(CallContext context, Action action, StreamListener<Result> listener) {
    try {
      if (action.getType().equals(FlightConstants.CANCEL_FLIGHT_INFO.getType())) {
        CancelFlightInfoRequest request =
            CancelFlightInfoRequest.deserialize(ByteBuffer.wrap(action.getBody()));
        CancelStatus cancelStatus = CancelStatus.UNSPECIFIED;
        for (FlightEndpoint endpoint : request.getInfo().getEndpoints()) {
          int index = parseIndexFromTicket(endpoint.getTicket());
          EndpointStatus status = statuses.get(index);
          if (status.cancelled) {
            cancelStatus = CancelStatus.NOT_CANCELLABLE;
          } else {
            status.cancelled = true;
            if (cancelStatus == CancelStatus.UNSPECIFIED) {
              cancelStatus = CancelStatus.CANCELLED;
            }
          }
        }
        listener.onNext(new Result(new CancelFlightInfoResult(cancelStatus).serialize().array()));
      } else if (action.getType().equals(FlightConstants.RENEW_FLIGHT_ENDPOINT.getType())) {
        RenewFlightEndpointRequest request =
            RenewFlightEndpointRequest.deserialize(ByteBuffer.wrap(action.getBody()));
        FlightEndpoint endpoint = request.getFlightEndpoint();
        int index = parseIndexFromTicket(endpoint.getTicket());
        EndpointStatus status = statuses.get(index);
        if (status.cancelled) {
          listener.onError(
              CallStatus.INVALID_ARGUMENT
                  .withDescription("Invalid flight: cancelled: " + index)
                  .toRuntimeException());
          return;
        }

        String ticketBody = new String(endpoint.getTicket().getBytes(), StandardCharsets.UTF_8);
        ticketBody += ": renewed (+ 10 seconds)";
        Ticket ticket = new Ticket(ticketBody.getBytes(StandardCharsets.UTF_8));
        Instant expiration = Instant.now().plus(10, ChronoUnit.SECONDS);
        status.expirationTime = expiration;
        FlightEndpoint newEndpoint =
            new FlightEndpoint(
                ticket, expiration, endpoint.getLocations().toArray(new Location[0]));
        listener.onNext(new Result(newEndpoint.serialize().array()));
      } else {
        listener.onError(
            CallStatus.INVALID_ARGUMENT
                .withDescription("Unknown action: " + action.getType())
                .toRuntimeException());
        return;
      }
    } catch (IOException | URISyntaxException e) {
      listener.onError(
          CallStatus.INTERNAL.withCause(e).withDescription(e.toString()).toRuntimeException());
      return;
    }
    listener.onCompleted();
  }

  @Override
  public void listActions(CallContext context, StreamListener<ActionType> listener) {
    listener.onNext(FlightConstants.CANCEL_FLIGHT_INFO);
    listener.onNext(FlightConstants.RENEW_FLIGHT_ENDPOINT);
    listener.onCompleted();
  }

  private FlightEndpoint addEndpoint(String ticket, Instant expirationTime) {
    Ticket flightTicket =
        new Ticket(
            String.format("%d: %s", statuses.size(), ticket).getBytes(StandardCharsets.UTF_8));
    statuses.add(new EndpointStatus(expirationTime));
    return new FlightEndpoint(flightTicket, expirationTime);
  }

  private int parseIndexFromTicket(Ticket ticket) {
    final String contents = new String(ticket.getBytes(), StandardCharsets.UTF_8);
    int index = contents.indexOf(':');
    if (index == -1) {
      throw CallStatus.INVALID_ARGUMENT
          .withDescription(
              "Invalid ticket: " + new String(ticket.getBytes(), StandardCharsets.UTF_8))
          .toRuntimeException();
    }
    int endpointIndex = Integer.parseInt(contents.substring(0, index));
    if (endpointIndex < 0 || endpointIndex >= statuses.size()) {
      throw CallStatus.NOT_FOUND.withDescription("Out of bounds").toRuntimeException();
    }
    return endpointIndex;
  }

  /** The status of a returned endpoint. */
  static final class EndpointStatus {
    Instant expirationTime;
    int numGets;
    boolean cancelled;

    EndpointStatus(Instant expirationTime) {
      this.expirationTime = expirationTime;
      this.numGets = 0;
      this.cancelled = false;
    }
  }
}
