package be.iminds.ilabt.jfed.lowlevel;


import be.iminds.ilabt.jfed.fedmon.webapi.service.json.*;
import be.iminds.ilabt.jfed.util.*;
import org.apache.xml.security.exceptions.XMLSecurityException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.w3c.dom.*;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;

import javax.annotation.Nonnull;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.stream.XMLInputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamReader;
import javax.xml.transform.TransformerException;
import javax.xml.transform.stream.StreamSource;
import java.io.IOException;
import java.io.StringReader;
import java.security.Key;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.cert.X509Certificate;
import java.security.interfaces.RSAPublicKey;
import java.text.ParseException;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Objects;

public class AbacCredential extends AnyCredential {
    private static final Logger LOG = LoggerFactory.getLogger(AbacCredential.class);

    protected final String name;
    @Nonnull
    protected final String credentialXml;

    protected final String type;
    protected final String version;

    //whether process() ran
    private boolean processed = false;
    //fields derived using process():
    private boolean speaksFor;
    private GeniUrn spokenForUrn;
    private PublicKey spokenForPubKey;
    private List<String> signerCerts;

    protected String headKeyId;
    protected String tailKeyId;
    protected String headRole;

    protected String expiresText;

    protected AbacCredential(String name, @Nonnull String credentialXml, String type, String version) {
        super(name, credentialXml, type, version);

        this.name = name;
        this.credentialXml = credentialXml;
        this.type = type;
        this.version = version;

        assert type.equalsIgnoreCase("abac") || type.equalsIgnoreCase("geni_abac") :
                "Created SfaCredential not of type sfa, but of type=\"" + type + "\" version=\"" + version + "\"";

    }

    @Override
    public String getExpires() {
        return expiresText;
    }

    @Override
    public Date getExpiresDate() {
        process();
        try {
            return expiresText == null ? null : RFC3339Util.iso8601StringToDate(expiresText);
        } catch (ParseException e) {
            throw new RuntimeException("Invalid date in ABAC credential expires: \""+expiresText+"\"", e);
        }
    }

    @Nonnull
    public String getCredentialXml() {
        return credentialXml;
    }

    @Override
    public boolean isSpeaksFor() {
        process();
        return speaksFor;
    }

    public GeniUrn getSpokenForUrn() {
        process();
        return spokenForUrn;
    }

    public PublicKey getSpokenForPubKey() {
        process();
        return spokenForPubKey;
    }

    public String getHeadKeyId() {
        return headKeyId;
    }

    public String getTailKeyId() {
        return tailKeyId;
    }

    public String getHeadRole() {
        return headRole;
    }

