package com.mailserver.smtp;

import com.mailserver.security.LoginGuard;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufUtil;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelPipeline;
import io.netty.handler.codec.LineBasedFrameDecoder;
import io.netty.handler.codec.MessageToMessageDecoder;
import io.netty.handler.ssl.SslContext;
import io.netty.handler.ssl.SslHandler;
import io.netty.handler.timeout.ReadTimeoutHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;
import reactor.netty.Connection;
import reactor.netty.DisposableServer;
import reactor.netty.NettyOutbound;
import reactor.netty.tcp.TcpServer;

import java.io.ByteArrayOutputStream;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.Locale;
import java.util.concurrent.TimeUnit;

/**
 * SMTP server (RFC 5321) on Reactor Netty.
 * <ul>
 *   <li>input is split into lines by {@link LineBasedFrameDecoder} — a network packet may hold half a line or many
 *       pipelined commands, both work;</li>
 *   <li>AUTH PLAIN / LOGIN against our mailboxes; someone else's address is accepted only from a logged-in sender
 *       (no open relay); a logged-in sender may use only his own address in MAIL FROM;</li>
 *   <li>the message is stored or queued before the client gets 250 — on failure 451 and the client retries;</li>
 *   <li>database and DNS work runs on a worker thread, not on the event loop.</li>
 * </ul>
 * TLS (RFC 3207 STARTTLS, or implicit TLS on a separate port, RFC 8314): with a {@code tls} context STARTTLS is
 * offered; AUTH only over TLS, or without it when {@code authWithoutTls} is on (a closed network: VM, docker network).
 * After STARTTLS the session starts over (new EHLO), and lines sent in plain text together with STARTTLS are not
 * executed — the connection is closed (command injection, CVE-2011-0411).
 */
public class SmtpServer {

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

    // RFC 5321 4.5.3.1.6: a text line is at most 1000 octets with CRLF; we allow more for sloppy clients
    static final int MAX_LINE = 64 * 1024;
    static final int MAX_RECIPIENTS = 100;
    private static final int MAX_BAD_COMMANDS = 20;
    /** Wrong passwords in one connection before it is closed. */
    private static final int MAX_AUTH_FAILURES = 3;
    // Compared by identity: TLS starts right after exactly this reply is written
    private static final String STARTTLS_READY = new String("220 2.0.0 Ready to start TLS\r\n");

    /**
     * @param tls         certificate for STARTTLS / implicit TLS; null — no TLS
     * @param implicitTls TLS from the first byte (port 465), no STARTTLS
     */
    public record Settings(String host, int port, String hostname, long maxMessageBytes, boolean authWithoutTls,
                           Duration idleTimeout, SslContext tls, boolean implicitTls) {

        public Settings(String host, int port, String hostname, long maxMessageBytes, boolean authWithoutTls,
                        Duration idleTimeout) {
            this(host, port, hostname, maxMessageBytes, authWithoutTls, idleTimeout, null, false);
        }
    }

    /** A line of input and whether it came over TLS. */
    record Line(byte[] bytes, boolean tls) {
    }

    /** Marks each line: decoded behind an SslHandler (encrypted) or not. */
    private static final class LineTagger extends MessageToMessageDecoder<ByteBuf> {
        @Override
        protected void decode(ChannelHandlerContext ctx, ByteBuf frame, List<Object> out) {
            out.add(new Line(ByteBufUtil.getBytes(frame), ctx.pipeline().get(SslHandler.class) != null));
        }
    }

    private final Settings settings;
    private final SmtpBackend backend;
    private final LoginGuard guard;
    private DisposableServer server;

    public SmtpServer(Settings settings, SmtpBackend backend) {
        this(settings, backend, LoginGuard.disabled());
    }

    public SmtpServer(Settings settings, SmtpBackend backend, LoginGuard guard) {
        this.settings = settings;
        this.backend = backend;
        this.guard = guard;
    }

