package model.ca;

import model.asn1.*;
import model.asn1.exceptions.ParseException;
import model.csr.*;
import model.pki.AlgorithmIdentifier;
import model.pki.SubjectPublicKeyInfo;
import model.pki.cert.*;
import model.pki.cert.Certificate;
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.RSAPrivateKeySpec;
import java.time.ZoneId;
import java.time.ZonedDateTime;
import java.util.*;
import java.util.stream.Stream;

/**
 * Holds a CA private key, its certificate, and signed / revoked list.
 */
public class CACertificate {
    /**
     * The key pair.
     */
    private KeyPair key;

    /**
     * The signed certificate.
     */
    private Certificate certificate;

    /**
     * Signed certificates.
     */
    private List<Certificate> signed;

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

    /**
     * Revoked certs.
     */
    private List<RevokedCertificate> revoked;

    /**
     * EFFECT: Init with a null key and null certificate, empty signed and revoked list, and serial at 1.
     */
    public CACertificate()  {
        this.key = null;
        this.certificate = null;
        this.serial = 1;
        this.signed = new ArrayList<>();
        this.revoked = new ArrayList<>();
    }

    /**
     * EFFECTS: Generate a new RSA2048 private key.
     * REQUIRES: getPublicKey() is null (i.e., no private key had been installed)
     */
    public void generateKey() throws NoSuchAlgorithmException {
        final KeyPairGenerator gen = KeyPairGenerator.getInstance("RSA");
        gen.initialize(2048);
        this.key = gen.generateKeyPair();
    }

    /**
     * EFFECT: Install the CA certificate.
     * MODIFIES: this
     * REQUIRES:
     *   - The new certificate must have the same algorithm and public key as getPublicKey(), except for testing purpose
     *   - It must be a v3 certificate
     *   - It must have basicConstraints { cA = TRUE }
     *   - It must contain key usage Digital Signature, Certificate Sign, CRL Sign
     *   - getCertificate() must be null (i.e., no certificate is installed yet).
     */
    public void installCertificate(Certificate certificate) {
        this.certificate = certificate;
    }

    /**
     * EFFECTS: Generate a CSR based on public key. It will have subject = CN=JCA.
     */
    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 = (RSAPublicKey) key.getPublic();
        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.
     * REQUIRES: The CA cert must not be installed.
     */
    public CertificationRequest signCSR()
            throws ParseException, NoSuchAlgorithmException, SignatureException, InvalidKeyException {
        final CertificationRequestInfo info = generateCSR();
        return new CertificationRequest(ASN1Object.TAG_SEQUENCE, null,
                info,
                getSigningAlgorithm(),
                new BitString(BitString.TAG, null, 0, signBytes(info.encodeDER())));
    }

    /**
     * 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.
     * 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);
        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.getPrivate());
        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
     */
    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.
     * 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);
    }

    /**
     * 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.
     * REQUIRES: The CA cert must be installed first.
     */
    public CertificateList signCRL()
            throws NoSuchAlgorithmException, SignatureException, InvalidKeyException {
        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]));
        return new CertificateList(ASN1Object.TAG_SEQUENCE, null,
                content,
                getSigningAlgorithm(),
                new BitString(BitString.TAG, null, 0,
                        signBytes(content.encodeValueDER())));
    }

    public Certificate getCertificate() {
        return certificate;
    }

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

    /**
     * EFFECTS: Get the public key, or null if no private key is installed.
     */
    public PublicKey getPublicKey() {
        if (key == null) {
            return null;
        }
        return key.getPublic();
    }

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

    public int getSerial() {
        return serial;
    }
}