    private void process() {
        if (processed) return;
        processed = true;

        speaksFor = false;
        signerCerts = new ArrayList<>();
        spokenForUrn = null;
        spokenForPubKey = null;

        /* speaksFor can be detected with:
                 <signed-credential xsi:noNamespaceSchemaLocation="http://www.protogeni.net/resources/credential/credential.xsd" >
                  <credential xml:id="_0">
                	<abac>
                		<head>
                            <role>speaks_for_*</role>

           signer cert is in:
              <signed-credential xsi:noNamespaceSchemaLocation="http://www.protogeni.net/resources/credential/credential.xsd">
              <signatures>
                  <Signature xmlns="http://www.w3.org/2000/09/xmldsig#">
                    <KeyInfo>
                      <X509Data>
                        <X509Certificate>
                * */

        XMLStreamReader streamReader = null;
        try {
            XMLInputFactory xif = XMLInputFactory.newFactory();
            StreamSource xml = new StreamSource(new StringReader(credentialXml));
            streamReader = xif.createXMLStreamReader(xml);

            boolean withinSignedCredential = false;

            boolean withinCredential = false;
            boolean withinAbac = false;
            boolean withinExpires = false;
            boolean withinHead = false;
            boolean withinTail = false;
            boolean withinKeyId = false;
            boolean withinRole = false;

            boolean withinSignatures = false;
            boolean withinSignature = false;
            boolean withinKeyInfo = false;
            boolean withinX509Data = false;
            boolean withinX509Certificate = false;

            String x509certText = "";
            while (streamReader.hasNext()) {
                streamReader.next();
                switch (streamReader.getEventType()) {
                    case XMLStreamReader.START_ELEMENT: {
                        if (Objects.equals(streamReader.getName().getLocalPart(), "signed-credential"))
                            withinSignedCredential = true;

                        if (withinSignedCredential && Objects.equals(streamReader.getName().getLocalPart(), "credential"))
                            withinCredential = true;
                        if (withinCredential && Objects.equals(streamReader.getName().getLocalPart(), "expires"))
                            withinExpires = true;
                        if (withinCredential && Objects.equals(streamReader.getName().getLocalPart(), "abac"))
                            withinAbac = true;
                        if (withinAbac && Objects.equals(streamReader.getName().getLocalPart(), "head"))
                            withinHead = true;
                        if (withinAbac && Objects.equals(streamReader.getName().getLocalPart(), "tail"))
                            withinTail = true;
                        if (withinHead && Objects.equals(streamReader.getName().getLocalPart(), "role"))
                            withinRole = true;
                        if ((withinHead || withinTail) && Objects.equals(streamReader.getName().getLocalPart(), "keyid"))
                            withinKeyId = true;

                        if (withinSignedCredential && Objects.equals(streamReader.getName().getLocalPart(), "signatures"))
                            withinSignatures = true;
                        if (withinSignatures && Objects.equals(streamReader.getName().getLocalPart(), "Signature"))
                            withinSignature = true;
                        if (withinSignature && Objects.equals(streamReader.getName().getLocalPart(), "KeyInfo"))
                            withinKeyInfo = true;
                        if (withinKeyInfo && Objects.equals(streamReader.getName().getLocalPart(), "X509Data"))
                            withinX509Data = true;
                        if (withinX509Data && Objects.equals(streamReader.getName().getLocalPart(), "X509Certificate")) {
                            withinX509Certificate = true;
                            x509certText = "";
                        }
                        break;
                    }
                    case XMLStreamReader.END_ELEMENT: {
                        if (withinSignedCredential && Objects.equals(streamReader.getName().getLocalPart(), "signed-credential"))
                            withinSignedCredential = false;

                        if (withinSignatures && Objects.equals(streamReader.getName().getLocalPart(), "credential"))
                            withinCredential = false;
                        if (withinCredential && Objects.equals(streamReader.getName().getLocalPart(), "abac"))
                            withinAbac = false;
                        if (withinCredential && Objects.equals(streamReader.getName().getLocalPart(), "expires"))
                            withinExpires = false;
                        if (withinAbac && Objects.equals(streamReader.getName().getLocalPart(), "head"))
                            withinHead = false;
                        if (withinAbac && Objects.equals(streamReader.getName().getLocalPart(), "tail"))
                            withinTail = false;
                        if (withinHead && Objects.equals(streamReader.getName().getLocalPart(), "role"))
                            withinRole = false;
                        if ((withinHead || withinTail) && Objects.equals(streamReader.getName().getLocalPart(), "keyid"))
                            withinKeyId = false;

                        if (withinSignedCredential && Objects.equals(streamReader.getName().getLocalPart(), "signatures"))
                            withinSignatures = false;
                        if (withinSignatures && Objects.equals(streamReader.getName().getLocalPart(), "Signature"))
                            withinSignature = false;
                        if (withinSignature && Objects.equals(streamReader.getName().getLocalPart(), "KeyInfo"))
                            withinKeyInfo = false;
                        if (withinKeyInfo && Objects.equals(streamReader.getName().getLocalPart(), "X509Data"))
                            withinX509Data = false;
                        if (withinX509Data && Objects.equals(streamReader.getName().getLocalPart(), "X509Certificate")) {
                            if (!x509certText.trim().isEmpty()) {
                                if (x509certText != null) {
                                    if (!x509certText.startsWith("-----BEGIN CERTIFICATE-----"))
                                        x509certText = "-----BEGIN CERTIFICATE-----\n" + x509certText + "\n-----END CERTIFICATE-----\n";
                                    signerCerts.add(x509certText.trim());
                                }
                            }

                            withinX509Certificate = false;
                            x509certText = "";
                        }
                        break;
                    }
                    case XMLStreamReader.CHARACTERS: {
                        if (withinRole) {
                            String roleText = streamReader.getText();
                            if (roleText != null) {
                                headRole = roleText.trim();
                                speaksFor = headRole.toLowerCase().startsWith("speaks_for_");
                            }
                        }
                        if (withinExpires) {
                            expiresText = streamReader.getText();
                        }
                        if (withinHead && withinKeyId) {
                            headKeyId = streamReader.getText();
                        }
                        if (withinTail && withinKeyId) {
                            tailKeyId = streamReader.getText();
                        }
                        if (withinX509Certificate) {
                            x509certText += streamReader.getText();
                        }
                        break;
                    }
                    default:
                        break;
                }
            }
        } catch (XMLStreamException e) {
            LOG.error("Exception while parsing ABAC credential", e);
            speaksFor = false;
        } finally {
            if (streamReader != null)
                try {
                    streamReader.close();
                } catch (XMLStreamException e) {
                    LOG.error("Exception closing streamReader is ignored", e);
                }
        }

        if (speaksFor) {
            List<GeniUrn> subjectUrns = new ArrayList<>();
            List<GeniUrn> issuerUrns = new ArrayList<>();
            for (String signerCert : signerCerts) {
                X509Certificate cert = KeyUtil.pemToX509Certificate(signerCert);
                if (cert == null) {
                    LOG.warn("Failed to convert signerCert to X509 certificate: \"\"\"\n" + signerCert + "\n\"\"\"");
                } else {
                    subjectUrns.addAll(KeyUtil.findUrnsInCertAltNames(cert, KeyUtil.AltNamesSource.SUBJECT_ALT_NAMES, false));
                    issuerUrns.addAll(KeyUtil.findUrnsInCertAltNames(cert, KeyUtil.AltNamesSource.ISSUES_ALT_NAMES, false));
                }
            }

            //simple way to handle most certificate chains
            List<GeniUrn> possibleUserUrns = new ArrayList<>(subjectUrns);
            possibleUserUrns.removeAll(issuerUrns);

            if (!possibleUserUrns.isEmpty()) {
                if (possibleUserUrns.size() > 1) {
                    System.err.println("Warning: more than 1 possible user urn in speaksFor credential signer subject alt names: " + possibleUserUrns);
                    LOG.warn("Warning: more than 1 possible user urn in speaksFor credential signer subject alt names: " + possibleUserUrns);
                }
                spokenForUrn = possibleUserUrns.get(0);

                //and finally, search for the publickey matching this URN
                for (String signerCert : signerCerts) {
                    X509Certificate cert = KeyUtil.pemToX509Certificate(signerCert);
                    if (cert == null) {
                        LOG.warn("Failed to convert signerCert to X509 certificate: \"\"\"\n" + signerCert + "\n\"\"\"");
                    } else {
                        spokenForPubKey = cert.getPublicKey();

                        assert speaksFor;
                        String signerKeyId = generateKeyId(cert);
                        String spokenForKeyId = headRole.substring("speaks_for_".length());
                        if (!Objects.equals(signerKeyId, spokenForKeyId)) {
                            throw new RuntimeException("AbacCredential credential is speaksfor, but signer does not match spoken for: signer="+signerKeyId+" spokenForKeyId="+spokenForKeyId);
                        }
                    }
                }

            }

        }
    }

