package com.mailserver.service;

import jakarta.annotation.PostConstruct;
import jakarta.mail.MessagingException;
import jakarta.mail.Session;
import jakarta.mail.internet.InternetAddress;
import jakarta.mail.internet.MimeMessage;
import org.apache.james.jdkim.DKIMVerifier;
import org.apache.james.jdkim.api.PublicKeyRecordRetriever;
import org.apache.james.jdkim.api.SignatureRecord;
import org.apache.james.jdkim.exceptions.FailException;
import org.apache.james.jdkim.impl.DNSPublicKeyRecordRetriever;
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.openssl.PEMKeyPair;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
import org.simplejavamail.utils.mail.dkim.Canonicalization;
import org.simplejavamail.utils.mail.dkim.DkimMessage;
import org.simplejavamail.utils.mail.dkim.DkimSigner;
import org.simplejavamail.utils.mail.dkim.SigningAlgorithm;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.Reader;
import java.io.StringReader;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.attribute.PosixFilePermissions;
import java.security.KeyFactory;
import java.security.KeyPairGenerator;
import java.security.interfaces.RSAPrivateCrtKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.spec.RSAPublicKeySpec;
import java.util.Base64;
import java.util.List;
import java.util.Locale;
import java.util.Properties;

/**
 * DKIM (RFC 6376): signs our outgoing mail and verifies incoming mail.
 * Signing — simple-java-mail utils-mail-dkim (rsa-sha256, relaxed/relaxed); verification — Apache James jDKIM
 * with the sender's public key from DNS. The key is generated once and kept in {@code keys-dir/dkim-private.pem};
 * the DNS record to publish is printed to the log at start and available from {@link #dnsRecord()}.
 */
@Service
public class DkimService {

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

    /** Result of checking an incoming message. */
    public record Verification(boolean signed, boolean passed, String domain) {
        static final Verification NONE = new Verification(false, false, null);
    }

    private final String domain;
    private final String selector;
    private final Path keysDir;
    private final boolean enabled;
    private final DKIMVerifier verifier;

    private RSAPrivateKey privateKey;
    private String publicKeyBase64;

    @Autowired
    public DkimService(@Value("${mailserver.dkim.domain:localhost}") String domain,
                       @Value("${mailserver.dkim.selector:mail}") String selector,
                       @Value("${mailserver.dkim.keys-dir:./dkim-keys}") String keysDir,
                       @Value("${mailserver.dkim.enabled:true}") boolean enabled) {
        this(domain, selector, Paths.get(keysDir), enabled, new DNSPublicKeyRecordRetriever());
    }

    // Tests pass their own public key lookup instead of DNS
    DkimService(String domain, String selector, Path keysDir, boolean enabled, PublicKeyRecordRetriever keys) {
        this.domain = domain.toLowerCase(Locale.ROOT);
        this.selector = selector;
        this.keysDir = keysDir;
        this.enabled = enabled;
        this.verifier = new DKIMVerifier(keys);
    }

    @PostConstruct
    public void init() {
        if (!enabled) {
            log.info("DKIM signing disabled");
            return;
        }
        try {
            Files.createDirectories(keysDir);
            Path keyFile = keysDir.resolve("dkim-private.pem");
            if (Files.exists(keyFile)) {
                privateKey = readPrivateKey(Files.readString(keyFile));
                log.info("DKIM: key loaded for {}", domain);
            } else {
                KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");
                generator.initialize(2048);
                privateKey = (RSAPrivateKey) generator.generateKeyPair().getPrivate();
                String pem = "-----BEGIN PRIVATE KEY-----\n"
                        + Base64.getMimeEncoder(64, "\n".getBytes()).encodeToString(privateKey.getEncoded())
                        + "\n-----END PRIVATE KEY-----\n";
                Files.writeString(keyFile, pem);
                restrict(keyFile);
                log.info("DKIM: new key generated for {}", domain);
            }
            publicKeyBase64 = Base64.getEncoder().encodeToString(publicKeyOf(privateKey));
            log.info("DKIM DNS record — TXT {}._domainkey.{}: {}", selector, domain, dnsRecord());
        } catch (Exception e) {
            throw new IllegalStateException("DKIM key cannot be loaded from " + keysDir, e);
        }
    }

