package io.grpc.gcs;

import static io.grpc.gcs.Args.METHOD_RANDOM;
import static io.grpc.gcs.Args.METHOD_READ;

import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
import com.google.cloud.hadoop.gcsio.GoogleCloudStorageFileSystem;
import com.google.cloud.hadoop.gcsio.GoogleCloudStorageFileSystemOptions;
import com.google.cloud.hadoop.gcsio.GoogleCloudStorageOptions;
import com.google.cloud.hadoop.gcsio.GoogleCloudStorageReadOptions;
import com.google.cloud.hadoop.gcsio.GoogleCloudStorageReadOptions.Fadvise;
import java.io.IOException;
import java.net.URI;
import java.nio.ByteBuffer;
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.SeekableByteChannel;
import java.util.Arrays;
import java.util.List;
import java.util.Random;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.logging.Logger;

public class GcsioClient {
  private static final String SCOPE = "https://www.googleapis.com/auth/cloud-platform";

  private Args args;
  private GoogleCloudStorageFileSystem gcsfs;
  private static final Logger logger = Logger.getLogger(GcsioClient.class.getName());

  public GcsioClient(Args args, boolean grpcEnabled) throws IOException {
    this.args = args;
    GoogleCredential creds = GoogleCredential.getApplicationDefault().createScoped(Arrays.asList(SCOPE));

    GoogleCloudStorageOptions gcsOpts =
        GoogleCloudStorageOptions.builder()
            .setAppName("weiranf-app")
            .setGrpcEnabled(grpcEnabled)
            .build();

    this.gcsfs = new GoogleCloudStorageFileSystem(creds,
        GoogleCloudStorageFileSystemOptions.builder()
            .setCloudStorageOptions(gcsOpts)
            .build()
    );

  }

  public void startCalls(List<Long> results) throws InterruptedException, IOException {
    if (args.threads == 0) {
      try {
        switch (args.method) {
          case METHOD_READ:
            makeMediaRequest(results);
            break;
          case METHOD_RANDOM:
            makeRandomMediaRequest(results);
            break;
          default:
            logger.warning("Please provide valid methods with --method");
        }
      } finally {
        gcsfs.close();
      }
    } else {
      ThreadPoolExecutor threadPoolExecutor =
          (ThreadPoolExecutor) Executors.newFixedThreadPool(args.threads);
      try {
        switch (args.method) {
          case METHOD_READ:
            for (int i = 0; i < args.threads; i++) {
              Runnable task = () -> {
                try {
                  makeMediaRequest(results);
                } catch (IOException e) {
                  e.printStackTrace();
                }
              };
              threadPoolExecutor.execute(task);
            }
            break;
          default:
            logger.warning("Please provide valid methods with --method");
        }
        threadPoolExecutor.shutdown();
        if (!threadPoolExecutor.awaitTermination(30, TimeUnit.MINUTES)) {
          threadPoolExecutor.shutdownNow();
        }
      } finally {
        threadPoolExecutor.shutdownNow();
        gcsfs.close();
      }
    }
  }

  private void makeMediaRequest(List<Long> results) throws IOException {
    int size = args.buffSize * 1024;

    URI uri = URI.create("gs://" + args.bkt + "/" + args.obj);

    ByteBuffer buff = ByteBuffer.allocate(size);
    for (int i = 0; i < args.calls; i++) {
      ReadableByteChannel readChannel = gcsfs.open(uri);
      long start = System.currentTimeMillis();
      readChannel.read(buff);
      long dur = System.currentTimeMillis() - start;
      if (buff.remaining() > 0) {
        logger.warning("Got remaining bytes: " + buff.remaining());
      }
      buff.clear();
      readChannel.close();
      //logger.info("time cost for reading bytes: " + dur + "ms");
      results.add(dur);
    }
  }

  private void makeRandomMediaRequest(List<Long> results) throws IOException {

    Random r = new Random();

    URI uri = URI.create("gs://" + args.bkt + "/" + args.obj);
    
    GoogleCloudStorageReadOptions readOpts = GoogleCloudStorageReadOptions.builder().build();

    SeekableByteChannel reader = gcsfs.open(uri, readOpts);
    for (int i = 0; i < args.calls; i++) {
      long offset = (long) r.nextInt(args.size - args.buffSize) * 1024;
      long start = System.currentTimeMillis();
      ByteBuffer buff = ByteBuffer.allocate(args.buffSize * 1024);
      reader.position(offset);
      reader.read(buff);
      long dur = System.currentTimeMillis() - start;
      if (buff.remaining() > 0) {
        logger.warning("Got remaining bytes: " + buff.remaining());
      }
      logger.info("time cost for reading bytes: " + dur + "ms");
      results.add(dur);
    }
    reader.close();
  }
}

@@ -104,16 +104,61 @@
 OD_READ;
+%0Aimport static io.grpc.gcs.Args.METHOD_WRITE;
 %0A%0Aimport
@@ -498,84 +498,8 @@
 ns;%0A
-import com.google.cloud.hadoop.gcsio.GoogleCloudStorageReadOptions.Fadvise;%0A
 impo
@@ -655,32 +655,78 @@
 bleByteChannel;%0A
+import java.nio.channels.WritableByteChannel;%0A
 import java.util
@@ -2029,32 +2029,100 @@
          break;%0A
+          case METHOD_WRITE:%0A            makeWriteRequest(results);%0A
           case M
@@ -3890,24 +3890,24 @@
 el.close();%0A
-
       //logg
@@ -3988,32 +3988,654 @@
 ur);%0A    %7D%0A  %7D%0A%0A
+  private void makeWriteRequest(List%3CLong%3E results) throws IOException %7B%0A    int size = args.buffSize * 1024;%0A    Random rd = new Random();%0A    byte%5B%5D randBytes = new byte%5Bsize%5D;%0A    rd.nextBytes(randBytes);%0A%0A    URI uri = URI.create(%22gs://%22 + args.bkt + %22/%22 + args.obj);%0A    for (int i = 0; i %3C args.calls; i++) %7B%0A      WritableByteChannel writeChannel = gcsfs.create(uri);%0A      ByteBuffer buff = ByteBuffer.wrap(randBytes);%0A      long start = System.currentTimeMillis();%0A      writeChannel.write(buff);%0A      long dur = System.currentTimeMillis() - start;%0A      writeChannel.close();%0A      results.add(dur);%0A    %7D%0A  %7D%0A%0A
   private void m

Add write operation for gcsio client