package com.atlassian.plugin.osgi.hostcomponents.impl;
   
   import com.atlassian.plugin.hostcontainer.HostContainer;
   import com.atlassian.plugin.osgi.hostcomponents.ComponentRegistrar;
   import com.atlassian.plugin.osgi.hostcomponents.ContextClassLoaderStrategy;
   import com.atlassian.plugin.osgi.hostcomponents.HostComponentRegistration;
   import com.atlassian.plugin.osgi.hostcomponents.InstanceBuilder;
   import com.atlassian.plugin.osgi.hostcomponents.PropertyBuilder;
   import com.atlassian.plugin.util.ContextClassLoaderSettingInvocationHandler;
  import com.google.common.collect.Maps;
  import org.osgi.framework.Bundle;
  import org.osgi.framework.BundleContext;
  import org.osgi.framework.ServiceFactory;
  import org.osgi.framework.ServiceRegistration;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.lang.reflect.InvocationHandler;
  import java.lang.reflect.Method;
  import java.lang.reflect.Proxy;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.List;
  import java.util.Map;
  import java.util.concurrent.CopyOnWriteArrayList;
  import java.util.function.Function;
  
  /**
   * Default component registrar that also can write registered host components into the OSGi service registry.
   */
  public class DefaultComponentRegistrar implements ComponentRegistrar {
      private static final Logger log = LoggerFactory.getLogger(DefaultComponentRegistrar.class);
  
      private final List<HostComponentRegistration> registry = new CopyOnWriteArrayList<>();
  
      public InstanceBuilder register(final Class<?>... mainInterfaces) {
          final Registration reg = new Registration(mainInterfaces);
          registry.add(reg);
          return new DefaultInstanceBuilder(reg);
      }
  
      public List<ServiceRegistration> writeRegistry(final BundleContext ctx) {
          final ArrayList<ServiceRegistration> services = new ArrayList<>();
  
          for (final HostComponentRegistration reg : new ArrayList<>(registry)) {
              if (Arrays.asList(reg.getMainInterfaceClasses()).contains(HostContainer.class)) {
                  log.warn("Cannot register a HostContainer as a host component, skipping");
                  registry.remove(reg);
                  continue;
              }
  
              final String[] names = reg.getMainInterfaces();
  
              reg.getProperties().put(HOST_COMPONENT_FLAG, Boolean.TRUE.toString());
  
              // If no bean name specified, generate one that will be consistent across restarts
              final String beanName = reg.getProperties().get(PropertyBuilder.BEAN_NAME);
              if (beanName == null) {
                  String genKey = String.valueOf(Arrays.asList(reg.getMainInterfaces()).hashCode());
                  reg.getProperties().put(PropertyBuilder.BEAN_NAME, "hostComponent-" + genKey);
              }
  
              if (log.isDebugEnabled()) {
                  log.debug("Registering: " + Arrays.asList(names) + " instance " + reg.getInstance() + "with properties: " + reg.getProperties());
              }
  
              if (names.length == 0) {
                  log.warn("Host component " + beanName + " of instance " + reg.getInstance() + " has no interfaces");
              }
  
              Object service = reg.getInstance();
  
              if (!ContextClassLoaderStrategy.USE_PLUGIN.name().equals(reg.getProperties().get(PropertyBuilder.CONTEXT_CLASS_LOADER_STRATEGY))) {
                  service = createContextClassLoaderSettingProxy(reg.getMainInterfaceClasses(), service);
              }
  
              if (Boolean.parseBoolean(reg.getProperties().get(PropertyBuilder.TRACK_BUNDLE))) {
                  service = createTrackBundleProxy(reg.getMainInterfaceClasses(), service);
              }
  
              final ServiceRegistration sreg = ctx.registerService(names, service, reg.getProperties());
              if (sreg != null) {
                  services.add(sreg);
              }
          }
          return Collections.unmodifiableList(services);
      }
  
      public List<HostComponentRegistration> getRegistry() {
          return Collections.unmodifiableList(registry);
      }
  
      /**
       * Wraps the service in a dynamic proxy that ensures all methods are executed with the object class's class loader
       * as the context class loader
       *
       * @param interfaces The interfaces to proxy
       * @param service    The instance to proxy
      * @return A proxy that wraps the service
      */
     private Object createContextClassLoaderSettingProxy(final Class<?>[] interfaces, final Object service) {
         final Function<Object, Object> transformer = service1 ->
                 Proxy.newProxyInstance(DefaultComponentRegistrar.class.getClassLoader(), interfaces, new ContextClassLoaderSettingInvocationHandler(service1));
 
         if (!(service instanceof ServiceFactory)) {
             return transformer.apply(service);
         }
         return new TransformingServiceFactory((ServiceFactory) service) {
             @Override
             protected Object transform(Bundle bundle, ServiceRegistration registration, Object service) {
                 return transformer.apply(service);
             }
         };
     }
 
     private ServiceFactory createTrackBundleProxy(final Class<?>[] interfaces, final Object service) {
         final ServiceFactory delegate = service instanceof ServiceFactory ?
                 (ServiceFactory) service :
                 new InstanceServiceFactory(service);
 
         return new TransformingServiceFactory(delegate) {
             @Override
             protected Object transform(final Bundle bundle, final ServiceRegistration registration, final Object service) {
                 return Proxy.newProxyInstance(DefaultComponentRegistrar.class.getClassLoader(), interfaces,
                         new BundleTrackingInvocationHandler(bundle, service));
             }
         };
     }
 
     private static class BundleTrackingInvocationHandler implements InvocationHandler {
 
         private final Bundle bundle;
         private final Object service;
 
         private BundleTrackingInvocationHandler(Bundle bundle, Object service) {
             this.bundle = bundle;
             this.service = service;
         }
 
         @Override
         public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
             Bundle original = CallingBundleStore.get();
             try {
                 CallingBundleStore.set(bundle);
                 return method.invoke(service, args);
             } finally {
                 CallingBundleStore.set(original);
             }
         }
     }
 
     private static class InstanceServiceFactory implements ServiceFactory {
 
         private final Object service;
 
         private InstanceServiceFactory(final Object service) {
             this.service = service;
         }
 
         @Override
         public Object getService(final Bundle bundle, final ServiceRegistration registration) {
             return service;
         }
 
         @Override
         public void ungetService(final Bundle bundle, final ServiceRegistration registration, final Object service) {
             // no-op
         }
     }
 
     private static abstract class TransformingServiceFactory implements ServiceFactory {
 
         private final ServiceFactory delegate;
         private final Map<Long, Object> bundleIdToOriginalService;
 
         private TransformingServiceFactory(ServiceFactory delegate) {
             this.delegate = delegate;
             this.bundleIdToOriginalService = Maps.newConcurrentMap();
         }
 
 
         @Override
         public final Object getService(final Bundle bundle, final ServiceRegistration registration) {
             final Object service = delegate.getService(bundle, registration);
             final Object transformed = transform(bundle, registration, service);
             // Store the original service after it has been transformed
             // in case the transformer throws an exception
             bundleIdToOriginalService.put(bundle.getBundleId(), service);
             return transformed;
         }
 
         @Override
         public final void ungetService(final Bundle bundle, final ServiceRegistration registration, final Object transformed) {
             final Object service = bundleIdToOriginalService.remove(bundle.getBundleId());
             if (service != null) {
                 delegate.ungetService(bundle, registration, service);
             }
         }
 
         protected abstract Object transform(Bundle bundle, ServiceRegistration registration, Object service);
 
     }
 
 }