    public void start() {
        if (settings.implicitTls() && settings.tls() == null) {
            throw new IllegalStateException("Implicit TLS needs a certificate");
        }
        TcpServer tcp = TcpServer.create()
                .host(settings.host())
                .port(settings.port())
                .doOnConnection(connection -> connection
                        .addHandlerFirst("smtp-idle", new ReadTimeoutHandler(settings.idleTimeout().toSeconds(), TimeUnit.SECONDS))
                        .addHandlerLast("smtp-lines", new LineBasedFrameDecoder(MAX_LINE, true, false))
                        .addHandlerLast("smtp-tag", new LineTagger()));
        if (settings.implicitTls()) {
            tcp = tcp.secure(spec -> spec.sslContext(settings.tls()));
        }
        server = tcp
                .handle((inbound, outbound) -> {
                    Connection[] connection = {null};
                    String[] remoteIp = {"unknown"};
                    inbound.withConnection(c -> {
                        connection[0] = c;
                        if (c.address() instanceof InetSocketAddress address && address.getAddress() != null) {
                            remoteIp[0] = address.getAddress().getHostAddress();
                        }
                    });
                    Session session = new Session(remoteIp[0], settings.implicitTls());
                    Flux<String> replies = inbound.receiveObject().cast(Line.class)
                            .concatMap(session::onLine)
                            .onErrorResume(e -> {
                                log.debug("SMTP {}: connection closed: {}", session.remoteIp, e.toString());
                                return Mono.just("421 4.4.2 " + settings.hostname() + " closing connection\r\n");
                            })
                            .takeUntil(reply -> session.closing);
                    // One reply at a time, each written before the next line is handled: after "220 Ready to start TLS"
                    // the TLS handler goes in only when the reply has left in plain text
                    return Flux.concat(Mono.just("220 " + settings.hostname() + " ESMTP ready\r\n"), replies)
                            .concatMap(reply -> send(outbound, reply)
                                    .then(Mono.defer(() -> reply == STARTTLS_READY ? startTls(connection[0], session) : Mono.<Void>empty())))
                            .then();
                })
                .bindNow(Duration.ofSeconds(10));
        log.info("SMTP server listening on {}:{} ({}{})", settings.host(), port(), settings.hostname(),
                settings.implicitTls() ? ", implicit TLS" : settings.tls() != null ? ", STARTTLS" : ", no TLS");
    }

    private static Mono<Void> send(NettyOutbound outbound, String reply) {
        return outbound.sendString(Mono.just(reply), StandardCharsets.US_ASCII).then();
    }

    /** TLS handler in front, fresh line decoder (no plain-text leftovers mixed into encrypted input). */
    private Mono<Void> startTls(Connection connection, Session session) {
        ChannelPipeline pipeline = connection.channel().pipeline();
        SslHandler ssl = settings.tls().newHandler(connection.channel().alloc());
        pipeline.replace("smtp-lines", "smtp-lines", new LineBasedFrameDecoder(MAX_LINE, true, false));
        pipeline.addFirst("smtp-tls", ssl);
        return Mono.create(sink -> ssl.handshakeFuture().addListener(future -> {
            if (future.isSuccess()) {
                log.debug("SMTP {}: TLS {}", session.remoteIp, ssl.engine().getSession().getProtocol());
                sink.success();
            } else {
                log.info("SMTP {}: TLS handshake failed: {}", session.remoteIp, future.cause().toString());
                session.closing = true;
                connection.dispose();
                sink.success();
            }
        }));
    }

    public int port() {
        return server.port();
    }

    public void stop() {
        if (server != null) {
            server.disposeNow(Duration.ofSeconds(5));
        }
    }

    private enum State { COMMAND, DATA, AUTH_PLAIN, AUTH_LOGIN_USER, AUTH_LOGIN_PASSWORD }

    private final class Session {
        final String remoteIp;
        boolean tls;
        State state = State.COMMAND;
        String helo;
        boolean extended;
        String authenticatedUser;
        String pendingAuthUser;
        String mailFrom;
        final List<SmtpBackend.Recipient> recipients = new ArrayList<>();
        final ByteArrayOutputStream data = new ByteArrayOutputStream();
        boolean dataTooBig;
        int badCommands;
        int authFailures;
        volatile boolean closing;

        Session(String remoteIp, boolean tls) {
            this.remoteIp = remoteIp;
            this.tls = tls;
        }

        Mono<String> onLine(Line input) {
            if (tls && !input.tls()) {
                // Sent in plain text after STARTTLS, before the handshake: never execute it
                return Mono.error(new IllegalStateException("plain-text input after STARTTLS"));
            }
            byte[] line = stripCr(input.bytes());
            return switch (state) {
                case DATA -> onDataLine(line);
                case AUTH_PLAIN -> authPlain(ascii(line));
                case AUTH_LOGIN_USER -> authLoginUser(ascii(line));
                case AUTH_LOGIN_PASSWORD -> authLoginPassword(ascii(line));
                case COMMAND -> onCommand(ascii(line));
            };
        }

