package com.mailserver.security;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.annotation.Scheduled;

import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.LinkedHashMap;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Brute-force protection for every place that checks a password: SMTP AUTH, IMAP LOGIN/AUTHENTICATE, web login,
 * password change. One instance is shared by all of them, so failures on SMTP and IMAP add up.
 * <ul>
 *   <li><b>By IP:</b> {@code maxPerIp} failures within {@code window} → the IP is blocked for {@code block};
 *       each next block of the same IP is twice as long, up to {@code maxBlock}.</li>
 *   <li><b>By mailbox:</b> {@code maxPerAccount} failures within {@code window} from any IPs (a distributed attack on
 *       one mailbox) → the mailbox is blocked for {@code block} — but not for IPs that logged in to it successfully
 *       before ({@code knownIpTtl}), so an attacker cannot lock the owner's own mail program out.</li>
 * </ul>
 * While blocked the password is not checked at all (no answer to a guess) and the caller gets {@link #retryAfter}.
 * State is in memory: a restart forgets it, which is fine for minutes-long blocks.
 */
public class LoginGuard {

    private static final Logger log = LoggerFactory.getLogger(LoginGuard.class);

    /** Upper bound of tracked keys of each kind, so a flood of random usernames/IPs cannot eat the memory. */
    static final int MAX_KEYS = 100_000;
    private static final int MAX_KNOWN_IPS_PER_ACCOUNT = 20;

    public record Settings(boolean enabled, int maxPerIp, int maxPerAccount, Duration window, Duration block,
                           Duration maxBlock, Duration knownIpTtl, Set<String> trustedIps, Duration failureDelay) {

        public static Settings defaults() {
            return new Settings(true, 10, 20, Duration.ofMinutes(15), Duration.ofMinutes(15), Duration.ofHours(24),
                    Duration.ofDays(30), Set.of(), Duration.ofSeconds(1));
        }
    }

    /** Failures in the window, current block and how many times the key was blocked in a row. */
    private static final class Counter {
        final Deque<Instant> failures = new ArrayDeque<>();
        volatile Instant blockedUntil = Instant.EPOCH;
        int strikes;
        Instant lastSeen = Instant.EPOCH;
    }

    private final Settings settings;
    private final Clock clock;
    private final Map<String, Counter> byIp = new ConcurrentHashMap<>();
    private final Map<String, Counter> byAccount = new ConcurrentHashMap<>();
    /** account → IPs that logged in to it successfully, with the time of the last success. */
    private final Map<String, Map<String, Instant>> knownIps = new ConcurrentHashMap<>();

    public LoginGuard(Settings settings) {
        this(settings, Clock.systemUTC());
    }

    LoginGuard(Settings settings, Clock clock) {
        this.settings = settings;
        this.clock = clock;
    }

    /** Guard that never blocks (tests, closed networks). */
    public static LoginGuard disabled() {
        Settings d = Settings.defaults();
        return new LoginGuard(new Settings(false, d.maxPerIp(), d.maxPerAccount(), d.window(), d.block(), d.maxBlock(),
                d.knownIpTtl(), Set.of(), Duration.ZERO));
    }

    /** Pause before answering "wrong password": slows down guessing inside one connection. */
    public Duration failureDelay() {
        return settings.enabled() ? settings.failureDelay() : Duration.ZERO;
    }

    /**
     * Zero when the attempt may go on and the password may be checked; otherwise how long the caller has to wait.
     */
    public Duration retryAfter(String ip, String account) {
        if (!settings.enabled()) {
            return Duration.ZERO;
        }
        Instant now = clock.instant();
        Duration wait = Duration.ZERO;
        if (!trusted(ip)) {
            wait = remaining(byIp.get(ip), now);
        }
        String acc = normalize(account);
        if (acc != null && !known(acc, ip, now)) {
            Duration accountWait = remaining(byAccount.get(acc), now);
            if (accountWait.compareTo(wait) > 0) {
                wait = accountWait;
            }
        }
        return wait;
    }

    public boolean blocked(String ip, String account) {
        return !retryAfter(ip, account).isZero();
    }

    /** A wrong password (or an unknown mailbox). */
    public void failure(String ip, String account) {
        if (!settings.enabled()) {
            return;
        }
        Instant now = clock.instant();
        if (!trusted(ip)) {
            count(byIp, ip, settings.maxPerIp(), now, "IP " + ip);
        }
        String acc = normalize(account);
        if (acc != null) {
            count(byAccount, acc, settings.maxPerAccount(), now, "mailbox " + acc);
        }
    }

    /** A right password: the IP becomes known for the mailbox, failures of this IP are forgiven. */
    public void success(String ip, String account) {
        if (!settings.enabled()) {
            return;
        }
        Instant now = clock.instant();
        String acc = normalize(account);
        if (acc != null) {
            if (knownIps.size() < MAX_KEYS || knownIps.containsKey(acc)) {
                knownIps.compute(acc, (k, ips) -> {
                    Map<String, Instant> map = ips != null ? ips : new LinkedHashMap<>();
                    synchronized (map) {
                        map.remove(ip);
                        map.put(ip, now);
                        while (map.size() > MAX_KNOWN_IPS_PER_ACCOUNT) {
                            map.remove(map.keySet().iterator().next());
                        }
                    }
                    return map;
                });
            }
        }
        byIp.computeIfPresent(ip, (k, c) -> {
            c.failures.clear();
            return c.blockedUntil.isAfter(now) ? c : null;
        });
    }

    /** Drops what expired; runs every minute. */
    @Scheduled(fixedDelay = 60_000)
    public void sweep() {
        Instant now = clock.instant();
        Instant idle = now.minus(settings.maxBlock()).minus(settings.window());
        byIp.entrySet().removeIf(e -> stale(e.getValue(), now, idle));
        byAccount.entrySet().removeIf(e -> stale(e.getValue(), now, idle));
        Instant knownSince = now.minus(settings.knownIpTtl());
        knownIps.entrySet().removeIf(e -> {
            synchronized (e.getValue()) {
                e.getValue().values().removeIf(t -> t.isBefore(knownSince));
                return e.getValue().isEmpty();
            }
        });
    }

    int trackedIps() {
        return byIp.size();
    }

    int trackedAccounts() {
        return byAccount.size();
    }

    // ---------------------------------------------------------------------------------------------

    private void count(Map<String, Counter> map, String key, int max, Instant now, String what) {
        if (map.size() >= MAX_KEYS && !map.containsKey(key)) {
            sweep();
            if (map.size() >= MAX_KEYS) {
                log.warn("login guard: {} keys tracked, {} not counted", map.size(), what);
                return;
            }
        }
        map.compute(key, (k, existing) -> {
            Counter c = existing != null ? existing : new Counter();
            c.lastSeen = now;
            Instant from = now.minus(settings.window());
            while (!c.failures.isEmpty() && c.failures.peekFirst().isBefore(from)) {
                c.failures.pollFirst();
            }
            if (c.blockedUntil.isAfter(now)) {
                return c; // already blocked: attempts are refused before the password check anyway
            }
            // A long quiet time after the last block forgives the strikes
            if (c.strikes > 0 && c.blockedUntil.plus(settings.maxBlock()).isBefore(now)) {
                c.strikes = 0;
            }
            c.failures.addLast(now);
            if (c.failures.size() >= max) {
                Duration block = settings.block().multipliedBy(1L << Math.min(c.strikes, 20));
                if (block.compareTo(settings.maxBlock()) > 0) {
                    block = settings.maxBlock();
                }
                c.strikes++;
                c.blockedUntil = now.plus(block);
                c.failures.clear();
                log.warn("login guard: {} blocked for {} after {} failed logins (block #{})",
                        what, human(block), max, c.strikes);
            }
            return c;
        });
    }

    private static Duration remaining(Counter c, Instant now) {
        if (c == null) {
            return Duration.ZERO;
        }
        Instant until = c.blockedUntil;
        return until.isAfter(now) ? Duration.between(now, until) : Duration.ZERO;
    }

    private boolean known(String account, String ip, Instant now) {
        Map<String, Instant> ips = knownIps.get(account);
        if (ips == null) {
            return false;
        }
        synchronized (ips) {
            Instant at = ips.get(ip);
            return at != null && at.plus(settings.knownIpTtl()).isAfter(now);
        }
    }

    private boolean trusted(String ip) {
        return settings.trustedIps().contains(ip);
    }

    private static boolean stale(Counter c, Instant now, Instant idle) {
        return !c.blockedUntil.isAfter(now) && c.lastSeen.isBefore(idle);
    }

    private static String normalize(String account) {
        if (account == null) {
            return null;
        }
        String acc = account.trim().toLowerCase(Locale.ROOT);
        return acc.isEmpty() ? null : acc;
    }

    /** "15 min" / "2 h" for messages to the user. */
    public static String human(Duration d) {
        long minutes = Math.max(1, (d.toSeconds() + 59) / 60);
        return minutes < 120 ? minutes + " min" : (minutes + 59) / 60 + " h";
    }
}