    @Override
    public boolean check(TestbedInfoSource testbedInfoSource) throws CredentialException {
        process();

        JFedTrustStore jFedTrustStore = new JFedTrustStore();

        for (String signerCert : signerCerts) {
            X509Certificate cert = KeyUtil.pemToX509Certificate(signerCert);
            if(cert == null){
                LOG.warn("Could not process signerCert {}..., skipping", signerCert.substring(0, Math.min(signerCert.length(),10)));
                continue;
            }

            List<GeniUrn> urns = KeyUtil.findUrnsInCertAltNames(cert, KeyUtil.AltNamesSource.SUBJECT_ALT_NAMES, false);
            urns.addAll(KeyUtil.findUrnsInCertAltNames(cert, KeyUtil.AltNamesSource.ISSUES_ALT_NAMES, false));
            for (GeniUrn urn : urns) {
                Server auth = testbedInfoSource.getFromAnyUrn(urn, TestbedInfoSource.SubAuthMatchAllowed.ALLOW_TOPLEVEL, TestbedInfoSource.SubAuthMatchPreference.PREFER_EXACT_SUBAUTHORITY);
                jFedTrustStore.addAuthorityCert(auth);
            }
        }

        return check(jFedTrustStore.getTrustStore());
    }

    /**
     * @param geniUserCert an X509 certificate, belonging to a geni user, authority or tool
     * @return the best found description of the subject of the cert. Typically, the (user) URN that is found in the alternative names.
     */
    public static String getGeniSubjectName(X509Certificate geniUserCert) {
        if (geniUserCert == null)
            return "null";
        //TODO
        return geniUserCert.getSubjectDN().getName();
    }

