package model.ca;

import annotations.Assoc;
import model.Event;
import model.EventLog;
import model.Observer;
import model.asn1.*;
import model.asn1.exceptions.InvalidCAException;
import model.asn1.exceptions.ParseException;
import model.asn1.parsing.BytesReader;
import model.csr.*;
import model.pki.AlgorithmIdentifier;
import model.pki.SubjectPublicKeyInfo;
import model.pki.cert.Certificate;
import model.pki.cert.Extension;
import model.pki.cert.TbsCertificate;
import model.pki.cert.Validity;
import model.pki.crl.CertificateList;
import model.pki.crl.CertificateListContent;
import model.pki.crl.RevokedCertificate;
import model.x501.AttributeTypeAndValue;
import model.x501.Name;
import model.x501.RelativeDistinguishedName;
import ui.Utils;

import java.math.BigInteger;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.time.ZoneId;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.util.*;
import java.util.stream.Stream;

/**
 * Holds a CA private key, its certificate, signed / revoked list, template list, and logs list. Data can be observed.
 */
public class CertificationAuthority {
    public static final int SERIAL_DEFAULT = 1;

    /**
     * The RSA2048 private key.
     */
    @Assoc(partOf = true, lowerBond = 0)
    private RSAPrivateKey key;

    /**
     * The public key.
     */
    @Assoc(partOf = true, lowerBond = 0)
    private RSAPublicKey publicKey;

    /**
     * The signed certificate.
     */
    @Assoc(partOf = true, lowerBond = 0)
    private Certificate certificate;

    /**
     * Signed certificates.
     */
    @Assoc(lowerBond = 0)
    private final List<Certificate> signed;

    /**
     * The next serial number.
     */
    private int serial;

    /**
     * Revoked certs.
     */
    @Assoc(lowerBond = 0)
    private final List<RevokedCertificate> revoked;

    /**
     * Certificate templates.
     */
    @Assoc(lowerBond = 0)
    private final List<Template> templates;

    /**
     * Audit logs.
     */
    @Assoc(lowerBond = 0)
    private final List<AuditLogEntry> logs;

    /**
     * Current operator.
     */
    private final String user;

    /**
     * Data observers.
     */
    @Assoc(lowerBond = 0)
    private final List<Observer> observers;

    /**
     * EFFECT: Init with the given parameters, user "yuuta", and no observers.
     * Throws {@link NoSuchAlgorithmException} if the key is specified but RSA is not supported.
     * Throws {@link InvalidKeySpecException} if the key specified is invalid.
     * Throws {@link InvalidCAException} or {@link ParseException} if the CA specified is invalid.
     * REQUIRES: n / p / e must be either all null or all non-null containing RSA2048 module and exponents.
     * If certificate is non-null, n / p / e must be non-null.
     */
    public CertificationAuthority(BigInteger n, BigInteger p, BigInteger e,
                                  Certificate certificate,
                                  List<Certificate> signed,
                                  int serial,
                                  List<RevokedCertificate> revoked,
                                  List<Template> templates,
                                  List<AuditLogEntry> logs)
            throws NoSuchAlgorithmException, InvalidKeySpecException, InvalidCAException, ParseException {
        if (n != null) {
            setKey(n, p, e);
        }
        if (certificate != null) {
            validateCertificate(certificate);
        }
        this.certificate = certificate;
        this.signed = new ArrayList<>(signed);
        this.serial = serial;
        this.revoked = new ArrayList<>(revoked);
        this.templates = new ArrayList<>(templates);
        this.logs = new ArrayList<>(logs);
        this.user = "yuuta";
        this.observers = new ArrayList<>();
    }

    /**
     * EFFECT: Init with a null key and null certificate, empty signed, revoked template, and log list,
     * serial at SERIAL_DEFAULT, user "yuuta", and no observers.
     */
    public CertificationAuthority() {
        this.key = null;
        this.publicKey = null;
        this.certificate = null;
        this.serial = SERIAL_DEFAULT;
        this.signed = new ArrayList<>();
        this.revoked = new ArrayList<>();
        this.templates = new ArrayList<>();
        this.logs = new ArrayList<>();
        this.user = "yuuta";
        this.observers = new ArrayList<>();
    }

