package de.mlessmann.certassist.openssl;

import static de.mlessmann.certassist.Constants.CERTASSIST_TMP_PREFIX;
import static java.util.concurrent.TimeUnit.*;
import static org.slf4j.LoggerFactory.getLogger;

import de.mlessmann.certassist.DeleteRecursiveFileVisitor;
import de.mlessmann.certassist.ExecutableResolver;
import de.mlessmann.certassist.except.CommandLineOperationException;
import de.mlessmann.certassist.except.UnresolvableCLIDependency;
import de.mlessmann.certassist.openssl.CertificateRequest.RequestType;
import de.mlessmann.certassist.openssl.CertificateSubject.CertificateSubjectBuilder;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.springframework.lang.NonNull;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import org.zeroturnaround.exec.ProcessExecutor;
import org.zeroturnaround.exec.ProcessResult;
import org.zeroturnaround.exec.StartedProcess;
import org.zeroturnaround.exec.stream.slf4j.Slf4jStream;

@Service
@RequiredArgsConstructor
@Slf4j
public class OpenSSLCertificateCreator {

    private static final Logger openSSLLogger = getLogger("OpenSSL-Logger");
    public static final String OPENSSL_CERT_SUBJECT_TEMPLATE =
        "/C=ISO-COUNTRY/ST=STATE/L=LOCALITY/O=ORGANIZATION/CN=COMMON-NAME";
    private static final String CSR_EXT_TEMPLATE =
        """
        authorityKeyIdentifier=keyid,issuer
        basicConstraints=CA:FALSE
        keyUsage = digitalSignature, nonRepudiation, keyEncipherment, dataEncipherment
        subjectAltName = @alt_names

        [alt_names]
        """;
    private static final Pattern FINGERPRINT_EXTRACTOR = Pattern.compile(
        "^(?<algo>[0-9a-zA-Z]+) (?i)Fingerprint(?-i)=(?<finger>[a-z:A-Z0-9]+)"
    );
    private static final String OSSL_ENV_KEY_PW = "KEY_PASS";
    private static final String OSSL_ARG_KEY_PW = "env:" + OSSL_ENV_KEY_PW;
    private final AtomicBoolean versionLogged = new AtomicBoolean(false);

    private final ExecutableResolver executableResolver;
    private final CertificatePasswordProvider passwordProvider;
    private final CertificateProvider certificateProvider;

    private static String buildSubjectArg(CertificateRequest request) {
        String certSubject = OPENSSL_CERT_SUBJECT_TEMPLATE
            .replace("ISO-COUNTRY", request.getSubject().getCountry())
            .replace("STATE", request.getSubject().getState())
            .replace("LOCALITY", request.getSubject().getLocality())
            .replace("ORGANIZATION", request.getSubject().getOrganization())
            .replace("COMMON-NAME", request.getSubject().getCommonName());

        if (StringUtils.isNotBlank(request.getSubject().getOrganizationalUnit())) {
            certSubject += "/OU=" + request.getSubject().getOrganizationalUnit();
        }

        if (StringUtils.isNotBlank(request.getSubject().getEmailAddress())) {
            certSubject += "/emailAddress=" + request.getSubject().getEmailAddress();
        }
        return certSubject;
    }

    private static void killIfActive(StartedProcess process) {
        if (process == null) {
            return;
        }
        Process sysProc = process.getProcess();
        if (!sysProc.isAlive()) {
            return;
        }

        try {
            log.debug("Process is still alive. Asking politely for it to destroy itself and waiting on exit for 2s");
            sysProc.destroy();
            sysProc.waitFor(2_000, MILLISECONDS);
        } catch (InterruptedException e) {
            log.debug("Interrupted while waiting for process to terminate. Registering forceful termination onExit", e);
            Runtime.getRuntime().addShutdownHook(new Thread(sysProc::destroyForcibly));
        }
    }