    /**
     * @return the "SHA-1 hash of the principal's public key"
     *
     * @link http://groups.geni.net/geni/wiki/TIEDABACCredential
     * */
    public static String generateKeyId(X509Certificate cert) {
        PublicKey pubKey = cert.getPublicKey();

        try {
            MessageDigest crypt;
            crypt = MessageDigest.getInstance("SHA-1");
            crypt.reset();
            RSAPublicKey rsaPubKey = (RSAPublicKey) pubKey;
            byte[] keyDataBytes = KeyUtil.getPublicKeyPKCS1(rsaPubKey);
//            byte[] keyDataBytes = pubKey.getEncoded(); //not this! These contain algorithm as well, which ABAC does not use, because it uses openssl's X509_pubkey_digest
            crypt.update(keyDataBytes);
            byte[] bytes = crypt.digest();

            String res = "";
            for (byte aByte : bytes) res += String.format("%02x", aByte);
            return res;
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("This java has no support for SHA-1", e);
        }
    }

    private final static String CREDENTIAL_NAMESPACE_URI = "http://www.protogeni.net/resources/credential/credential.xsd";
    public static AbacCredential createSpeaksFor(Date expireDate, X509Certificate speakerToolCert, X509Certificate spokenForCert, Key spokenForPrivateKey) throws CredentialException {
        String userKeyhash = generateKeyId(spokenForCert).toLowerCase();
        String toolKeyhash = generateKeyId(speakerToolCert).toLowerCase();
        String expiresString = RFC3339Util.dateToRFC3339String(expireDate, true/*forceZulu*/);
        String credentialId = "_0";
        String template = "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n" +
                "<signed-credential xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\"\n" +
                "xsi:noNamespaceSchemaLocation=\"http://www.protogeni.net/resources/credential/credential.xsd\"\n" +
                "xsi:schemaLocation=\"http://www.protogeni.net/resources/credential/ext/policy/1\n" +
                "http://www.protogeni.net/resources/credential/ext/policy/1/policy.xsd\">\n" +
                "  <credential xml:id=\""+credentialId+"\">\n" +
                "\n" +
                "        <type>abac</type>\n" +
                "        <serial/>\n" +
                "        <owner_gid/>\n" +
                "        <target_gid/>\n" +
                "        <uuid/>\n" +
                "        <expires>"+expiresString+"</expires>\n" +
                "        <abac>\n" +
                "            <rt0>\n" +
                "                <version>1.1</version>\n" +
                "                <head>\n" +
                "   <ABACprincipal><keyid>"+userKeyhash+"</keyid></ABACprincipal>\n" +
                "   <role>speaks_for_"+userKeyhash+"</role>\n" +
                "</head>\n" +
                "<tail>\n" +
                "   <ABACprincipal><keyid>"+toolKeyhash+"</keyid></ABACprincipal>\n" +
                "</tail>\n" +
                "\n" +
                "            </rt0>\n" +
                "        </abac>\n" +
                "\n" +
                "</credential>\n" +
                "  <signatures>\n" +
                "  </signatures>\n" +
                "</signed-credential>";

        DocumentBuilderFactory docFactory = DocumentBuilderFactory.newInstance();
        docFactory.setNamespaceAware(true);
        DocumentBuilder docBuilder;
        Document doc;
        Element rootElement;
        try {
            docBuilder = docFactory.newDocumentBuilder();
            doc = docBuilder.parse(new InputSource(new StringReader(template)));

            rootElement = doc.getDocumentElement();
            assert rootElement != null;
            assert rootElement.getTagName() != null;
            assert Objects.equals(rootElement.getTagName(), "signed-credential");
//            assert rootElement.getNamespaceURI().equals(CREDENTIAL_NAMESPACE_URI);

//            NodeList credentialNl = rootElement.getElementsByTagNameNS(CREDENTIAL_NAMESPACE_URI, "credential");
            NodeList credentialNl = rootElement.getElementsByTagName("credential");
            assert credentialNl.getLength() == 1;
            LOG.debug("Need to mark " + credentialNl.getLength() + " <credential> xml:id attributes as xml IDs");
            for (int j = 0; j < credentialNl.getLength(); j++) {
                Node n = credentialNl.item(j);
                if (n.getNodeType() == Node.ELEMENT_NODE) {
                    Element credentialEl = (Element) n;

                    NamedNodeMap l = credentialEl.getAttributes();
                    for (int i = 0; i < l.getLength(); i++) {
                        Attr a = (Attr) l.item(i);
                        if (Objects.equals(a.getName(), "xml:id")) {
                            LOG.debug("Marking <credential> Attribute as id: " + a);
                            credentialEl.setIdAttributeNode(a, true);
                        } else
                            LOG.debug("<credential> Attribute is not id: '" + a.getNamespaceURI() + "' : '" + a.getName() + "' -> " + a);
                    }
                }
            }
        } catch (ParserConfigurationException | IOException | SAXException e) {
            throw new CredentialException("Failed to create credential: " + e.getMessage(), e);
        }

//        NodeList signaturesNodeList = doc.getElementsByTagNameNS(CREDENTIAL_NAMESPACE_URI, "signatures");
//        if (signaturesNodeList.getLength() <= 0) {
//            throw new CredentialException("Cannot find any Signature element: not a valid credential template.");
//        }
//        Element signaturesElement = (Element) signaturesNodeList.item(0);

        String credentialXmlString;
        try {
            credentialXmlString = XmlUtil.signXml(doc, credentialId, "signatures", spokenForCert, spokenForPrivateKey);
        } catch (TransformerException | XMLSecurityException e) {
            throw new CredentialException("Failed to sign generated credential: " + e.getMessage(), e);
        }

        return new AbacCredential(
                "speaksFor credential allowing "+getGeniSubjectName(speakerToolCert)+" to speak for "+getGeniSubjectName(spokenForCert)/*name*/,
                credentialXmlString,
                "abac"/*type*/, "1.1"/*version*/);
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;

        AbacCredential that = (AbacCredential) o;

        if (!Objects.equals(credentialXml, that.credentialXml)) return false;
        if (!Objects.equals(name, that.name)) return false;

        return true;
    }

    @Override
    public int hashCode() {
        int result = name.hashCode();
        result = 31 * result + credentialXml.hashCode();
        return result;
    }
}