    /**
     * EFFECTS: Generate a new RSA2048 private key. This action will be logged.
     *          Observers will be notified for (RSAPublicKey.class, DIRECTION_CHANGE, INDEX_NOT_IN_LIST).
     *          Observers will be notified for (AuditLogEntry.class, DIRECTION_ADD, i).
     * REQUIRES: getPublicKey() is null (i.e., no private key had been installed)
     * MODIFIES: this
     */
    public void generateKey() throws NoSuchAlgorithmException {
        final KeyPairGenerator gen = KeyPairGenerator.getInstance("RSA");
        gen.initialize(2048);
        final KeyPair pair = gen.generateKeyPair();
        this.key = (RSAPrivateKey) pair.getPrivate();
        this.publicKey = (RSAPublicKey) pair.getPublic();
        notif(getPublicKey(), Observer.DIRECTION_CHANGE, Observer.INDEX_NOT_IN_LIST);
        log("Generated CA private key.");
    }

    /**
     * EFFECTS: Load the RSA private and public exponents.
     * Throws {@link NoSuchAlgorithmException} if RSA is not available on the platform.
     * Throws {@link InvalidKeySpecException} if the input is invalid.
     * REQUIRES: getPublicKey() is null (i.e., no private key had been installed)
     * MODIFIES: this
     */
    private void setKey(BigInteger n, BigInteger p, BigInteger e)
            throws NoSuchAlgorithmException, InvalidKeySpecException {
        this.key = (RSAPrivateKey) KeyFactory.getInstance("RSA").generatePrivate(new RSAPrivateKeySpec(n, p));
        this.publicKey =
                (RSAPublicKey) KeyFactory.getInstance("RSA").generatePublic(new RSAPublicKeySpec(n, e));
    }

    /**
     * EFFECTS: Load the RSA private and public exponents. This action will be logged.
     *          Observers will be notified for (RSAPublicKey.class, DIRECTION_CHANGE, INDEX_NOT_IN_LIST).
     *          Observers will be notified for (AuditLogEntry.class, DIRECTION_ADD, i).
     * Throws {@link NoSuchAlgorithmException} if RSA is not available on the platform.
     * Throws {@link InvalidKeySpecException} if the input is invalid.
     * REQUIRES: getPublicKey() is null (i.e., no private key had been installed)
     * MODIFIES: this
     */
    public void loadKey(BigInteger n, BigInteger p, BigInteger e)
            throws NoSuchAlgorithmException, InvalidKeySpecException {
        setKey(n, p, e);
        notif(getPublicKey(), Observer.DIRECTION_CHANGE, Observer.INDEX_NOT_IN_LIST);
        log("Installed CA private key.");
    }

    /**
     * EFFECTS: Throw {@link InvalidCAException} if the incoming cert is not v3.
     */
    private void validateCACertificateVersion(Certificate cert) throws InvalidCAException {
        if (cert.getCertificate().getVersion() == null
                || cert.getCertificate().getVersion().getLong() != TbsCertificate.VERSION_V3) {
            throw new InvalidCAException("The input certificate must be V3");
        }
    }

    /**
     * EFFECTS: Throw {@link InvalidCAException} if the incoming cert does not have the matching public key.
     */
    private void validateCACertificatePublicKey(Certificate cert) throws InvalidCAException {
        final SubjectPublicKeyInfo expectedPKInfo = getCAPublicKeyInfo();
        if (!Arrays.equals(cert.getCertificate().getSubjectPublicKeyInfo().getAlgorithm().getType().getInts(),
                expectedPKInfo.getAlgorithm().getType().getInts())
                || !Arrays.equals(cert.getCertificate().getSubjectPublicKeyInfo().getSubjectPublicKey().getVal(),
                expectedPKInfo.getSubjectPublicKey().getVal())) {
            throw new InvalidCAException("The input certificate does not have the corresponding public key");
        }
    }