        private Mono<String> onCommand(String line) {
            String[] parts = line.trim().split("\\s+", 2);
            String verb = parts[0].toUpperCase(Locale.ROOT);
            String arg = parts.length > 1 ? parts[1].trim() : "";
            return switch (verb) {
                case "EHLO", "HELO" -> hello(verb.equals("EHLO"), arg);
                case "AUTH" -> auth(arg);
                case "MAIL" -> mail(arg);
                case "RCPT" -> rcpt(arg);
                case "DATA" -> dataStart();
                case "RSET" -> {
                    resetTransaction();
                    yield reply("250 2.0.0 OK");
                }
                case "NOOP" -> reply("250 2.0.0 OK");
                case "VRFY" -> reply("252 2.5.0 Cannot VRFY user, but will accept message");
                case "STARTTLS" -> startTls(arg);
                case "QUIT" -> {
                    closing = true;
                    yield reply("221 2.0.0 " + settings.hostname() + " Bye");
                }
                case "" -> reply("500 5.5.2 Empty command");
                default -> bad("500 5.5.1 Command not recognized");
            };
        }

        private Mono<String> hello(boolean ehlo, String arg) {
            if (arg.isEmpty()) {
                return reply("501 5.5.4 Hostname required");
            }
            helo = arg;
            extended = ehlo;
            resetTransaction();
            if (!ehlo) {
                return reply("250 " + settings.hostname());
            }
            StringBuilder sb = new StringBuilder();
            sb.append("250-").append(settings.hostname()).append(" Hello ").append(arg).append("\r\n");
            sb.append("250-SIZE ").append(settings.maxMessageBytes()).append("\r\n");
            sb.append("250-8BITMIME\r\n");
            sb.append("250-PIPELINING\r\n");
            if (settings.tls() != null && !tls) {
                sb.append("250-STARTTLS\r\n");
            }
            if (authAllowed()) {
                sb.append("250-AUTH PLAIN LOGIN\r\n");
            }
            sb.append("250 ENHANCEDSTATUSCODES\r\n");
            return Mono.just(sb.toString());
        }

        private boolean authAllowed() {
            return tls || settings.authWithoutTls();
        }

        private Mono<String> startTls(String arg) {
            if (settings.tls() == null) {
                return reply("454 4.7.0 TLS not available");
            }
            if (tls) {
                return bad("503 5.5.1 TLS already active");
            }
            if (!arg.isEmpty()) {
                return reply("501 5.5.4 Syntax: STARTTLS");
            }
            // RFC 3207 4.2: everything known about the client is forgotten, a new EHLO is expected
            tls = true;
            helo = null;
            extended = false;
            authenticatedUser = null;
            resetTransaction();
            return Mono.just(STARTTLS_READY);
        }

        private Mono<String> auth(String arg) {
            if (!authAllowed()) {
                return reply("538 5.7.11 Encryption required for requested authentication mechanism");
            }
            if (helo == null || !extended) {
                return reply("503 5.5.1 EHLO first");
            }
            if (authenticatedUser != null) {
                return reply("503 5.5.1 Already authenticated");
            }
            if (mailFrom != null) {
                return reply("503 5.5.1 AUTH not allowed during a mail transaction");
            }
            String[] parts = arg.split("\\s+", 2);
            String mechanism = parts[0].toUpperCase(Locale.ROOT);
            String initial = parts.length > 1 ? parts[1] : null;
            if (mechanism.equals("PLAIN")) {
                if (initial == null) {
                    state = State.AUTH_PLAIN;
                    return reply("334 ");
                }
                return authPlain(initial);
            }
            if (mechanism.equals("LOGIN")) {
                if (initial != null) {
                    return authLoginUser(initial);
                }
                state = State.AUTH_LOGIN_USER;
                return reply("334 " + b64("Username:"));
            }
            return reply("504 5.5.4 Unrecognized authentication type");
        }

        private Mono<String> authPlain(String response) {
            state = State.COMMAND;
            if (response.equals("*")) {
                return reply("501 5.0.0 Authentication cancelled");
            }
            String decoded = unb64(response);
            // authzid \0 authcid \0 password
            String[] fields = decoded == null ? new String[0] : decoded.split("\0", -1);
            if (fields.length != 3) {
                return reply("501 5.5.2 Cannot decode AUTH PLAIN response");
            }
            return checkPassword(fields[1], fields[2]);
        }

        private Mono<String> authLoginUser(String response) {
            if (response.equals("*")) {
                state = State.COMMAND;
                return reply("501 5.0.0 Authentication cancelled");
            }
            pendingAuthUser = unb64(response);
            if (pendingAuthUser == null) {
                state = State.COMMAND;
                return reply("501 5.5.2 Cannot decode username");
            }
            state = State.AUTH_LOGIN_PASSWORD;
            return reply("334 " + b64("Password:"));
        }

