/**
 * 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
 *
 *     https://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.avro.grpc;

import org.apache.avro.AvroRuntimeException;
import org.apache.avro.grpc.test.Kind;
import org.apache.avro.grpc.test.MD5;
import org.apache.avro.grpc.test.TestError;
import org.apache.avro.grpc.test.TestRecord;
import org.apache.avro.grpc.test.TestService;
import org.apache.avro.ipc.CallFuture;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.util.Arrays;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.Server;
import io.grpc.ServerBuilder;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;

public class TestAvroProtocolGrpc {
  private final TestRecord record = TestRecord.newBuilder().setName("foo").setKind(Kind.FOO)
      .setArrayOfLongs(Arrays.asList(42L, 424L, 4242L)).setHash(new MD5(new byte[] { 4, 2, 4, 2 }))
      .setNullableHash(null).build();
  private final String declaredErrMsg = "Declared error";
  private final String undeclaredErrMsg = "Undeclared error";
  private final TestError declaredError = TestError.newBuilder().setMessage$(declaredErrMsg).build();
  private final RuntimeException undeclaredError = new RuntimeException(undeclaredErrMsg);
  private CountDownLatch oneWayStart;
  private CountDownLatch oneWayDone;
  private AtomicInteger oneWayCount;
  private TestService stub;
  private TestService.Callback callbackStub;
  private Server server;
  private ManagedChannel channel;

  @BeforeEach
  public void setUp() throws IOException {
    TestService serviceImpl = new TestServiceImplBase();
    setUpServerAndClient(serviceImpl);
  }

  private void setUpServerAndClient(TestService serviceImpl) throws IOException {
    if (server != null && !server.isShutdown()) {
      server.shutdown();
    }
    if (channel != null && !channel.isShutdown()) {
      channel.shutdownNow();
    }
    server = ServerBuilder.forPort(0).addService(AvroGrpcServer.createServiceDefinition(TestService.class, serviceImpl))
        .build();
    server.start();
    int port = server.getPort();
    channel = ManagedChannelBuilder.forAddress("localhost", port).usePlaintext().build();
    stub = AvroGrpcClient.create(channel, TestService.class);
    callbackStub = AvroGrpcClient.create(channel, TestService.Callback.class);
  }

  @AfterEach
  public void cleanUp() {
    channel.shutdownNow();
    server.shutdownNow();
  }

  @Test
  void echoRecord() throws Exception {
    TestRecord echoedRecord = stub.echo(record);
    assertEquals(record, echoedRecord);
  }

  @Test
  void multipleArgsAdd() throws Exception {
    int result = stub.add(3, 5, 2);
    assertEquals(10, result);
  }

  @Test
  void multipleArgsConcatenate() throws Exception {
    String val1 = "foo-bar";
    Boolean val2 = true;
    long val3 = 123321L;
    int val4 = 42;
    assertEquals(val1 + val2 + val3 + val4, stub.concatenate(val1, val2, val3, val4));
  }

  @Test
  void callbackInterface() throws Exception {
    CallFuture<TestRecord> future = new CallFuture<>();
    callbackStub.echo(record, future);
    assertEquals(record, future.get(1, TimeUnit.SECONDS));
  }

  @Test
  void oneWayRpc() throws Exception {
    oneWayStart = new CountDownLatch(1);
    oneWayDone = new CountDownLatch(3);
    oneWayCount = new AtomicInteger();
    stub.ping();
    stub.ping();
    // client is not stalled while server is waiting for processing requests
    assertEquals(0, oneWayCount.get());
    oneWayStart.countDown();
    stub.ping();
    oneWayDone.await(1, TimeUnit.SECONDS);
    assertEquals(3, oneWayCount.get());
  }

  @Test
  void declaredError() throws Exception {
    try {
      stub.error(true);
      fail("Expected exception but none thrown");
    } catch (TestError te) {
      assertEquals(declaredErrMsg, te.getMessage$());
    }
  }

  @Test
  void undeclaredError() throws Exception {
    try {
      stub.error(false);
      fail("Expected exception but none thrown");
    } catch (AvroRuntimeException e) {
      assertTrue(e.getMessage().contains(undeclaredErrMsg));
    }
  }

  @Test
  void nullableResponse() throws Exception {
    setUpServerAndClient(new TestServiceImplBase() {
      @Override
      public String concatenate(String val1, boolean val2, long val3, int val4) {
        return null;
      }
    });
    assertNull(stub.concatenate("foo", true, 42L, 42));
  }

  @Test
  void grpcConnectionError() throws Exception {
    assertThrows(AvroRuntimeException.class, () -> {
      // close the channel and initiate request
      channel.shutdownNow();
      stub.add(0, 1, 2);
    });
  }

  @Test
  void repeatedRequests() throws Exception {
    TestRecord[] echoedRecords = new TestRecord[5];
    // validate results after all requests are done
    for (int i = 0; i < 5; i++) {
      echoedRecords[i] = stub.echo(record);
    }
    for (TestRecord result : echoedRecords) {
      assertEquals(record, result);
    }
  }

  @Test
  void concurrentClientAccess() throws Exception {
    ExecutorService es = Executors.newCachedThreadPool();
    Future<TestRecord>[] records = new Future[5];
    Future<Integer>[] adds = new Future[5];
    // submit requests in parallel
    for (int i = 0; i < 5; i++) {
      records[i] = es.submit(() -> stub.echo(record));
      int j = i;
      adds[i] = es.submit(() -> stub.add(j, 2 * j, 3 * j));
    }
    // validate all results
    for (int i = 0; i < 5; i++) {
      assertEquals(record, records[i].get());
      assertEquals(6 * i, (long) adds[i].get());
    }
  }

  @Test
  void concurrentChannels() throws Exception {
    ManagedChannel otherChannel = ManagedChannelBuilder.forAddress("localhost", server.getPort()).usePlaintext()
        .build();
    TestService otherStub = AvroGrpcClient.create(otherChannel, TestService.class);
    Future<Integer>[] adds = new Future[5];
    Future<Integer>[] otherAdds = new Future[5];
    ExecutorService es = Executors.newCachedThreadPool();
    // submit requests on clients with different channels
    for (int i = 0; i < 5; i++) {
      int j = i;
      adds[i] = es.submit(() -> stub.add(j, j - 1, j - 2));
      otherAdds[i] = es.submit(() -> otherStub.add(j, j + 1, j + 2));
    }
    // validate all results
    for (int i = 0; i < 5; i++) {
      assertEquals((3 * i) - 3, (long) adds[i].get());
      assertEquals((3 * i) + 3, (long) otherAdds[i].get());
    }
    otherChannel.shutdownNow();
  }

  private class TestServiceImplBase implements TestService {
    @Override
    public TestRecord echo(TestRecord record) {
      return record;
    }

    @Override
    public int add(int arg1, int arg2, int arg3) {
      return arg1 + arg2 + arg3;
    }

    @Override
    public void error(boolean declared) throws TestError {
      if (declared) {
        throw declaredError;
      }
      throw undeclaredError;
    }

    @Override
    public void ping() {
      try {
        oneWayStart.await();
        oneWayCount.incrementAndGet();
        oneWayDone.countDown();
      } catch (InterruptedException e) {
        fail("thread interrupted when waiting for all one-way messages");
      }
    }

    @Override
    public String concatenate(String val1, boolean val2, long val3, int val4) {
      return val1 + val2 + val3 + val4;
    }
  }
}
