package su.xserver.iikocon.handler;

import io.vertx.core.Handler;
import io.vertx.core.json.JsonObject;
import io.vertx.ext.web.RoutingContext;
import io.vertx.redis.client.Command;
import io.vertx.redis.client.Redis;
import io.vertx.redis.client.Request;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.NavigableMap;
import java.util.Set;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;

public class RedisRateLimiter implements Handler<RoutingContext> {

  private final Logger logger;
  private final Redis redis;
  private final int limitPerWindow;
  private final long windowMillis;
  private static final String PREFIX = "ip:limit:";

  // Основной кэш: clientKey -> время окончания блокировки
  private final ConcurrentHashMap<String, Long> blockedClients = new ConcurrentHashMap<>();
  // Индекс по времени: время окончания -> множество клиентов
  private final ConcurrentSkipListMap<Long, Set<String>> expiryIndex = new ConcurrentSkipListMap<>();
  private final ScheduledExecutorService cleaner = Executors.newSingleThreadScheduledExecutor();

  private final AtomicLong allowedRequests = new AtomicLong(0);
  private final AtomicLong blockedRequests = new AtomicLong(0);
  private final AtomicLong redisCalls = new AtomicLong(0);
  private final AtomicLong redisFailures = new AtomicLong(0);
  private final AtomicLong totalRedisLatency = new AtomicLong(0);
  private final AtomicLong redisLatencyCount = new AtomicLong(0);

  // Частота блокировок по IP
  private final ConcurrentHashMap<String, AtomicLong> blockedByClient = new ConcurrentHashMap<>();

  public RedisRateLimiter(Redis redis, int limitPerWindow, long windowMillis) {
    this.logger = LoggerFactory.getLogger("[RedisRateLimiter]");
    this.redis = redis;
    this.limitPerWindow = limitPerWindow;
    this.windowMillis = windowMillis;

    // Периодическая очистка только истёкших блокировок
    cleaner.scheduleAtFixedRate(this::cleanupExpiredClients, windowMillis, windowMillis / 2, TimeUnit.MILLISECONDS);
  }

  @Override
  public void handle(RoutingContext context) {
    String clientKey = getClientKey(context);
    long now = System.currentTimeMillis();

    // Проверяем локальную блокировку
    Long blockedUntil = blockedClients.get(clientKey);
    if (blockedUntil != null) {
      if (blockedUntil > now) {
        blockedRequests.incrementAndGet();
        incrementBlockCount(clientKey);
        sendTooManyRequests(context);
        return;
      } else {
        unblockClient(clientKey, blockedUntil);
      }
    }

    String redisKey = PREFIX + clientKey;
    checkRateLimit(context, redisKey, clientKey);
  }

  private void checkRateLimit(RoutingContext context, String redisKey, String clientKey) {
    String luaScript = """
        local key = KEYS[1]
        local limit = tonumber(ARGV[1])
        local ttl = tonumber(ARGV[2])

        local current = redis.call('INCR', key)
        if current == 1 then
            redis.call('PEXPIRE', key, ttl)
        end

        if current > limit then
            return 'TOO_MANY_REQUESTS'
        else
            return 'OK'
        end
    """;

    redisCalls.incrementAndGet();
    long start = System.nanoTime();

    Request request = Request.cmd(Command.EVAL)
      .arg(luaScript)
      .arg(1)
      .arg(redisKey)
      .arg(limitPerWindow)
      .arg(windowMillis);

    redis.send(request)
      .onSuccess(response -> {
        long duration = System.nanoTime() - start;
        redisLatencyCount.incrementAndGet();
        totalRedisLatency.addAndGet(TimeUnit.NANOSECONDS.toMillis(duration));

        String result = response.toString();
        if ("TOO_MANY_REQUESTS".equals(result)) {
          blockClient(clientKey);
          blockedRequests.incrementAndGet();
          incrementBlockCount(clientKey);
          sendTooManyRequests(context);
        } else {
          allowedRequests.incrementAndGet();
          context.next();
        }
      }).onFailure(error -> {
        redisFailures.incrementAndGet();
        context.response()
          .setStatusCode(503)
          .putHeader("Content-Type", "application/json")
          .end(new JsonObject()
            .put("error", "503 Service Unavailable")
            .put("message", "Redis is not connected")
            .encodePrettily()
          );

        logger.error(error.getMessage());
      });
  }

  private void blockClient(String clientKey) {
    long blockedUntil = System.currentTimeMillis() + windowMillis;
    blockedClients.put(clientKey, blockedUntil);
    expiryIndex.computeIfAbsent(blockedUntil, t -> ConcurrentHashMap.newKeySet()).add(clientKey);
  }

  private void unblockClient(String clientKey, long expiryTime) {
    blockedClients.remove(clientKey);
    Set<String> clients = expiryIndex.get(expiryTime);
    if (clients != null) {
      clients.remove(clientKey);
      if (clients.isEmpty()) {
        expiryIndex.remove(expiryTime);
      }
    }
  }

  private void incrementBlockCount(String clientKey) {
    blockedByClient.computeIfAbsent(clientKey, k -> new AtomicLong(0)).incrementAndGet();
  }

  private void cleanupExpiredClients() {
    long now = System.currentTimeMillis();
    // Получаем все записи, у которых время истечения <= now
    NavigableMap<Long, Set<String>> expired = expiryIndex.headMap(now, true);

    if (expired.isEmpty()) return;

    for (Map.Entry<Long, Set<String>> entry : expired.entrySet()) {
      Set<String> clients = entry.getValue();
      for (String client : clients) {
        blockedClients.remove(client);
      }
    }

    expired.clear(); // очищаем диапазон из индекса
  }

  private void sendTooManyRequests(RoutingContext context) {
    context.response()
      .setStatusCode(429)
      .putHeader("Content-Type", "application/json")
      .end(new JsonObject()
        .put("error", "429 Too Many Requests")
        .put("message", "Try again later")
        .encodePrettily()
      );
  }

  private String getClientKey(RoutingContext context) {
    return context.request().remoteAddress().host().replace(':', '.');
  }
}