        private Mono<String> authLoginPassword(String response) {
            state = State.COMMAND;
            String user = pendingAuthUser;
            pendingAuthUser = null;
            if (response.equals("*")) {
                return reply("501 5.0.0 Authentication cancelled");
            }
            String password = unb64(response);
            if (password == null) {
                return reply("501 5.5.2 Cannot decode password");
            }
            return checkPassword(user, password);
        }

        private Mono<String> checkPassword(String user, String password) {
            String username = user.trim().toLowerCase(Locale.ROOT);
            Duration wait = guard.retryAfter(remoteIp, username);
            if (!wait.isZero()) {
                // Blocked: the password is not even checked, the connection is closed
                log.warn("SMTP {}: login as {} refused, blocked for {}", remoteIp, username, LoginGuard.human(wait));
                closing = true;
                return reply("421 4.7.0 Too many failed logins, try again in " + LoginGuard.human(wait));
            }
            return blocking(() -> backend.authenticate(username, password))
                    .flatMap(ok -> {
                        if (ok) {
                            guard.success(remoteIp, username);
                            authenticatedUser = username;
                            log.info("SMTP {}: authenticated as {}", remoteIp, username);
                            return Mono.just("235 2.7.0 Authentication successful\r\n");
                        }
                        guard.failure(remoteIp, username);
                        log.warn("SMTP {}: authentication failed for {}", remoteIp, username);
                        badCommands++;
                        if (++authFailures >= MAX_AUTH_FAILURES) {
                            closing = true;
                        }
                        String answer = closing
                                ? "535 5.7.8 Authentication credentials invalid, closing connection\r\n"
                                : "535 5.7.8 Authentication credentials invalid\r\n";
                        Duration delay = guard.failureDelay();
                        return delay.isZero() ? Mono.just(answer) : Mono.just(answer).delayElement(delay);
                    });
        }

        private Mono<String> mail(String arg) {
            if (helo == null) {
                return reply("503 5.5.1 EHLO/HELO first");
            }
            if (mailFrom != null) {
                return reply("503 5.5.1 Sender already specified");
            }
            if (!arg.regionMatches(true, 0, "FROM:", 0, 5)) {
                return reply("501 5.5.4 Syntax: MAIL FROM:<address>");
            }
            String rest = arg.substring(5).trim();
            String address = path(rest);
            if (address == null) {
                return reply("501 5.1.7 Bad sender address syntax");
            }
            Long size = sizeParam(rest);
            if (size != null && size > settings.maxMessageBytes()) {
                return reply("552 5.3.4 Message size exceeds fixed limit");
            }
            if (authenticatedUser != null && !address.isEmpty() && !address.equals(authenticatedUser)) {
                return reply("553 5.7.1 Sender address not owned by authenticated user");
            }
            mailFrom = address;
            return reply("250 2.1.0 Sender OK");
        }

        private Mono<String> rcpt(String arg) {
            if (mailFrom == null) {
                return reply("503 5.5.1 MAIL first");
            }
            if (!arg.regionMatches(true, 0, "TO:", 0, 3)) {
                return reply("501 5.5.4 Syntax: RCPT TO:<address>");
            }
            String address = path(arg.substring(3).trim());
            if (address == null || address.isEmpty() || address.indexOf('@') < 1) {
                return reply("501 5.1.3 Bad recipient address syntax");
            }
            if (recipients.size() >= MAX_RECIPIENTS) {
                return reply("452 4.5.3 Too many recipients");
            }
            String user = authenticatedUser;
            return blocking(() -> backend.route(address, user))
                    .map(route -> switch (route) {
                        case LOCAL, OUTBOUND, TRAP -> {
                            recipients.add(new SmtpBackend.Recipient(address, route));
                            yield "250 2.1.5 Recipient OK\r\n";
                        }
                        case UNKNOWN_USER -> "550 5.1.1 No such user here\r\n";
                        case RELAY_DENIED -> {
                            log.warn("SMTP {}: relay denied {} -> {}", remoteIp, mailFrom, address);
                            yield "550 5.7.1 Relaying denied: authentication required\r\n";
                        }
                        case OUTBOUND_DISABLED -> "550 5.7.1 Outbound delivery is disabled on this server\r\n";
                        case LIST_CLOSED -> "550 5.7.1 This list takes mail only from logged-in senders\r\n";
                    });
        }

        private Mono<String> dataStart() {
            if (mailFrom == null) {
                return reply("503 5.5.1 MAIL first");
            }
            if (recipients.isEmpty()) {
                return reply("554 5.5.1 No valid recipients");
            }
            state = State.DATA;
            data.reset();
            dataTooBig = false;
            return reply("354 End data with <CR><LF>.<CR><LF>");
        }

