package io.scalecube.services.registry;

import io.scalecube.services.ServiceEndpoint;
import io.scalecube.services.ServiceReference;
import io.scalecube.services.registry.api.RegistryEvent;
import io.scalecube.services.registry.api.ServiceRegistry;
import io.scalecube.services.registry.api.RegistryEvent.Type;

import org.jctools.maps.NonBlockingHashMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import reactor.core.publisher.DirectProcessor;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxProcessor;
import reactor.core.publisher.FluxSink;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class ServiceRegistryImpl implements ServiceRegistry {

  private static final Logger LOGGER = LoggerFactory.getLogger(ServiceRegistryImpl.class);

  // todo how to remove it (tags problem)?
  private final Map<String, ServiceEndpoint> serviceEndpoints = new NonBlockingHashMap<>();
  private final Map<String, List<ServiceReference>> referencesByQualifier = new NonBlockingHashMap<>();

  private final FluxProcessor<RegistryEvent, RegistryEvent> events =
      DirectProcessor.<RegistryEvent>create().serialize();

  private final FluxSink<RegistryEvent> sink = events.sink();

  @Override
  public List<ServiceEndpoint> listServiceEndpoints() {
    // todo how to collect tags correctly?
    return new ArrayList<>(serviceEndpoints.values());
  }

  @Override
  public List<ServiceReference> listServiceReferences() {
    return serviceReferenceStream().collect(Collectors.toList());
  }

  @Override
  public List<ServiceReference> lookupService(String qualifier) {
    List<ServiceReference> result = referencesByQualifier.get(qualifier);
    return result != null ? Collections.unmodifiableList(result) : Collections.emptyList();
  }

  @Override
  public List<ServiceReference> lookupService(Predicate<? super ServiceReference> filter) {
    // Convert to stream of service references
    Stream<ServiceReference> stream = serviceReferenceStream();

    // Filter by filter
    if (filter != null) {
      stream = stream.filter(filter);
    }

    // Collect results
    return stream.collect(Collectors.toList());
  }

  @Override
  public boolean registerService(ServiceEndpoint serviceEndpoint) {
    boolean success = serviceEndpoints.putIfAbsent(serviceEndpoint.id(), serviceEndpoint) == null;
    if (success) {
      serviceEndpoint
          .serviceRegistrations()
          .stream()
          .flatMap(serviceRegistration -> serviceRegistration.methods().stream()
              .map(sm -> new ServiceReference(sm, serviceRegistration, serviceEndpoint)))
          .forEach(serviceReference -> {
            referencesByQualifier
                .computeIfAbsent(serviceReference.qualifier(), key -> new CopyOnWriteArrayList<>())
                .add(serviceReference);
            sink.next(new RegistryEvent(Type.ADDED, serviceReference));
          });
    }
    return success;
  }

  @Override
  public ServiceEndpoint unregisterService(String endpointId) {
    ServiceEndpoint serviceEndpoint = serviceEndpoints.remove(endpointId);
    if (serviceEndpoint != null) {
      referencesByQualifier.values().forEach(list -> list.removeIf(sr -> sr.endpointId().equals(endpointId)));
    }
    return serviceEndpoint;
  }

  Stream<ServiceReference> serviceReferenceStream() {
    return referencesByQualifier.values().stream().flatMap(Collection::stream);
  }

  public Flux<RegistryEvent> listen() {
    return events;
  }
}

@@ -226,31 +226,24 @@
 try.api.
-Service
 Registry
 ;%0Aimport
@@ -234,16 +234,26 @@
 Registry
+Event.Type
 ;%0Aimport
@@ -284,32 +284,39 @@
 egistry.api.
+Service
 Registry
 Event.Type;%0A
@@ -303,26 +303,16 @@
 Registry
-Event.Type
 ;%0A%0Aimpor
@@ -3399,16 +3399,27 @@
 values()
+%0A          
 .forEach
@@ -3431,21 +3431,42 @@
  -%3E 
-list.removeIf
+%7B%0A            list.stream().filter
 (sr 
@@ -3503,16 +3503,210 @@
 ointId))
+.collect(Collectors.toSet())%0A                .forEach(sr -%3E %7B%0A                  list.remove(sr);%0A                  sink.next(new RegistryEvent(Type.REMOVED, sr));%0A                %7D);%0A          %7D
 );%0A    %7D
@@ -3929,22 +3929,175 @@
  return 
+Flux.fromIterable(referencesByQualifier.values()).flatMap(Flux::fromIterable)%0A        .map(sr -%3E new RegistryEvent(Type.ADDED, sr))%0A        .concatWith(
 events
+)
 ;%0A  %7D%0A%7D%0A

update remove events sink