package ai.verta.modeldb.common.connections;

import ai.verta.modeldb.common.CommonConstants;
import ai.verta.modeldb.common.CommonMessages;
import ai.verta.modeldb.common.authservice.AuthInterceptor;
import ai.verta.modeldb.common.config.Config;
import ai.verta.modeldb.common.exceptions.UnavailableException;
import ai.verta.uac.CollaboratorServiceGrpc;
import io.grpc.ClientInterceptor;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.Metadata;
import io.grpc.stub.MetadataUtils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

public class UAC {
  private static final Logger LOGGER = LogManager.getLogger(UAC.class);

  private final ManagedChannel authServiceChannel;
  private final String serviceUserEmail;
  private final String serviceUserDevKey;

  private ClientInterceptor clientInterceptor = null;

  private final CollaboratorServiceGrpc.CollaboratorServiceFutureStub collaboratorServiceFutureStub;

  public static UAC FromConfig(Config config) {
    if (!config.hasAuth()) return null;
    else return new UAC(config);
  }

  private UAC(Config config) {
    this(
        config.authService.host,
        config.authService.port,
        config.service_user.email,
        config.service_user.devKey);
  }

  public UAC(String host, Integer port, String serviceUserEmail, String serviceUserDevKey) {
    LOGGER.trace(CommonMessages.HOST_PORT_INFO_STR, host, port);
    if (host != null && port != null) { // AuthService not available.
      authServiceChannel =
          ManagedChannelBuilder.forTarget(host + CommonConstants.STRING_COLON + port)
              .usePlaintext()
              .build();

      this.serviceUserEmail = serviceUserEmail;
      this.serviceUserDevKey = serviceUserDevKey;
    } else {
      throw new UnavailableException(
          "Host OR Port not found for contacting authentication service");
    }

    collaboratorServiceFutureStub = CollaboratorServiceGrpc.newFutureStub(authServiceChannel);
  }

  private UAC(UAC other) {
    authServiceChannel = other.authServiceChannel;
    serviceUserDevKey = other.serviceUserDevKey;
    serviceUserEmail = other.serviceUserEmail;

    collaboratorServiceFutureStub = other.collaboratorServiceFutureStub;
  }

  public UAC withServiceAccount() {
    UAC c = new UAC(this);
    Metadata requestHeaders = new 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> dev_key_hyphen =
        Metadata.Key.of("developer-key", Metadata.ASCII_STRING_MARSHALLER);
    Metadata.Key<String> source_key = Metadata.Key.of("source", Metadata.ASCII_STRING_MARSHALLER);

    requestHeaders.put(email_key, this.serviceUserEmail);
    requestHeaders.put(dev_key, this.serviceUserDevKey);
    requestHeaders.put(dev_key_hyphen, this.serviceUserDevKey);
    requestHeaders.put(source_key, "PythonClient");

    c.clientInterceptor = MetadataUtils.newAttachHeadersInterceptor(requestHeaders);

    return c;
  }

  public CollaboratorServiceGrpc.CollaboratorServiceFutureStub getCollaboratorService() {
    if (clientInterceptor != null) {
      return collaboratorServiceFutureStub.withInterceptors(clientInterceptor);
    }
    return collaboratorServiceFutureStub.withInterceptors(
        MetadataUtils.newAttachHeadersInterceptor(AuthInterceptor.METADATA_INFO.get()));
  }
}

@@ -349,16 +349,94 @@
 ceGrpc;%0A
+import ai.verta.uac.UACServiceGrpc;%0Aimport ai.verta.uac.WorkspaceServiceGrpc;%0A
 import i
@@ -1057,16 +1057,182 @@
 ureStub;
+%0A  private final UACServiceGrpc.UACServiceFutureStub uacServiceFutureStub;%0A  private final WorkspaceServiceGrpc.WorkspaceServiceFutureStub workspaceServiceFutureStub;
 %0A%0A  publ
@@ -1392,32 +1392,36 @@
   this(%0A        
+    
 config.authServi
@@ -1429,16 +1429,20 @@
 e.host,%0A
+    
         
@@ -1466,32 +1466,36 @@
 e.port,%0A        
+    
 config.service_u
@@ -1505,16 +1505,20 @@
 .email,%0A
+    
         
@@ -1816,16 +1816,20 @@
         
+    
 ManagedC
@@ -1910,16 +1910,24 @@
         
+        
 .usePlai
@@ -1934,16 +1934,24 @@
 ntext()%0A
+        
         
@@ -2126,16 +2126,20 @@
         
+    
 %22Host OR
@@ -2297,16 +2297,182 @@
 annel);%0A
+    uacServiceFutureStub = UACServiceGrpc.newFutureStub(authServiceChannel);%0A    workspaceServiceFutureStub = WorkspaceServiceGrpc.newFutureStub(authServiceChannel);%0A
   %7D%0A%0A  p
@@ -2716,16 +2716,138 @@
 reStub;%0A
+    uacServiceFutureStub = other.uacServiceFutureStub;%0A    workspaceServiceFutureStub = other.workspaceServiceFutureStub;%0A
   %7D%0A%0A  p
@@ -3084,32 +3084,36 @@
 v_key =%0A        
+    
 Metadata.Key.of(
@@ -3206,16 +3206,20 @@
 yphen =%0A
+    
         
@@ -3995,16 +3995,700 @@
 eptors(%0A
+            MetadataUtils.newAttachHeadersInterceptor(AuthInterceptor.METADATA_INFO.get()));%0A  %7D%0A%0A  public UACServiceGrpc.UACServiceFutureStub getUACService() %7B%0A    if (clientInterceptor != null) %7B%0A      return uacServiceFutureStub.withInterceptors(clientInterceptor);%0A    %7D%0A    return uacServiceFutureStub.withInterceptors(%0A            MetadataUtils.newAttachHeadersInterceptor(AuthInterceptor.METADATA_INFO.get()));%0A  %7D%0A%0A  public WorkspaceServiceGrpc.WorkspaceServiceFutureStub getWorkspaceService() %7B%0A    if (clientInterceptor != null) %7B%0A      return workspaceServiceFutureStub.withInterceptors(clientInterceptor);%0A    %7D%0A    return workspaceServiceFutureStub.withInterceptors(%0A    
         

Add getters that do NOT block for the UAC and Workspace services (#2041)