    /**
     * EFFECTS: Throw {@link InvalidCAException} if the incoming cert does not have cA = true in its basicConstraints.
     */
    private void validateCACertificateBasicConstraints(Certificate cert) throws InvalidCAException, ParseException {
        final Extension basicConstraints = cert.getCertificate().getExtension(ObjectIdentifier.OID_BASIC_CONSTRAINTS);
        if (basicConstraints == null) {
            throw new InvalidCAException("The certificate does not have a valid basicConstraints extension.");
        }
        final ASN1Object basicConstraintsValue =
                new ASN1Object(new BytesReader(basicConstraints.getExtnValue().getBytes()), false);
        if (basicConstraintsValue.getLength() <= 0) {
            throw new InvalidCAException("The certificate does not have a valid basicConstraints extension.");
        }
        final ASN1Object bool =
                ASN1Object.parse(new BytesReader(basicConstraintsValue.encodeValueDER()), false);
        if (!((Bool) bool).getValue()) {
            throw new InvalidCAException("The certificate does not have a valid basicConstraints extension.");
        }
    }

    /**
     * EFFECTS: Throw {@link InvalidCAException} if the incoming cert does not have valid key usages.
     */
    private void validateCACertificateKeyUsage(Certificate cert) throws InvalidCAException, ParseException {
        final Extension keyUsage = cert.getCertificate().getExtension(ObjectIdentifier.OID_KEY_USAGE);
        if (keyUsage == null) {
            throw new InvalidCAException("The certificate does not have a valid keyUsage extension.");
        }
        final ASN1Object keyUsageValue =
                ASN1Object.parse(new BytesReader(keyUsage.getExtnValue().getBytes()), false);
        final BitSet bitSet = BitSet.valueOf(Utils.byteToByte(((BitString) keyUsageValue).getVal()));
        if (!bitSet.get(7) || !bitSet.get(2) || !bitSet.get(1)) {
            throw new InvalidCAException("The certificate does not have a valid keyUsage extension.");
        }
    }

    /**
     * EFFECT: Validate the CA certificate. Throws {@link InvalidCAException} if any of the
     * following are violated:
     * - It must be a v3 certificate
     * - The new certificate must have the same algorithm and public key as getPublicKey()
     * - It must have basicConstraints { cA = TRUE }
     * - It must contain key usage Digital Signature, Certificate Sign, CRL Sign
     * Throws {@link ParseException} if the cert has invalid extension values.
     */
    private void validateCertificate(Certificate certificate) throws InvalidCAException, ParseException {
        validateCACertificateVersion(certificate);
        validateCACertificatePublicKey(certificate);
        validateCACertificateBasicConstraints(certificate);
        validateCACertificateKeyUsage(certificate);
    }

    /**
     * EFFECT: Install the CA certificate. Throws {@link InvalidCAException} if any of the
     * following are violated:
     * - It must be a v3 certificate
     * - The new certificate must have the same algorithm and public key as getPublicKey()
     * - It must have basicConstraints { cA = TRUE }
     * - It must contain key usage Digital Signature, Certificate Sign, CRL Sign
     *          Observers will be notified for (Certificate.class, DIRECTION_CHANGE, INDEX_NOT_IN_LIST).
     *          Observers will be notified for (AuditLogEntry.class, DIRECTION_ADD, i).
     * Throws {@link ParseException} if the cert has invalid extension values.
     * This action will be logged.
     * REQUIRES:
     * - getCertificate() must be null (i.e., no certificate is installed yet).
     * MODIFIES: this
     */
    public void installCertificate(Certificate certificate) throws InvalidCAException, ParseException {
        validateCertificate(certificate);
        this.certificate = certificate;
        notif(certificate, Observer.DIRECTION_CHANGE, Observer.INDEX_NOT_IN_LIST);
        log("CA certificate is installed.");
    }

