/*
 * Copyright © 2021 Celeral.
 *
 * Licensed 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 com.celeral.netlet.rpc;

import java.io.Closeable;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.UnaryOperator;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import com.celeral.netlet.codec.StatefulStreamCodec;
import com.celeral.utils.WeakIdentityHashMap;

public class DelegationTransport implements InvocationHandler, Closeable {
  private final ConcurrentLinkedQueue<ProxyProvider.RPCFuture> futureResponses;
  public final ProxyProvider.DelegatingClient client;
  private ConnectionAgent agent;
  private TimeoutPolicy policy;

  private final ConcurrentLinkedQueue<Object> deletedIdentifiers;
  private final WeakIdentityHashMap<Object, Object> identityMap;

  public DelegationTransport(
      ConnectionAgent agent,
      TimeoutPolicy policy,
      MethodSerializer<?> methodSerializer,
      UnaryOperator<StatefulStreamCodec<Object>> serdesConfigurator,
      Executor executor) {
    deletedIdentifiers = new ConcurrentLinkedQueue<>();
    identityMap = new WeakIdentityHashMap<>(1, deletedIdentifiers::add);

    this.agent = agent;
    this.policy = policy;
    this.futureResponses = new ConcurrentLinkedQueue<>();
    this.client =
        new ProxyProvider.DelegatingClient(
            futureResponses, agent, methodSerializer, executor, serdesConfigurator);
  }

  public void register(Object identifier, Object proxy) {
    synchronized (identityMap) {
      identityMap.put(proxy, identifier);
    }
  }

  @Override
  public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
    Object identity = identityMap.get(proxy);
    do {
      if (!client.isConnected()) {
        agent.connect(client);
      }

      logger.trace("calling {}", method);
      Object[] deletedIdentifiers;
      if (this.deletedIdentifiers.isEmpty()) {
        deletedIdentifiers = null;
      } else {
        deletedIdentifiers = this.deletedIdentifiers.toArray();
        logger.debug(
            "requesting deletion of objects with identifiers = {}", this.deletedIdentifiers);
      }
      ProxyProvider.RPCFuture future =
          new ProxyProvider.RPCFuture(client.send(identity, method, args, deletedIdentifiers));
      futureResponses.add(future);

      try {
        final Object response = future.get(policy.getTimeoutMillis(), TimeUnit.MILLISECONDS);
        if (deletedIdentifiers != null) {
          for (Object identifier : deletedIdentifiers) {
            if (!this.deletedIdentifiers.remove(identifier)) {
              logger.warn(
                  "identifier {} just removed from the list is missing from the list!", identifier);
            }
          }
        }
        return response;
      } catch (TimeoutException ex) {
        policy.handleTimeout(this, ex);
      } catch (ExecutionException ex) {
        throw ex.getCause();
      }
    } while (true);
  }

  @Override
  public void close() {
    if (client.isConnected()) {
      agent.disconnect(client);
    }
  }

  public ConnectionAgent getAgent()
  {
    return agent;
  }

  public void setAgent(ConnectionAgent agent)
  {
    this.agent = agent;
  }

  public TimeoutPolicy getPolicy()
  {
    return policy;
  }

  public void setPolicy(TimeoutPolicy policy)
  {
    this.policy = policy;
  }

  public static final Logger logger = LogManager.getLogger(DelegationTransport.class);
}

@@ -719,16 +719,41 @@
 Method;%0A
+import java.util.Arrays;%0A
 import j
@@ -2665,16 +2665,65 @@
 ifiers;%0A
+      synchronized (this.deletedIdentifiers) %7B%0A  
       if
@@ -2757,24 +2757,26 @@
 sEmpty()) %7B%0A
+  
         dele
@@ -2808,17 +2808,27 @@
     
-%7D
+  %7D%0A       
  else %7B%0A
     
@@ -2819,24 +2819,26 @@
      else %7B%0A
+  
         dele
@@ -2893,24 +2893,69 @@
 ();%0A        
+  this.deletedIdentifiers.clear();%0A          
 logger.debug
@@ -2955,29 +2955,16 @@
 r.debug(
-%0A            
 %22request
@@ -3011,21 +3011,24 @@
  = %7B%7D%22, 
-this.
+(Object)
 deletedI
@@ -3040,24 +3040,34 @@
 fiers);%0A
+        %7D%0A
       %7D%0A
       Pr
@@ -3054,24 +3054,25 @@
   %7D%0A      %7D%0A
+%0A
       ProxyP
@@ -3263,24 +3263,17 @@
  Object 
-response
+o
  = futur
@@ -3337,20 +3337,16 @@
         
-if (
 deletedI
@@ -3360,18 +3360,15 @@
 ers 
-!
 = null
-) %7B
+;
 %0A   
@@ -3376,236 +3376,307 @@
     
-  for (Object identifier : deletedIdentifiers) %7B%0A            if (!this.deletedIdentifiers.remove(identifier)) %7B%0A              logger.warn(%0A                  %22identifier %7B%7D just removed from the list is missing from the list!%22, i
+return o;%0A      %7D catch (TimeoutException ex) %7B%0A        policy.handleTimeout(this, ex);%0A      %7D catch (ExecutionException ex) %7B%0A        throw ex.getCause();%0A      %7D finally %7B%0A        if (deletedIdentifiers != null) %7B%0A          logger.info(%22Adding back garbage collected objects %7B%7D%22, (Object)deletedI
 dent
@@ -3676,24 +3676,25 @@
 edIdentifier
+s
 );%0A         
@@ -3698,205 +3698,149 @@
     
-  %7D%0A          %7D%0A        %7D%0A        return response;%0A      %7D catch (TimeoutException ex) %7B%0A        policy.handleTimeout(this, ex);%0A      %7D catch (ExecutionException ex) %7B%0A        throw ex.getCause();
+synchronized (this.deletedIdentifiers) %7B%0A            this.deletedIdentifiers.addAll(Arrays.asList(deletedIdentifiers));%0A          %7D%0A        %7D
 %0A   

Fix the race condition for deletion of the garbage collected object on the delegation side which resulted in repeated request for deletion on the execution side.