        private Mono<String> onDataLine(byte[] line) {
            if (line.length == 1 && line[0] == '.') {
                state = State.COMMAND;
                if (dataTooBig) {
                    resetTransaction();
                    return reply("552 5.3.4 Message size exceeds fixed limit");
                }
                SmtpBackend.Envelope envelope = new SmtpBackend.Envelope(
                        remoteIp, helo, mailFrom, authenticatedUser, List.copyOf(recipients));
                byte[] message = withReceivedHeader(data.toByteArray());
                resetTransaction();
                return blocking(() -> {
                    backend.accept(envelope, message);
                    return true;
                })
                        .map(ok -> "250 2.0.0 Message accepted\r\n")
                        .onErrorResume(e -> {
                            log.error("SMTP {}: message from {} not stored: {}", remoteIp, envelope.mailFrom(), e.toString());
                            return Mono.just("451 4.3.0 Message not stored, try again later\r\n");
                        });
            }
            if (!dataTooBig) {
                int offset = line.length > 0 && line[0] == '.' ? 1 : 0; // dot-unstuffing (RFC 5321 4.5.2)
                data.write(line, offset, line.length - offset);
                data.write('\r');
                data.write('\n');
                if (data.size() > settings.maxMessageBytes()) {
                    dataTooBig = true;
                    data.reset();
                }
            }
            return Mono.empty();
        }

        private byte[] withReceivedHeader(byte[] message) {
            String received = "Received: from " + helo + " ([" + remoteIp + "])\r\n"
                    + "\tby " + settings.hostname() + " with " + (tls ? "ESMTPS" : "ESMTP") + (authenticatedUser != null ? "A" : "") + ";\r\n"
                    + "\t" + DateTimeFormatter.RFC_1123_DATE_TIME.format(ZonedDateTime.now()) + "\r\n";
            byte[] header = received.getBytes(StandardCharsets.US_ASCII);
            byte[] all = new byte[header.length + message.length];
            System.arraycopy(header, 0, all, 0, header.length);
            System.arraycopy(message, 0, all, header.length, message.length);
            return all;
        }

        private void resetTransaction() {
            mailFrom = null;
            recipients.clear();
            data.reset();
            dataTooBig = false;
        }

        private Mono<String> bad(String text) {
            if (++badCommands > MAX_BAD_COMMANDS) {
                closing = true;
                return reply("421 4.7.0 Too many errors, closing connection");
            }
            return reply(text);
        }
    }

    private static Mono<String> reply(String line) {
        return Mono.just(line + "\r\n");
    }

    private static <T> Mono<T> blocking(java.util.concurrent.Callable<T> work) {
        return Mono.fromCallable(work).subscribeOn(Schedulers.boundedElastic());
    }

    /** "<a@b> SIZE=1" → "a@b"; "<>" → "" (null sender of bounces); malformed → null. */
    static String path(String text) {
        int start = text.indexOf('<');
        int end = text.indexOf('>', start + 1);
        if (start != 0 || end < 0) {
            return null;
        }
        String address = text.substring(1, end).trim();
        int colon = address.indexOf(':'); // source route "@a,@b:user@host" (RFC 5321 4.1.2) — ignored
        if (address.startsWith("@") && colon > 0) {
            address = address.substring(colon + 1);
        }
        if (address.contains(" ") || address.contains("<")) {
            return null;
        }
        return address.toLowerCase(Locale.ROOT);
    }

    private static Long sizeParam(String text) {
        for (String param : text.substring(text.indexOf('>') + 1).trim().split("\\s+")) {
            if (param.regionMatches(true, 0, "SIZE=", 0, 5)) {
                try {
                    return Long.parseLong(param.substring(5));
                } catch (NumberFormatException e) {
                    return null;
                }
            }
        }
        return null;
    }

    private static byte[] stripCr(byte[] line) {
        if (line.length > 0 && line[line.length - 1] == '\r') {
            byte[] copy = new byte[line.length - 1];
            System.arraycopy(line, 0, copy, 0, copy.length);
            return copy;
        }
        return line;
    }

    private static String ascii(byte[] line) {
        return new String(line, StandardCharsets.UTF_8);
    }

    private static String b64(String text) {
        return Base64.getEncoder().encodeToString(text.getBytes(StandardCharsets.UTF_8));
    }

    private static String unb64(String text) {
        try {
            return new String(Base64.getDecoder().decode(text.trim()), StandardCharsets.UTF_8);
        } catch (IllegalArgumentException e) {
            return null;
        }
    }
}