    /**
     * EFFECTS: Generate a CSR based on public key. It will have subject = CN=JCA.
     * REQUIRES:
     * - getCertificate() must be null (i.e., no certificate is installed yet).
     */
    private CertificationRequestInfo generateCSR() throws ParseException {
        return new CertificationRequestInfo(ASN1Object.TAG_SEQUENCE, null,
                new Int(Int.TAG, null, CertificationRequestInfo.VERSION_V1),
                new Name(ASN1Object.TAG_SEQUENCE, null, new RelativeDistinguishedName[]{
                        new RelativeDistinguishedName(ASN1Object.TAG_SET, null, new AttributeTypeAndValue[]{
                                new AttributeTypeAndValue(ASN1Object.TAG_SEQUENCE, null,
                                        new ObjectIdentifier(ObjectIdentifier.TAG, null,
                                                ObjectIdentifier.OID_CN),
                                        new PrintableString(PrintableString.TAG, null, "JCA"))
                        })
                }),
                getCAPublicKeyInfo(),
                new Attributes(new Tag(TagClass.CONTEXT_SPECIFIC, true, 0), // IMPLICIT
                        null,
                        new Attribute[]{
                                new Attribute(ASN1Object.TAG_SEQUENCE, null,
                                        new ObjectIdentifier(ObjectIdentifier.TAG, null,
                                                new Integer[]{1, 3, 6, 1, 4, 1, 311, 13, 2, 3}),
                                        new Values(ASN1Object.TAG_SET, null,
                                                new ASN1Object[]{
                                                        new IA5String(IA5String.TAG, null,
                                                                "10.0.20348.2")
                                                }))}));
    }

    private Byte[] getPubKeyBitStream() {
        final RSAPublicKey pub = getPublicKey();
        final BigInteger exponent = pub.getPublicExponent();
        byte[] modules = pub.getModulus().toByteArray();
        final Int asn1Exponent = new Int(Int.TAG, null, exponent);
        // Use OctetString to avoid leading zero issues.
        final ASN1Object asn1Modules = new OctetString(Int.TAG, null, Utils.byteToByte(modules));
        final Byte[] asn1ExponentDER = asn1Exponent.encodeDER();
        final Byte[] asn1ModulesDER = asn1Modules.encodeDER();
        return Stream.of(Arrays.asList(ASN1Object.TAG_SEQUENCE.encodeDER()),
                        Arrays.asList(new ASN1Length(asn1ModulesDER.length + asn1ExponentDER.length).encodeDER()),
                        Arrays.asList(asn1ModulesDER),
                        Arrays.asList(asn1ExponentDER))
                .flatMap(Collection::stream)
                .toArray(Byte[]::new);
    }

    /**
     * EFFECTS: Encode the RSA public key into SubjectPubicKeyInfo format (BIT STRING -> SEQUENCE -> { INT INT }).
     */
    public SubjectPublicKeyInfo getCAPublicKeyInfo() {
        return new SubjectPublicKeyInfo(ASN1Object.TAG_SEQUENCE, null,
                new AlgorithmIdentifier(ASN1Object.TAG_SEQUENCE, null,
                        new ObjectIdentifier(ObjectIdentifier.TAG, null,
                                ObjectIdentifier.OID_RSA_ENCRYPTION),
                        new Null(Null.TAG, null)),
                new BitString(BitString.TAG, null, 0, getPubKeyBitStream()));
    }

    /**
     * EFFECT: Generate CSR and sign it, so the CA can request itself a certificate.
     *          Observers will be notified for (AuditLogEntry.class, DIRECTION_ADD, i).
     * REQUIRES: The CA cert must not be installed.
     * MODIFIES: this (This action will be logged)
     */
    public CertificationRequest signCSR()
            throws ParseException, NoSuchAlgorithmException, SignatureException, InvalidKeyException {
        final CertificationRequestInfo info = generateCSR();
        final CertificationRequest csr = new CertificationRequest(ASN1Object.TAG_SEQUENCE, null,
                info,
                getSigningAlgorithm(),
                new BitString(BitString.TAG, null, 0, signBytes(info.encodeDER())));
        log("Signed CA csr");
        return csr;
    }

    /**
     * EFFECT: Return SHA256withRSA.
     */
    private AlgorithmIdentifier getSigningAlgorithm() {
        return new AlgorithmIdentifier(ASN1Object.TAG_SEQUENCE, null,
                new ObjectIdentifier(ObjectIdentifier.TAG, null,
                        ObjectIdentifier.OID_SHA256_WITH_RSA_ENCRYPTION),
                new Null(Null.TAG, null));
    }