    @NonNull
    public OpenSSLCertificateResult createCertificate(CertificateRequest request) throws CommandLineOperationException {
        Path tmpDir;
        try {
            tmpDir = Files.createTempDirectory(CERTASSIST_TMP_PREFIX);
        } catch (IOException e) {
            throw new CommandLineOperationException("Could not create temporary directory for certificate creation", e);
        }

        String keyPassphrase = passwordProvider.generateNewPassword();
        Path keyFile = createKeyfile(request, tmpDir.resolve("certificate.key"), keyPassphrase);
        if (
            request.getType() == RequestType.ROOT_AUTHORITY || request.getType() == RequestType.STANDALONE_CERTIFICATE
        ) {
            Path certificate = createCertificate(request, keyFile, tmpDir.resolve("certificate.crt"), keyPassphrase);
            String fingerprint = getCertificateFingerprint(certificate);
            passwordProvider.setPasswordFor(fingerprint, keyPassphrase);
            return new OpenSSLCertificateResult(tmpDir, certificate, keyFile, certificate, fingerprint);
        }

        try (var certAuthority = certificateProvider.requestCertificateUsage(request.getTrustingAuthority())) {
            Path signingRequest = createSigningRequest(request, keyFile, tmpDir.resolve("child.csr"), keyPassphrase);
            Path signedCert = signCertificate(
                request,
                certAuthority.certificatePath(),
                certAuthority.certificateKeyPath(),
                passwordProvider.getPasswordFor(certAuthority.fingerprint()),
                signingRequest
            );
            String fingerprint = getCertificateFingerprint(signedCert);
            passwordProvider.setPasswordFor(fingerprint, keyPassphrase);

            Path fullchain = tmpDir.resolve("fullchain.pem");
            try {
                Path certAuthFullchain = Optional
                    .ofNullable(certAuthority.fullchainPath())
                    .orElse(certAuthority.certificatePath());
                Files.write(fullchain, Files.readAllBytes(certAuthFullchain), StandardOpenOption.CREATE);
                Files.write(fullchain, Files.readAllBytes(signedCert), StandardOpenOption.APPEND);
            } catch (IOException e) {
                throw new CommandLineOperationException("Failed to create fullchain file.", e);
            }

            return new OpenSSLCertificateResult(tmpDir, signedCert, keyFile, fullchain, fingerprint);
        }
    }

    private Path createKeyfile(CertificateRequest request, Path outFile, String filePassword)
        throws CommandLineOperationException {
        Path keyFile = outFile.toAbsolutePath();
        log.debug("Writing new certificate key to {}", keyFile);

        StartedProcess keygenProc = null;
        try {
            keygenProc =
            new ProcessExecutor()
                .command(
                    resolveOpenSSL(),
                    "genrsa",
                    "-out",
                    keyFile.toString(),
                    "-aes256",
                    "-passout",
                    OSSL_ARG_KEY_PW,
                    Integer.toString(request.getRequestedKeyLength())
                )
                .environment(OSSL_ENV_KEY_PW, filePassword)
                .redirectOutput(Slf4jStream.of(openSSLLogger).asDebug())
                .redirectError(Slf4jStream.of(openSSLLogger).asError())
                .start();
            keygenProc.getFuture().get(3, MINUTES);
        } catch (IOException | ExecutionException | InterruptedException | TimeoutException e) {
            throw new CommandLineOperationException("Failure running OpenSSL keygen command.", e);
        } finally {
            killIfActive(keygenProc);
        }
        return keyFile;
    }

