package ervu_business_metrics.service;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicInteger;

import ervu_business_metrics.config.KafkaEnabledCondition;
import ervu_business_metrics.model.sso.UserSessionInfo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Conditional;
import org.springframework.stereotype.Component;

/**
 * @author Adel Kalimullin
 */
@Component
@Conditional(KafkaEnabledCondition.class)
public class SessionResponseHolder {
  private final Logger LOGGER = LoggerFactory.getLogger(SessionResponseHolder.class);
  private final Map<String, SessionAccumulator> sessionsMap = new ConcurrentHashMap<>();

  public void initIfAbsent(String requestKey, int totalExpected) {
    sessionsMap.computeIfAbsent(requestKey, key -> new SessionAccumulator(totalExpected));
  }

  public void add(String requestKey, UserSessionInfo sessionInfo) {
    SessionAccumulator acc = sessionsMap.get(requestKey);
    if (acc != null) {
      acc.addSession(sessionInfo);
    }
  }

  public List<UserSessionInfo> awaitSessions(String requestKey, long timeout, TimeUnit unit)
      throws ExecutionException, InterruptedException, TimeoutException {
    SessionAccumulator acc = sessionsMap.get(requestKey);

    long startTime = System.nanoTime();
    LOGGER.info("Starting to await sessions for requestKey = {}", requestKey);

    while (acc == null && (System.nanoTime() - startTime) < unit.toNanos(timeout)) {
      LOGGER.debug("Accumulator for requestKey = {} not found yet, retrying", requestKey);
      TimeUnit.MILLISECONDS.sleep(50);
      acc = sessionsMap.get(requestKey);
    }

    if (acc == null) {
      LOGGER.error("Accumulator for requestKey: {} was not initialized within the timeout period.", requestKey);
      return List.of();
    }

    return acc.getFuture().get(timeout, unit);
  }

  public void remove(String requestKey) {
    sessionsMap.remove(requestKey);
  }

  private static class SessionAccumulator {
    private final List<UserSessionInfo> sessions = new CopyOnWriteArrayList<>();
    private final CompletableFuture<List<UserSessionInfo>> future = new CompletableFuture<>();
    private final AtomicInteger current = new AtomicInteger();
    private final int totalExpected;

    public SessionAccumulator(int totalExpected) {
      this.totalExpected = totalExpected;
      if (totalExpected == 0) {
        future.complete(List.of());
      }
    }

    public void addSession(UserSessionInfo session) {
      sessions.add(session);
      if (current.incrementAndGet() == totalExpected) {
        future.complete(sessions);
      }
    }

    public CompletableFuture<List<UserSessionInfo>> getFuture() {
      return future;
    }
  }
}