    /**
     * EFFECTS: Sign the CSR based on the template.
     *          Observers will be notified for (Certificate.class, DIRECTION_ADD, i).
     *          Observers will be notified for (AuditLogEntry.class, DIRECTION_ADD, i).
     * REQUIRES: The CA cert must be installed first, req must have a subject, template must be enabled.
     * MODIFIES: this
     */
    public Certificate signCert(CertificationRequestInfo req, Template template)
            throws NoSuchAlgorithmException, SignatureException, InvalidKeyException {
        final TbsCertificate newCert = generateCert(req, template);
        final Certificate cert = new Certificate(ASN1Object.TAG_SEQUENCE, null,
                newCert,
                getSigningAlgorithm(),
                new BitString(BitString.TAG, null, 0,
                        signBytes(newCert.encodeValueDER())));
        this.signed.add(cert);
        notif(cert, Observer.DIRECTION_ADD, this.signed.size() - 1);
        log("Signed a cert with serial number " + cert.getCertificate().getSerialNumber().getLong());
        return cert;
    }

    /**
     * EFFECTS: Hash the input message with SHA256 and sign it with RSA and get the signature.
     */
    private Byte[] signBytes(Byte[] message)
            throws NoSuchAlgorithmException, SignatureException, InvalidKeyException {
        final Signature signature = Signature.getInstance("SHA256withRSA");
        signature.initSign(key);
        signature.update(Utils.byteToByte(message));
        return Utils.byteToByte(signature.sign());
    }

    /**
     * EFFECTS: Apply the template.
     * For the new certificate:
     * - Issuer will be set to CA#getCertificate()#getSubject()
     * - The template will be applied (subject, validity, cdp)
     * - A serial number will be generated
     * MODIFIES: this
     */
    private TbsCertificate generateCert(CertificationRequestInfo req, Template template) {
        final ZonedDateTime now = ZonedDateTime.now(ZoneId.of("UTC"));
        return new TbsCertificate(ASN1Object.TAG_SEQUENCE, null,
                new Int(Int.TAG, new Tag(TagClass.CONTEXT_SPECIFIC, true, 0),
                        TbsCertificate.VERSION_V3),
                new Int(Int.TAG, null, serial++),
                getSigningAlgorithm(),
                certificate.getCertificate().getSubject(),
                new Validity(ASN1Object.TAG_SEQUENCE, null,
                        new GeneralizedTime(GeneralizedTime.TAG, null, now),
                        new UtcTime(UtcTime.TAG, null,
                                now.plusDays(template.getValidity()))),
                template.getSubject() == null ? req.getSubject() :
                        template.getSubject(),
                req.getSubjectPKInfo(),
                null);
    }

    /**
     * EFFECTS: Add the revocation info to revoked list. This action will be logged.
     *          Observers will be notified for (RevokedCertificate.class, DIRECTION_ADD, i).
     *          Observers will be notified for (AuditLogEntry.class, DIRECTION_ADD, i).
     * REQUIRES: revoked should have the serial of an issued certificate; its date should be current.
     * MODIFIES: this
     */
    public void revoke(RevokedCertificate rev) {
        revoked.add(rev);
        notif(rev, Observer.DIRECTION_ADD, revoked.size() - 1);
        log("Certificate " + rev.getSerialNumber().getLong() + " is revoked with reason " + rev.getReason()
                + " at " + rev.getRevocationDate().getTimestamp().format(DateTimeFormatter.ISO_LOCAL_DATE_TIME));
    }

    /**
     * EFFECTS: Generate and sign the CRL, based on getRevokedCerts(). The CSR will have current time as thisUpdate with
     * no nextUptime, and it will have issuer same as the CA's subject.
     *          Observers will be notified for (AuditLogEntry.class, DIRECTION_ADD, i).
     * REQUIRES: The CA cert must be installed first.
     * MODIFIES: this (This action will be logged)
     */
    public CertificateList signCRL()
            throws NoSuchAlgorithmException, SignatureException, InvalidKeyException {
        final CertificateListContent content = new CertificateListContent(ASN1Object.TAG_SEQUENCE, null,
                certificate.getCertificate().getSubject(),
                getSigningAlgorithm(),
                new GeneralizedTime(GeneralizedTime.TAG, null, ZonedDateTime.now(ZoneId.of("UTC"))),
                null,
                revoked.toArray(new RevokedCertificate[0]));
        final CertificateList crl = new CertificateList(ASN1Object.TAG_SEQUENCE, null,
                content,
                getSigningAlgorithm(),
                new BitString(BitString.TAG, null, 0,
                        signBytes(content.encodeValueDER())));
        log("Signed CRL with " + revoked.size() + " revoked certs.");
        return crl;
    }