    private Path createCertificate(CertificateRequest request, Path keyFile, Path outFile, String keyPassphrase)
        throws CommandLineOperationException {
        log.debug("Writing new certificate file {}", outFile);

        String certSubject = buildSubjectArg(request);
        StartedProcess certGenProc = null;
        try {
            certGenProc =
            new ProcessExecutor()
                .command(
                    resolveOpenSSL(),
                    "req",
                    "-x509",
                    "-new",
                    "-passin",
                    OSSL_ARG_KEY_PW,
                    "-key",
                    keyFile.toString(),
                    "-sha256",
                    "-days",
                    Integer.toString(request.getRequestedValidityDays()),
                    "-out",
                    outFile.toString(),
                    "-utf8",
                    "-subj",
                    certSubject
                )
                .environment(OSSL_ENV_KEY_PW, keyPassphrase)
                .redirectOutput(Slf4jStream.of(openSSLLogger).asDebug())
                .redirectError(Slf4jStream.of(openSSLLogger).asError())
                .start();
            certGenProc.getFuture().get(30, SECONDS);
        } catch (IOException | ExecutionException | InterruptedException | TimeoutException e) {
            throw new CommandLineOperationException("Failure running OpenSSL req command.", e);
        } finally {
            killIfActive(certGenProc);
        }
        return outFile;
    }

    private Path createSigningRequest(CertificateRequest request, Path keyFile, Path outFile, String certPassword)
        throws CommandLineOperationException {
        log.atDebug().log("Writing new certificate signing request file {}", outFile);

        String certSubject = buildSubjectArg(request);
        StartedProcess certGenProc = null;
        try {
            certGenProc =
            new ProcessExecutor()
                .command(
                    resolveOpenSSL(),
                    "req",
                    "-new",
                    "-passin",
                    OSSL_ARG_KEY_PW,
                    "-key",
                    keyFile.toString(),
                    "-sha256",
                    "-out",
                    outFile.toString(),
                    "-utf8",
                    "-subj",
                    certSubject
                )
                .environment(OSSL_ENV_KEY_PW, certPassword)
                .redirectOutput(Slf4jStream.of(openSSLLogger).asDebug())
                .redirectError(Slf4jStream.of(openSSLLogger).asError())
                .start();
            certGenProc.getFuture().get(30, SECONDS);
        } catch (IOException | ExecutionException | InterruptedException | TimeoutException e) {
            throw new CommandLineOperationException("Failure running OpenSSL req command.", e);
        } finally {
            killIfActive(certGenProc);
        }
        return outFile;
    }

    public boolean verifyCertificate(@NonNull Path fullChainFile, @NonNull Path trustedCA)
        throws CommandLineOperationException {
        return verifyCertificate(fullChainFile, List.of(trustedCA));
    }

    public boolean verifyCertificate(@NonNull Path fullChainFile, @NonNull List<Path> trustedCAs)
        throws CommandLineOperationException {
        if (CollectionUtils.isEmpty(trustedCAs)) {
            throw new IllegalArgumentException(
                "At least one trusted CA certificate must be provided to run the verification command."
            );
        }

        Path tmpDir = null;
        StartedProcess verifyCommand = null;
        try {
            Path tempTrustedBundle;
            if (trustedCAs.size() == 1) {
                tempTrustedBundle = trustedCAs.getFirst();
            } else {
                tmpDir = Files.createTempDirectory(CERTASSIST_TMP_PREFIX);
                tempTrustedBundle = tmpDir.resolve("trusted_bundle.pem");
                for (Path ca : trustedCAs) {
                    Files.writeString(
                        tempTrustedBundle,
                        Files.readString(ca),
                        StandardCharsets.UTF_8,
                        StandardOpenOption.CREATE,
                        StandardOpenOption.APPEND
                    );
                }
            }

            verifyCommand =
            new ProcessExecutor()
                .command(resolveOpenSSL(), "verify", "-CAfile", tempTrustedBundle.toString(), fullChainFile.toString())
                .redirectOutput(Slf4jStream.of(openSSLLogger).asError())
                .redirectError(Slf4jStream.of(openSSLLogger).asError())
                .start();
            var verifyResult = verifyCommand.getFuture().get(30, SECONDS);
            return verifyResult.getExitValue() == 0;
        } catch (IOException | InterruptedException | ExecutionException | TimeoutException e) {
            throw new RuntimeException(e);
        } finally {
            killIfActive(verifyCommand);
            if (tmpDir != null) {
                try {
                    Files.walkFileTree(tmpDir, new DeleteRecursiveFileVisitor());
                    Files.deleteIfExists(tmpDir);
                } catch (IOException e) {
                    log.error("Failed to clean up temporary verification directory: {}", tmpDir, e);
                }
            }
        }
    }

