package ai.verta.modeldb;

import ai.verta.modeldb.common.config.TestConfig;
import ai.verta.modeldb.common.config.TestUser;
import io.grpc.*;

@SuppressWarnings("unchecked")
public class AuthClientInterceptor {

  private String client1Email;
  private String client1DevKey;
  private String client2Email;
  private String client2DevKey;
  private Client1AuthInterceptor client1AuthInterceptor;
  private Client2AuthInterceptor client2AuthInterceptor;

  public AuthClientInterceptor(TestConfig testConfig) {
    client1AuthInterceptor =
        new Client1AuthInterceptor(testConfig.testUsers.getOrDefault("primaryUser", null));
    client2AuthInterceptor =
        new Client2AuthInterceptor(testConfig.testUsers.getOrDefault("secondaryUser", null));
  }

  public Client1AuthInterceptor getClient1AuthInterceptor() {
    return client1AuthInterceptor;
  }

  public Client2AuthInterceptor getClient2AuthInterceptor() {
    return client2AuthInterceptor;
  }

  public String getClient1Email() {
    return client1Email;
  }

  public String getClient2Email() {
    return client2Email;
  }

  private class Client1AuthInterceptor implements ClientInterceptor {

    public Client1AuthInterceptor(TestUser user) {
      client1Email = user.email;
      client1DevKey = user.devKey;
    }

    @Override
    public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
        MethodDescriptor<ReqT, RespT> methodDescriptor, CallOptions callOptions, Channel channel) {
      return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(
          channel.newCall(methodDescriptor, callOptions)) {
        @Override
        public void start(Listener<RespT> responseListener, Metadata headers) {
          // TODO: Here set request metadata
          Metadata.Key<String> email_key =
              Metadata.Key.of("email", Metadata.ASCII_STRING_MARSHALLER);
          Metadata.Key<String> dev_key =
              Metadata.Key.of("developer_key", Metadata.ASCII_STRING_MARSHALLER);
          Metadata.Key<String> source_key =
              Metadata.Key.of("source", Metadata.ASCII_STRING_MARSHALLER);

          headers.put(email_key, client1Email);
          headers.put(dev_key, client1DevKey);
          headers.put(source_key, "PythonClient");
          super.start(responseListener, headers);
        }
      };
    }
  }

  private class Client2AuthInterceptor implements ClientInterceptor {

    Client2AuthInterceptor(TestUser user) {
      client2Email = user.email;
      client2DevKey = user.devKey;
    }

    public String getClient2Email() {
      return client2Email;
    }

    @Override
    public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
        MethodDescriptor<ReqT, RespT> methodDescriptor, CallOptions callOptions, Channel channel) {
      return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(
          channel.newCall(methodDescriptor, callOptions)) {
        @Override
        public void start(Listener<RespT> responseListener, Metadata headers) {
          // TODO: Here set request metadata
          Metadata.Key<String> email_key =
              Metadata.Key.of("email", Metadata.ASCII_STRING_MARSHALLER);
          Metadata.Key<String> dev_key =
              Metadata.Key.of("developer_key", Metadata.ASCII_STRING_MARSHALLER);
          Metadata.Key<String> source_key =
              Metadata.Key.of("source", Metadata.ASCII_STRING_MARSHALLER);

          headers.put(email_key, client2Email);
          headers.put(dev_key, client2DevKey);
          headers.put(source_key, "PythonClient");
          super.start(responseListener, headers);
        }
      };
    }
  }
}

@@ -58,20 +58,27 @@
 .config.
-Test
+ServiceUser
 Config;%0A
@@ -119,20 +119,22 @@
 fig.Test
-User
+Config
 ;%0Aimport
@@ -1198,32 +1198,41 @@
 Interceptor(
-TestUser
+ServiceUserConfig
  user) %7B%0A   
@@ -2451,16 +2451,25 @@
 tor(
-TestUser
+ServiceUserConfig
  use

Remove references to TestUser so tests can compile (#2261)