package org.dotwebstack.framework.backend.postgres;

import io.r2dbc.pool.ConnectionPool;
import io.r2dbc.postgresql.api.PostgresqlConnection;
import io.r2dbc.spi.Closeable;
import io.r2dbc.spi.Connection;
import io.r2dbc.spi.ConnectionFactory;
import io.r2dbc.spi.ConnectionFactoryMetadata;
import io.r2dbc.spi.Wrapped;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import reactor.core.Disposable;
import reactor.core.publisher.Mono;
import reactor.util.context.ContextView;

@Slf4j
public class ContextAwareConnectionPool
    implements ConnectionFactory, Disposable, Closeable, Wrapped<ConnectionFactory> {

  public static final String CTX_CONNECTION_UUID = "CONNECTION_UUID";

  private final Map<String, Connection> connectionMap = new HashMap<>();

  private final ConnectionPool connectionPool;

  public ContextAwareConnectionPool(ConnectionPool connectionPool) {
    this.connectionPool = connectionPool;
  }

  @Override
  public Publisher<? extends Connection> create() {
    return Mono.deferContextual(ctx -> connectionPool.create()
        .doOnNext(connection -> addToConnectionMap(ctx, connection)));
  }

  private void addToConnectionMap(ContextView ctx, Connection connection) {
    if (ctx.hasKey(CTX_CONNECTION_UUID)) {
      connectionMap.put(ctx.get(CTX_CONNECTION_UUID), connection);
    } else {
      LOG.debug("Context does not contain key: {}", CTX_CONNECTION_UUID);
    }
  }

  public void cancelRequest(String uuid) {
    LOG.debug("Cancel connection for uuid {} (exists: {})", uuid, connectionMap.containsKey(uuid));
    getConnection(uuid).filter(Wrapped.class::isInstance)
        .map(Wrapped.class::cast)
        .map(Wrapped::unwrap)
        .filter(PostgresqlConnection.class::isInstance)
        .map(PostgresqlConnection.class::cast)
        .ifPresent(this::cancelRequest);
  }

  private void cancelRequest(PostgresqlConnection conn) {
    conn.cancelRequest()
        .subscribe();
  }

  public void cleanUp(String uuid) {
    LOG.debug("Clean up connection for uuid {} (exists: {})", uuid, connectionMap.containsKey(uuid));
    connectionMap.remove(uuid);
  }

  private Optional<Connection> getConnection(String uuid) {
    if (connectionMap.containsKey(uuid)) {
      return Optional.of(connectionMap.get(uuid));
    }
    return Optional.empty();
  }

  boolean hasConnections() {
    return !connectionMap.isEmpty();
  }

  @Override
  public ConnectionFactoryMetadata getMetadata() {
    return connectionPool.getMetadata();
  }

  @Override
  public Publisher<Void> close() {
    return connectionPool.close();
  }

  @Override
  public ConnectionFactory unwrap() {
    return connectionPool.unwrap();
  }

  @Override
  public void dispose() {
    connectionPool.dispose();
  }
}

@@ -140,39 +140,8 @@
 on;%0A
-import io.r2dbc.spi.Closeable;%0A
 impo
@@ -434,40 +434,8 @@
 er;%0A
-import reactor.core.Disposable;%0A
 impo
@@ -466,16 +466,16 @@
 r.Mono;%0A
+
 import r
@@ -554,20 +554,16 @@
 tionPool
-%0A   
  impleme
@@ -587,59 +587,8 @@
 tory
-, Disposable, Closeable, Wrapped%3CConnectionFactory%3E
  %7B%0A%0A
@@ -2459,258 +2459,7 @@
 %0A  %7D
-%0A%0A  @Override%0A  public Publisher%3CVoid%3E close() %7B%0A    return connectionPool.close();%0A  %7D%0A%0A  @Override%0A  public ConnectionFactory unwrap() %7B%0A    return connectionPool.unwrap();%0A  %7D%0A%0A  @Override%0A  public void dispose() %7B%0A    connectionPool.dispose();%0A  %7D
 %0A%7D%0A

Revert "Extended ContextAwareConnectionPool with Disposable, Closeable, Wrapped<ConnectionFactory> implementations. (#633)"