/*
    * 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.Closeable;
  import java.io.IOException;
  import java.net.MalformedURLException;
  import java.rmi.AccessException;
  import java.rmi.AlreadyBoundException;
  import java.rmi.ConnectIOException;
  import java.rmi.Naming;
  import java.rmi.NoSuchObjectException;
  import java.rmi.NotBoundException;
  import java.rmi.Remote;
  import java.rmi.RemoteException;
  import java.rmi.ServerException;
  import java.rmi.registry.LocateRegistry;
  import java.rmi.registry.Registry;
  import java.rmi.server.RMISocketFactory;
  import java.rmi.server.UnicastRemoteObject;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Objects;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicBoolean;
  
  import org.eclipse.jdt.annotation.NonNull;
  import org.newsclub.net.unix.AFSocket;
  import org.newsclub.net.unix.rmi.ShutdownHookSupport.ShutdownHook;
  
  import com.kohlschutter.annotations.compiletime.SuppressFBWarnings;
  
  /**
   * The {@link AFSocket}-compatible equivalent of {@link Naming}. Use this class for accessing RMI
   * registries that are reachable by {@link AFSocket}s.
   *
   * @author Christian Kohlschütter
   */
  public abstract class AFNaming extends AFRegistryAccess {
    private static final String RMI_SERVICE_NAME = AFRMIService.class.getName();
  
    private static final Map<AFNamingRef, AFNaming> INSTANCES = new HashMap<>();
  
    private AFRegistry registry = null;
    private AFRMIService rmiService = null;
    private final int registryPort;
    private final int servicePort;
    AFRMISocketFactory socketFactory;
    private final AtomicBoolean remoteShutdownAllowed = new AtomicBoolean(true);
    private final AtomicBoolean shutdownInProgress = new AtomicBoolean(false);
    private final AtomicBoolean addedShutdownHook = new AtomicBoolean(false);
  
    /**
     * Creates a new naming instance with the given ports.
     *
     * @param registryPort The registry port.
     * @param servicePort The port for AFRMIService.
     */
    protected AFNaming(final int registryPort, final int servicePort) {
      super();
      this.registryPort = registryPort;
      this.servicePort = servicePort;
    }
  
    /**
     * Creates a new {@link AFRegistry} given a {@link Registry} implementation.
     *
     * @param impl The implementation.
     * @return The new {@link AFRegistry} instance.
     * @throws RemoteException on error.
     */
    protected abstract AFRegistry newAFRegistry(Registry impl) throws RemoteException;
  
    /**
     * Creates or returns the {@link AFRMISocketFactory} to be used with this instance.
     *
     * @return The socket factory.
     * @throws IOException on error.
     */
    protected abstract AFRMISocketFactory initSocketFactory() throws IOException;
  
    @SuppressWarnings("unchecked")
    static <T extends AFNaming> T getInstance(final int registryPort,
        @NonNull AFNamingProvider<T> provider) throws RemoteException {
     Objects.requireNonNull(provider);
     final AFNamingRef sap = new AFNamingRef(provider, registryPort);
     T instance;
     synchronized (AFNaming.class) {
       instance = (T) INSTANCES.get(sap);
       if (instance == null) {
         try {
           instance = provider.newInstance(registryPort);
           Objects.requireNonNull(instance);
           synchronized (instance) {
             instance.socketFactory = instance.initSocketFactory();
           }
         } catch (RemoteException e) {
           throw e;
         } catch (IOException e) {
           throw new RemoteException(e.getMessage(), e);
         }
         INSTANCES.put(sap, instance);
       }
     }
     return instance;
   }
 
   /**
    * Returns the {@link AFRMISocketFactory} associated with this instance.
    *
    * @return The {@link AFRMISocketFactory}.
    */
   @SuppressFBWarnings("EI_EXPOSE_REP")
   public synchronized AFRMISocketFactory getSocketFactory() {
     return socketFactory;
   }
 
   /**
    * Returns the registry port.
    *
    * @return The port.
    */
   public final int getRegistryPort() {
     return registryPort;
   }
 
   AFRMIService getRMIService() throws RemoteException, NotBoundException {
     return getRMIService(getRegistry());
   }
 
   synchronized AFRMIService getRMIService(AFRegistry reg) throws RemoteException,
       NotBoundException {
     if (rmiService == null) {
       this.rmiService = getRMIServiceFromRegistry(reg);
     }
     return rmiService;
   }
 
   AFRMIService getRMIServiceFromRegistry(AFRegistry reg) throws RemoteException, NotBoundException {
     AFRMIService service;
     service = (AFRMIService) reg.lookup(RMI_SERVICE_NAME, 5, TimeUnit.SECONDS);
     this.remoteShutdownAllowed.set(service.isShutdownAllowed());
     return service;
   }
 
   private void closeUponRuntimeShutdown() {
     if (addedShutdownHook.compareAndSet(false, true)) {
       ShutdownHookSupport.addWeakShutdownHook(new ShutdownHook() {
 
         @Override
         @SuppressWarnings("LockOnNonEnclosingClassLiteral" /* errorprone */)
         public synchronized void onRuntimeShutdown(Thread thread) throws IOException {
           if (registry != null && registry.isLocal()) {
             shutdownRegistry();
           }
         }
       });
     }
   }
 
   private synchronized void rebindRMIService(final AFRMIService assigner) throws RemoteException {
     rmiService = assigner;
     getRegistry().rebind(RMI_SERVICE_NAME, assigner);
   }
 
   @Override
   public AFRegistry getRegistry() throws RemoteException {
     return getRegistry(0, TimeUnit.SECONDS);
   }
 
   /**
    * Returns a reference to the existing RMI registry.
    *
    * If there's no registry running at this port after waiting for up to the given time, an
    * exception is thrown.
    *
    * @param timeout The timeout value.
    * @param unit The timeout unit.
    * @return The registry.
    * @throws RemoteException If there was a problem.
    */
   public AFRegistry getRegistry(long timeout, TimeUnit unit) throws RemoteException {
     if (shutdownInProgress.get()) {
       throw new ShutdownException();
     }
     synchronized (this) {
       AFRegistry reg = getRegistry(false);
       if (reg == null) {
         reg = openRegistry(timeout, unit);
       }
       return reg;
     }
   }
 
   /**
    * Tries to access the registry, waiting some time if necessary.
    *
    * @param timeout The timeout.
    * @param unit The unit for the timeout.
    * @return The registry instance.
    * @throws RemoteException on error.
    */
   protected abstract AFRegistry openRegistry(long timeout, TimeUnit unit) throws RemoteException;
 
   /**
    * Returns a reference to the RMI registry, or {@code null}.
    *
    * If there's no registry running at this port, and {@code create} is set to {@code true}, a new
    * one is created; when {@code create} is set to {@code false}, {@code null} is returned.
    *
    * @param create {@code true} if a new register may be created if necessary.
    * @return The registry, or {@code null}
    * @throws RemoteException If there was a problem.
    */
   public AFRegistry getRegistry(boolean create) throws RemoteException {
     if (shutdownInProgress.get()) {
       throw new ShutdownException();
     }
     synchronized (this) {
       if (registry != null) {
         return registry;
       } else if (!socketFactory.hasRegisteredPort(registryPort)) {
         return create ? createRegistry() : null;
       }
 
       AFRegistry reg = locateRegistry();
       setRegistry(reg);
 
       try {
         getRMIService(reg);
       } catch (NotBoundException | NoSuchObjectException | ConnectIOException e) {
         if (create) {
           setRegistry(null);
           return createRegistry();
         } else {
           throw new ServerException("Could not access " + AFRMIService.class.getName(), e);
         }
       }
 
       return registry;
     }
   }
 
   private AFRegistry locateRegistry() throws RemoteException {
     Registry regImpl = LocateRegistry.getRegistry(null, registryPort, socketFactory);
     return regImpl == null ? null : newAFRegistry(regImpl);
   }
 
   /**
    * Shuts this RMI Registry down.
    *
    * @throws RemoteException if the operation fails.
    */
   public void shutdownRegistry() throws RemoteException {
     synchronized (this) {
       if (registry == null) {
         return;
       }
 
       AFRegistry registryToBeClosed = registry;
       AFRMIService rmiServiceToBeClosed = rmiService;
 
       if (!registryToBeClosed.isLocal()) {
         if (!isRemoteShutdownAllowed()) {
           throw new ServerException("The server refuses to be shutdown remotely");
         }
         setRegistry(null);
 
         try {
           shutdownViaRMIService(registryToBeClosed, rmiServiceToBeClosed);
         } catch (Exception e) {
           // ignore
         }
         return;
       }
 
       setRegistry(null);
 
       if (!shutdownInProgress.compareAndSet(false, true)) {
         return;
       }
       try {
         unexportRMIService(registryToBeClosed, (AFRMIServiceImpl) rmiServiceToBeClosed);
         forceUnexportBound(registryToBeClosed);
         closeSocketFactory();
         shutdownRegistryFinishingTouches();
       } finally {
         shutdownInProgress.set(false);
       }
     }
   }
 
   /**
    * Called by {@link #shutdownRegistry()} as the final step.
    */
   protected abstract void shutdownRegistryFinishingTouches();
 
   private synchronized void unexportRMIService(AFRegistry reg, AFRMIServiceImpl serv)
       throws AccessException, RemoteException {
     if (serv != null) {
       serv.shutdownRegisteredCloseables();
     }
 
     try {
       if (serv != null) {
         unexportObject(serv);
       }
       reg.unbind(RMI_SERVICE_NAME);
     } catch (ShutdownException | NotBoundException e) {
       // ignore
     }
     this.rmiService = null;
   }
 
   private void forceUnexportBound(AFRegistry reg) {
     try {
       reg.forceUnexportBound();
     } catch (Exception e) {
       // ignore
     }
   }
 
   private void closeSocketFactory() {
     if (socketFactory != null) {
       try {
         socketFactory.close();
       } catch (IOException e) {
         // ignore
       }
     }
   }
 
   private void shutdownViaRMIService(AFRegistry reg, AFRMIService serv) throws RemoteException {
     try {
       if (serv == null) {
         serv = getRMIService(reg);
       }
       if (serv.isShutdownAllowed()) {
         serv.shutdown();
       }
     } catch (ServerException | ConnectIOException | NotBoundException e) {
       // ignore
     }
   }
 
   /**
    * Creates a new RMI {@link Registry}.
    *
    * If there already was a registry created previously, it is shut down and replaced by the current
    * one.
    *
    * Use {@link #getRegistry()} to try to reuse an existing registry.
    *
    * @return The registry
    * @throws RemoteException if the operation fails.
    * @see #getRegistry()
    */
   public synchronized AFRegistry createRegistry() throws RemoteException {
     AFRegistry existingRegistry = registry;
     if (existingRegistry == null) {
       try {
         existingRegistry = getRegistry(false);
       } catch (ServerException e) {
         Throwable cause = e.getCause();
         if (cause instanceof NotBoundException || cause instanceof ConnectIOException) {
           existingRegistry = null;
         } else {
           throw e;
         }
       }
     }
     if (existingRegistry != null) {
       if (!isRemoteShutdownAllowed()) {
         throw new ServerException("The server refuses to be shutdown remotely");
       }
       shutdownRegistry();
     }
 
     initRegistryPrerequisites();
     AFRegistry newAFRegistry = newAFRegistry(LocateRegistry.createRegistry(registryPort,
         socketFactory, socketFactory));
     setRegistry(newAFRegistry);
 
     final AFRMIService service = new AFRMIServiceImpl(this);
     UnicastRemoteObject.exportObject(service, servicePort, socketFactory, socketFactory);
 
     rebindRMIService(service);
 
     return registry;
   }
 
   /**
    * Called by {@link #createRegistry()} right before creating/setting the registry.
    *
    * @throws ServerException on error.
    */
   protected abstract void initRegistryPrerequisites() throws ServerException;
 
   /**
    * Checks if this {@link AFNaming} instance can be shut down remotely.
    *
    * @return {@code true} if remote shutdown is allowed.
    */
   public boolean isRemoteShutdownAllowed() {
     return remoteShutdownAllowed.get();
   }
 
   /**
    * Controls whether this {@link AFNaming} instance can be shut down remotely.
    *
    * @param remoteShutdownAllowed {@code true} if remote shutdown is allowed.
    */
   public void setRemoteShutdownAllowed(boolean remoteShutdownAllowed) {
     this.remoteShutdownAllowed.set(remoteShutdownAllowed);
   }
 
   /**
    * Exports and binds the given Remote object to the given name, using the given {@link AFNaming}
    * setup.
    *
    * @param name The name to use to bind the object in the registry.
    * @param obj The object to export and bind.
    * @throws RemoteException if the operation fails.
    * @throws AlreadyBoundException if there already was something bound at that name
    */
   public void exportAndBind(String name, Remote obj) throws RemoteException, AlreadyBoundException {
     exportObject(obj, getSocketFactory());
 
     getRegistry().bind(name, obj);
   }
 
   /**
    * Exports and re-binds the given Remote object to the given name, using the given
    * {@link AFNaming} setup.
    *
    * @param name The name to use to bind the object in the registry.
    * @param obj The object to export and bind.
    * @throws RemoteException if the operation fails.
    */
   public void exportAndRebind(String name, Remote obj) throws RemoteException {
     exportObject(obj, getSocketFactory());
 
     getRegistry().rebind(name, obj);
   }
 
   /**
    * Forcibly un-exports the given object, if it exists, and unbinds the object from the registry
    * (otherwise returns without an error).
    *
    * @param name The name used to bind the object.
    * @param obj The object to un-export.
    * @throws RemoteException if the operation fails.
    */
   public void unexportAndUnbind(String name, Remote obj) throws RemoteException {
     unexportObject(obj);
     try {
       unbind(name);
     } catch (MalformedURLException | NotBoundException e) {
       // ignore
     }
   }
 
   /**
    * Exports the given Remote object, using the given socket factory and a randomly assigned port.
    *
    * NOTE: This helper function can also be used for regular RMI servers.
    *
    * @param obj The object to export.
    * @param socketFactory The socket factory to use.
    * @return The remote stub.
    * @throws RemoteException if the operation fails.
    */
   public static Remote exportObject(Remote obj, RMISocketFactory socketFactory)
       throws RemoteException {
     return UnicastRemoteObject.exportObject(obj, 0, socketFactory, socketFactory);
   }
 
   /**
    * Forcibly un-exports the given object, if it exists (otherwise returns without an error). This
    * should be called upon closing a {@link Closeable} {@link Remote} object.
    *
    * NOTE: This helper function can also be used for regular RMI servers.
    *
    * @param obj The object to un-export.
    */
   public static void unexportObject(Remote obj) {
     try {
       UnicastRemoteObject.unexportObject(obj, true);
     } catch (NoSuchObjectException e) {
       // ignore
     }
   }
 
   private synchronized void setRegistry(AFRegistry registry) {
     this.registry = registry;
     if (registry == null) {
       rmiService = null;
     } else if (registry.isLocal()) {
       closeUponRuntimeShutdown();
     }
   }
 }