    /**
     * EFFECTS: Log the action with the current date and user.
     *          Observers will be notified for (AuditLogEntry.class, DIRECTION_ADD, i).
     * MODIFIES: this
     */
    private void log(String message) {
        final AuditLogEntry i = new AuditLogEntry(user, ZonedDateTime.now(), message);
        this.logs.add(i);
        notif(i, Observer.DIRECTION_ADD, logs.size() - 1);
        EventLog.getInstance().logEvent(new Event(user + ": " + message));
    }

    /**
     * EFFECTS: Register the given observer, so it will be called upon changes.
     * MODIFIES: this
     */
    public void registerObserver(final Observer<?> observer) {
        this.observers.add(observer);
    }

    /**
     * EFFECTS: Notify the observers.
     * REQUIRES: direction must be valid Observer constants, i must be either >= 0 or Observer.INDEX_NOT_IN_LIST.
     */
    private void notif(Object o, int direction, int i) {
        observers.forEach(e -> e.accept(o, direction, i));
    }

    /**
     * EFFECTS: Find the template based on name, or null if not found.
     */
    public Template findTemplate(String name, boolean requireEnabled) {
        Optional<Template> opt = templates.stream().filter(temp -> {
            if (requireEnabled && !temp.isEnabled()) {
                return false;
            }
            return temp.getName().equals(name);
        }).findFirst();
        return opt.orElse(null);
    }

    /**
     * EFFECTS: Install the new template. This action will be logged.
     *          Observers will be notified for (Template.class, DIRECTION_ADD, i).
     *          Observers will be notified for (AuditLogEntry.class, DIRECTION_ADD, i).
     * REQUIRES: findTemplate(template.getName(), false) == null
     * MODIFIES: this
     */
    public void addTemplate(Template template) {
        this.templates.add(template);
        notif(template, Observer.DIRECTION_ADD, templates.size() - 1);
        log("Added a new template: " + template.getName());
    }

    /**
     * EFFECTS: Set the given template to enabled / disabled, order will be kept. This action will be logged.
     *          Observers will be notified for (Template.class, DIRECTION_CHANGE, i).
     *          Observers will be notified for (AuditLogEntry.class, DIRECTION_ADD, i).
     * REQUIRES: the template is valid (findTemplate does not return null)
     * MODIFIES: this
     */
    public void setTemplateEnable(Template template, boolean enable) {
        final Template t = findTemplate(template.getName(), false);
        int i = templates.indexOf(t);
        templates.set(i, new Template(t.getName(), enable, t.getSubject(), t.getValidity()));
        notif(template, Observer.DIRECTION_CHANGE, i);
        log("Template " + template.getName() + " has been " + (enable ? "enabled" : "disabled"));
    }

    /**
     * EFFECTS: Remove the given template. This action will be logged.
     *          Observers will be notified for (Template.class, DIRECTION_REMOVE, i).
     *          Observers will be notified for (AuditLogEntry.class, DIRECTION_ADD, i).
     * REQUIRES: the template is valid (findTemplate does not return null)
     * MODIFIES: this
     */
    public void removeTemplate(Template template) {
        int i = templates.indexOf(template);
        templates.remove(findTemplate(template.getName(), false));
        notif(template, Observer.DIRECTION_REMOVE, i);
        log("Template " + template.getName() + " is removed");
    }

    // Getters

    public Certificate getCertificate() {
        return certificate;
    }

    public List<Certificate> getSigned() {
        return signed;
    }

    public List<RevokedCertificate> getRevoked() {
        return revoked;
    }

    public int getSerial() {
        return serial;
    }

    public List<Template> getTemplates() {
        return templates;
    }

    public String getUser() {
        return user;
    }

    public List<AuditLogEntry> getLogs() {
        return logs;
    }

    public RSAPublicKey getPublicKey() {
        return publicKey;
    }

    public RSAPrivateKey getKey() {
        return key;
    }
}