/*
    * junixsocket
    *
    * Copyright 2009-2024 Christian Kohlschütter
    *
    * 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 org.newsclub.net.unix.rmi;
  
  import java.io.IOException;
  import java.rmi.AccessException;
  import java.rmi.AlreadyBoundException;
  import java.rmi.ConnectIOException;
  import java.rmi.NoSuchObjectException;
  import java.rmi.NotBoundException;
  import java.rmi.Remote;
  import java.rmi.RemoteException;
  import java.rmi.registry.Registry;
  import java.rmi.server.RemoteObject;
  import java.rmi.server.RemoteServer;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicBoolean;
  
  import org.newsclub.net.unix.rmi.ShutdownHookSupport.ShutdownHook;
  
  import com.kohlschutter.annotations.compiletime.SuppressFBWarnings;
  
  /**
   * A wrapper for RMI registries, both remote and local, to allow for a clean removal of bound
   * resources upon shutdown.
   *
   * @author Christian Kohlschütter
   */
  public abstract class AFRegistry implements Registry {
    final RemoteCloseable<?> boundCloser;
  
    private final Registry impl;
    private final Map<String, Remote> bound = new HashMap<>();
    private final AFNaming naming;
    private final AtomicBoolean boundCloserExported = new AtomicBoolean(false);
  
    private AFRMIService rmiService = null;
  
    AFRegistry(AFNaming naming, Registry impl) {
      this.naming = naming;
      this.impl = impl;
      this.boundCloser = new RemoteCloseable<Void>() {
        @Override
        public Void get() throws IOException {
          return null;
        }
  
        @Override
        public void close() throws IOException {
          AFRegistry.this.forceUnexportBound();
        }
      };
  
      if (isLocal()) {
        ShutdownHookSupport.addWeakShutdownHook(new ShutdownHook() {
          @Override
          public void onRuntimeShutdown(Thread thread) {
            forceUnexportBound();
          }
        });
      }
    }
  
    /**
     * Returns {@code true} if the wrapped Registry instance is a locally created
     * {@link RemoteServer}.
     *
     * @return {@code true} if wrapped instance is a locally created {@link RemoteServer}.
     * @see #isLocal()
     */
    @Deprecated
    public boolean isRemoteServer() {
      return isLocal();
    }
  
    /**
     * Returns {@code true} if the wrapped Registry instance is locally created.
     *
     * @return {@code true} if wrapped instance is locally created.
     */
    public final boolean isLocal() {
      return (impl instanceof RemoteServer);
   }
 
   /**
    * Returns the {@link AFNaming} instance responsible for this registry.
    *
    * @return The {@link AFNaming} instance.
    */
   @SuppressFBWarnings("EI_EXPOSE_REP")
   public AFNaming getNaming() {
     return naming;
   }
 
   @Override
   public Remote lookup(String name) throws RemoteException, NotBoundException, AccessException {
     return impl.lookup(name);
   }
 
   @Override
   public void bind(String name, Remote obj) throws RemoteException, AlreadyBoundException,
       AccessException {
     impl.bind(name, RemoteObject.toStub(obj));
     synchronized (bound) {
       bound.put(name, obj);
     }
     checkBound();
   }
 
   @Override
   public void unbind(String name) throws RemoteException, NotBoundException, AccessException {
     impl.unbind(name);
     synchronized (bound) {
       bound.remove(name);
     }
     checkBound();
   }
 
   @Override
   public void rebind(String name, Remote obj) throws RemoteException, AccessException {
     impl.rebind(name, RemoteObject.toStub(obj));
     synchronized (bound) {
       bound.put(name, obj);
     }
     checkBound();
   }
 
   @Override
   public String[] list() throws RemoteException, AccessException {
     return impl == null ? new String[0] : impl.list();
   }
 
   /**
    * Returns the remote reference bound to the specified <code>name</code> in this registry. If the
    * reference has not been bound yet, repeated attempts to resolve it are made until the specified
    * time elapses.
    *
    * @param name the name for the remote reference to look up
    * @param timeout The timeout value.
    * @param unit The timeout unit.
    *
    * @return a reference to a remote object
    *
    * @throws NotBoundException if <code>name</code> is not currently bound and couldn't be resolved
    *           in the specified time.
    *
    * @throws RemoteException if remote communication with the registry failed; if exception is a
    *           <code>ServerException</code> containing an <code>AccessException</code>, then the
    *           registry denies the caller access to perform this operation
    *
    * @throws AccessException if this registry is local and it denies the caller access to perform
    *           this operation
    *
    * @throws NullPointerException if <code>name</code> is <code>null</code>
    */
   public Remote lookup(String name, long timeout, TimeUnit unit) throws NotBoundException,
       RemoteException {
     long timeWait = unit.toMillis(timeout);
 
     Exception exFirst = null;
     do {
       try {
         return impl.lookup(name);
       } catch (NotBoundException | ConnectIOException | NoSuchObjectException e) {
         if (exFirst == null) {
           exFirst = e;
         }
       }
 
       try {
         Thread.sleep(Math.min(timeWait, 50));
       } catch (InterruptedException e1) {
         exFirst.addSuppressed(e1);
         break;
       }
       timeWait -= 50;
     } while (timeWait > 0);
 
     if (exFirst instanceof NotBoundException) {
       throw (NotBoundException) exFirst;
     } else if (exFirst instanceof RemoteException) {
       throw (RemoteException) exFirst;
     } else {
       throw new RemoteException("Lookup timed out");
     }
   }
 
   void forceUnexportBound() {
     final Map<String, Remote> map;
     synchronized (bound) {
       map = new HashMap<>(bound);
       bound.clear();
     }
     try {
       checkBound();
     } catch (RemoteException e1) {
       // ignore
     }
     for (Map.Entry<String, Remote> en : map.entrySet()) {
       String name = en.getKey();
       Remote obj = en.getValue();
       if (obj == null) {
         continue;
       }
       AFNaming.unexportObject(obj);
       try {
         unbind(name);
       } catch (RemoteException | NotBoundException e) {
         // ignore
       }
 
     }
     try {
       for (String list : list()) {
         try {
           unbind(list);
         } catch (RemoteException | NotBoundException e) {
           // ignore
         }
       }
     } catch (RemoteException e) {
       // ignore
     }
 
     AFNaming.unexportObject(this.impl);
     AFNaming.unexportObject(this);
   }
 
   private void checkBound() throws RemoteException {
     boolean empty;
     synchronized (bound) {
       empty = bound.isEmpty();
     }
     if (empty) {
       if (boundCloserExported.compareAndSet(true, false)) {
         AFRMIService service;
         try {
           service = getRMIService();
           service.unregisterForShutdown(boundCloser);
         } catch (NoSuchObjectException | NotBoundException e) {
           return;
         } finally {
           AFNaming.unexportObject(boundCloser);
         }
       }
     } else if (boundCloserExported.compareAndSet(false, true)) {
       AFNaming.exportObject(boundCloser, naming.getSocketFactory());
 
       AFRMIService service;
       try {
         service = getRMIService();
       } catch (NotBoundException e) {
         return;
       }
       service.registerForShutdown(boundCloser);
     }
   }
 
   private AFRMIService getRMIService() throws RemoteException, NotBoundException {
     if (rmiService == null) {
       rmiService = naming.getRMIService(this);
     }
     return rmiService;
   }
 }