    /**
     * Checks whether the provided key file is encrypted using a passphrase
     */
    public boolean isKeyEncrypted(@NonNull Path keyFile) throws CommandLineOperationException {
        // If the key is not encrypted, any passphrase will work -> so generate a random one to check.
        String passphrase = UUID.randomUUID().toString();
        boolean firstPass = verifyKeyPassphrase(keyFile, passphrase);
        if (firstPass) {
            // Try with another random passphrase in case we randomly got lucky guessing the passphrase the first time.
            passphrase = UUID.randomUUID().toString();
            return !verifyKeyPassphrase(keyFile, passphrase);
        }
        return true;
    }

    /**
     * Verifies a passphrase against a provided key.
     * @implNote Due to the implementation of the OpenSSL cli, <em>any password</em> will be valid for unencrypted keys. (Check with {@link #isKeyEncrypted(Path).)
     */
    public boolean verifyKeyPassphrase(@NonNull Path keyFile, @NonNull String passphrase)
        throws CommandLineOperationException {
        StartedProcess verifyCommand = null;
        try {
            verifyCommand =
            new ProcessExecutor()
                .command(
                    resolveOpenSSL(),
                    "rsa",
                    "-check",
                    "-in",
                    keyFile.toString(),
                    "-passin",
                    "pass:" + passphrase,
                    "-noout"
                )
                .redirectOutput(Slf4jStream.of(openSSLLogger).asError())
                .redirectError(Slf4jStream.of(openSSLLogger).asError())
                .start();
            var verifyResult = verifyCommand.getFuture().get(30, SECONDS);
            return verifyResult.getExitValue() == 0;
        } catch (IOException | InterruptedException | ExecutionException | TimeoutException e) {
            throw new CommandLineOperationException("Failed to verify key encryption", e);
        } finally {
            killIfActive(verifyCommand);
        }
    }