    /** TXT value to publish at {@code <selector>._domainkey.<domain>}. */
    public String dnsRecord() {
        return publicKeyBase64 == null ? null : "v=DKIM1; k=rsa; p=" + publicKeyBase64;
    }

    /**
     * Signs the message when it is from our domain; otherwise returns it unchanged.
     * The result is the exact byte sequence to send — do not re-serialize it.
     */
    public byte[] sign(byte[] message) throws MessagingException, IOException {
        if (!enabled || privateKey == null) {
            return message;
        }
        MimeMessage mime = new MimeMessage(Session.getInstance(new Properties()), new ByteArrayInputStream(message));
        if (mime.getFrom() == null || mime.getFrom().length == 0
                || !(mime.getFrom()[0] instanceof InternetAddress from)
                || !domain.equals(domainOf(from.getAddress()))) {
            return message;
        }
        DkimSigner signer = new DkimSigner(domain, selector, privateKey);
        signer.setCheckDomainKey(false); // the DNS record may not be published yet (tests, new server)
        signer.setSigningAlgorithm(SigningAlgorithm.SHA256_WITH_RSA);
        signer.setHeaderCanonicalization(Canonicalization.RELAXED);
        signer.setBodyCanonicalization(Canonicalization.RELAXED);
        signer.setLengthParam(false);
        signer.setCopyHeaderFields(false);
        ByteArrayOutputStream out = new ByteArrayOutputStream(message.length + 1024);
        DkimMessage signed = new DkimMessage(mime, signer);
        signed.setAllow8bitMIME(true);
        signed.writeTo(out);
        return out.toByteArray();
    }

    /** Checks DKIM signatures of an incoming message (public key from the signer's DNS). */
    public Verification verify(byte[] message) {
        try {
            List<SignatureRecord> valid;
            synchronized (verifier) { // DKIMVerifier keeps results between calls
                verifier.resetResults();
                valid = verifier.verify(new ByteArrayInputStream(message));
            }
            if (valid == null) {
                return Verification.NONE;
            }
            if (valid.isEmpty()) {
                return new Verification(true, false, null);
            }
            return new Verification(true, true, valid.get(0).getDToken().toString().toLowerCase(Locale.ROOT));
        } catch (FailException e) {
            log.info("DKIM: signature failed: {}", e.getMessage());
            return new Verification(true, false, null);
        } catch (IOException e) {
            log.warn("DKIM: message cannot be read: {}", e.getMessage());
            return new Verification(true, false, null);
        }
    }

    static String domainOf(String address) {
        int at = address == null ? -1 : address.lastIndexOf('@');
        return at < 0 ? "" : address.substring(at + 1).toLowerCase(Locale.ROOT);
    }

    private static RSAPrivateKey readPrivateKey(String pem) throws IOException {
        try (Reader reader = new StringReader(pem); PEMParser parser = new PEMParser(reader)) {
            Object object = parser.readObject();
            JcaPEMKeyConverter converter = new JcaPEMKeyConverter();
            // PKCS#8 ("PRIVATE KEY") and the older PKCS#1 ("RSA PRIVATE KEY") written by earlier versions
            if (object instanceof PEMKeyPair pair) {
                return (RSAPrivateKey) converter.getKeyPair(pair).getPrivate();
            }
            if (object instanceof PrivateKeyInfo info) {
                return (RSAPrivateKey) converter.getPrivateKey(info);
            }
            throw new IOException("Not an RSA private key: " + (object == null ? "empty file" : object.getClass().getSimpleName()));
        }
    }

    private static byte[] publicKeyOf(RSAPrivateKey key) throws Exception {
        RSAPrivateCrtKey crt = (RSAPrivateCrtKey) key;
        return KeyFactory.getInstance("RSA")
                .generatePublic(new RSAPublicKeySpec(crt.getModulus(), crt.getPublicExponent()))
                .getEncoded();
    }

    private static void restrict(Path file) {
        try {
            Files.setPosixFilePermissions(file, PosixFilePermissions.fromString("rw-------"));
        } catch (UnsupportedOperationException | IOException ignored) {
            // not a POSIX file system
        }
    }
}