    private Path signCertificate(
        CertificateRequest request,
        Path caCert,
        Path caKey,
        String caKeyPassphrase,
        Path csrFile
    ) throws CommandLineOperationException {
        Path outFile = csrFile.resolveSibling(csrFile.getFileName().toString().replace(".csr", ".crt"));
        log.debug("Writing new signed certificate file {}", outFile);
        Path extFile = csrFile.resolveSibling(csrFile.getFileName().toString().replace(".csr", ".ext"));

        try {
            String extContent = CSR_EXT_TEMPLATE;
            List<String> altNames = Optional
                .ofNullable(request.getExtension())
                .map(CertificateRequestExtension::getAlternativeNames)
                .orElse(List.of());
            if (!altNames.isEmpty()) {
                AtomicInteger counter = new AtomicInteger(1);
                String altNamesContent = altNames
                    .stream()
                    .map(name -> "DNS.%d = %s".formatted(counter.getAndIncrement(), name))
                    .collect(Collectors.joining("\n"));
                extContent = extContent.replaceAll("\\[alt_names]\n?", "[alt_names]\n" + altNamesContent);
            } else {
                extContent = extContent.replaceAll("\\s*subjectAltName\\s+=\\s+@alt_names\n?", "");
                extContent = extContent.replaceAll("\\[alt_names]\n?", "");
            }

            log.debug("Writing extension file content: \n {}", extContent);
            Files.writeString(
                extFile,
                extContent,
                StandardCharsets.UTF_8,
                StandardOpenOption.CREATE,
                StandardOpenOption.TRUNCATE_EXISTING
            );
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        StartedProcess certGenProc = null;
        try {
            certGenProc =
            new ProcessExecutor()
                .command(
                    resolveOpenSSL(),
                    "x509",
                    "-req",
                    "-days",
                    Integer.toString(request.getRequestedValidityDays()),
                    "-in",
                    csrFile.toString(),
                    "-CA",
                    caCert.toString(),
                    "-CAkey",
                    caKey.toString(),
                    "-CAcreateserial",
                    "-passin",
                    OSSL_ARG_KEY_PW,
                    "-out",
                    outFile.toString(),
                    "-extfile",
                    extFile.toString()
                )
                .environment(OSSL_ENV_KEY_PW, caKeyPassphrase)
                .redirectOutput(Slf4jStream.of(openSSLLogger).asDebug())
                .redirectError(Slf4jStream.of(openSSLLogger).asError())
                .start();
            ProcessResult result = certGenProc.getFuture().get(30, SECONDS);
            // Check exit code
            if (result.getExitValue() != 0) {
                throw new CommandLineOperationException(
                    "Failed to sign certificate. Exit code: " + result.getExitValue()
                );
            }
        } catch (IOException | TimeoutException | ExecutionException | InterruptedException e) {
            throw new CommandLineOperationException("Failure running OpenSSL x509 command.", e);
        } finally {
            killIfActive(certGenProc);
        }
        return outFile;
    }

    public String getCertificateFingerprint(Path certificate) throws CommandLineOperationException {
        StartedProcess fingerprintProc = null;
        try {
            fingerprintProc =
            new ProcessExecutor()
                .command(resolveOpenSSL(), "x509", "-in", certificate.toString(), "-noout", "-fingerprint")
                .readOutput(true)
                .redirectError(Slf4jStream.of(openSSLLogger).asError())
                .start();
            var fingerprintResult = fingerprintProc.getFuture().get(30, SECONDS);
            String output = fingerprintResult.getOutput().getUTF8();

            if (fingerprintResult.getExitValue() != 0) {
                log.debug("Fingerprint command output:\n{}", output);
                throw new CommandLineOperationException(
                    "Failed to get fingerprint of certificate. Exit code: %d".formatted(
                            fingerprintResult.getExitValue()
                        )
                );
            }

            Matcher matcher = FINGERPRINT_EXTRACTOR.matcher(output);
            if (!matcher.find()) {
                log.debug(output);
                throw new CommandLineOperationException(
                    "Unexpected output of fingerprint command. (See log for more details)"
                );
            }
            String algorithm = matcher.group("algo");
            String fingerprint = matcher.group("finger");
            if (StringUtils.isBlank(algorithm) || StringUtils.isBlank(fingerprint)) {
                throw new CommandLineOperationException(
                    "Unexpected output of fingerprint command: %s %s".formatted(algorithm, fingerprint)
                );
            }
            return "%s;%s".formatted(algorithm, fingerprint);
        } catch (IOException | ExecutionException | TimeoutException | InterruptedException e) {
            throw new RuntimeException(e);
        } finally {
            killIfActive(fingerprintProc);
        }
    }

    public CertificateRequest getCertificateInfo(Path path) throws CommandLineOperationException {
        StartedProcess infoProc = null;
        try {
            infoProc =
            new ProcessExecutor()
                .command(
                    resolveOpenSSL(),
                    "x509",
                    "-in",
                    path.toString(),
                    "-noout",
                    "-dateopt",
                    "iso_8601",
                    "-fingerprint",
                    "-subject",
                    "-issuer",
                    "-serial",
                    "-dates",
                    "-alias",
                    "-email",
                    "-purpose",
                    "-ext",
                    "subjectAltName",
                    "-nameopt",
                    "sep_multiline",
                    "-nameopt",
                    "lname"
                )
                .readOutput(true)
                .redirectError(Slf4jStream.of(openSSLLogger).asError())
                .start();
            var infoResult = infoProc.getFuture().get(30, SECONDS);
            String output = infoResult.getOutput().getUTF8();
            if (infoResult.getExitValue() != 0) {
                log.debug("Certificate info command output:\n{}", output);
                throw new CommandLineOperationException(
                    "Failed to get info of path. Exit code: %d".formatted(infoResult.getExitValue())
                );
            }
            return getCertificateInfo(output.lines().toArray(String[]::new));
        } catch (IOException | ExecutionException | InterruptedException | TimeoutException e) {
            throw new RuntimeException(e);
        }
    }

    private String resolveOpenSSL() throws CommandLineOperationException {
        try {
            String path = executableResolver.getOpenSSLPath();
            if (!versionLogged.get()) {
                try {
                    StartedProcess versionProc = new ProcessExecutor()
                        .command(path, "version")
                        .readOutput(true)
                        .redirectError(Slf4jStream.of(openSSLLogger).asError())
                        .start();
                    var versionResult = versionProc.getFuture().get();
                    if (versionResult.getExitValue() != 0) {
                        throw new CommandLineOperationException(
                            "Failed to get OpenSSL version. Exit code: " + versionResult.getExitValue()
                        );
                    }

                    log.info("Using OpenSSL version: {}", versionResult.getOutput().getUTF8());
                    versionLogged.set(true);
                } catch (IOException | InterruptedException | ExecutionException e) {
                    throw new CommandLineOperationException("Failed to get OpenSSL version", e);
                }
            }
            return path;
        } catch (UnresolvableCLIDependency e) {
            throw new CommandLineOperationException(e);
        }
    }

    private CertificateRequest getCertificateInfo(String[] lines) {
        var builder = CertificateRequest.builder();
        boolean hasIssuer = false;

        for (int i = 0; i < lines.length; i++) {
            String line = lines[i];
            if (line.startsWith("subject=")) {
                CertificateSubjectBuilder subjectBuilder = CertificateSubject.builder();

                line = lines[++i];
                while (line.startsWith("    ")) {
                    subjectBuilder = readSubjectInfo(line, subjectBuilder);
                    line = lines[++i];
                }
                builder = builder.subject(subjectBuilder);
            } else if (line.startsWith("issuer=")) {
                hasIssuer = true;
                CertificateSubjectBuilder issuerBuilder = CertificateSubject.builder();

                line = lines[++i];
                while (line.startsWith("    ")) {
                    issuerBuilder = readSubjectInfo(line, issuerBuilder);
                    line = lines[++i];
                }

                builder = builder.issuer(issuerBuilder);
            } else if (line.startsWith("X509v3 Subject Alternative Name")) {
                String[] altNames = lines[++i].split(",");
                builder = builder.extension(CertificateRequestExtension.builder().alternativeNames(altNames));
            }
        }

        builder = builder.type(hasIssuer ? RequestType.NORMAL_CERTIFICATE : RequestType.STANDALONE_CERTIFICATE);
        return builder.build();
    }

    private CertificateSubjectBuilder readSubjectInfo(String line, CertificateSubjectBuilder builder) {
        String[] parts = line.split("=", 2);
        if (parts.length != 2) {
            return builder;
        }
        String key = parts[0];
        String value = parts[1];
        return switch (key.trim()) {
            case "countryName" -> builder.country(value);
            case "stateOrProvinceName" -> builder.state(value);
            case "localityName" -> builder.locality(value);
            case "organizationName" -> builder.organization(value);
            case "organizationalUnitName" -> builder.organizationalUnit(value);
            case "commonName" -> builder.commonName(value);
            case "emailAddress" -> builder.emailAddress(value);
            default -> throw new IllegalStateException("Unexpected subject key: %s in line: %s".formatted(key, line));
        };
    }
}