package be.iminds.ilabt.jfed.lowlevel;


import be.iminds.ilabt.jfed.util.JFedTrustStore;
import be.iminds.ilabt.jfed.util.KeyUtil;
import org.apache.xml.security.Init;
import org.apache.xml.security.keys.KeyInfo;
import org.apache.xml.security.keys.storage.StorageResolver;
import org.apache.xml.security.keys.storage.implementations.KeyStoreResolver;
import org.apache.xml.security.signature.XMLSignature;
import org.apache.xml.security.utils.Constants;
import org.apache.xml.security.utils.ElementProxy;
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.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.stream.StreamSource;
import java.io.IOException;
import java.io.StringReader;
import java.security.KeyStore;
import java.security.PublicKey;
import java.security.cert.X509Certificate;
import java.util.*;

/**
 * <p>AnyCredential is a wrapper for any XML Credential.
 * <p/>
 * <p>It does not parse anything as the format is unknown <p>A name is stored alongside the actual credential, which is useful for listing the credential in GUI lists and comboboxes.
 */
public class AnyCredential {
    private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(AnyCredential.class);

    protected String name;
    @Nonnull
    protected final String credentialXml;

    protected Document xmlDoc;

    //see Geni Aggregate Manager API v3: http://groups.geni.net/geni/wiki/GAPI_AM_API_V3/CommonConcepts#credentials
    protected final String type;
    protected final String version;

    protected AnyCredential(String name, @Nonnull String credentialXml, String type, String version) {
        Objects.requireNonNull(credentialXml, "AnyCredential credentialXml may not be null");

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

    /**
     * Pre Geni API v3 credentials: type=geni_sfa version=2
     */
    public static AnyCredential createSfa2(String name, String credentialXml) throws CredentialException {
        return new SfaCredential(name, credentialXml, "geni_sfa", "2");
    }

    public static AnyCredential createAutoDetect(String name, String credentialXml) throws CredentialException {
        String type = null;
        String version = null;

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

            boolean withinType = false;
            boolean withinAbac = false;
            boolean withinAbacVersion = false;
            while (streamReader.hasNext()) {
                streamReader.next();
                switch (streamReader.getEventType()) {
                    case XMLStreamReader.START_ELEMENT: {
                        if (Objects.equals(streamReader.getName().getLocalPart(), "type")) {
                            withinType = true;
                        }
                        if (Objects.equals(streamReader.getName().getLocalPart(), "abac")) {
                            withinAbac = true;
                        }
                        if (withinAbac && Objects.equals(streamReader.getName().getLocalPart(), "version")) {
                            withinAbacVersion = true;
                        }
                        break;
                    }
                    case XMLStreamReader.END_ELEMENT: {
                        if (Objects.equals(streamReader.getName().getLocalPart(), "abac"))
                            withinType = false;
                        if (Objects.equals(streamReader.getName().getLocalPart(), "type"))
                            withinAbac = false;
                        if (Objects.equals(streamReader.getName().getLocalPart(), "version"))
                            withinAbacVersion = false;
                        break;
                    }
                    case XMLStreamReader.CHARACTERS: {
                        if (withinType) {
                            String typeText = streamReader.getText();
                            if (typeText != null) {
                                if (Objects.equals(typeText.trim().toLowerCase(), "abac"))
                                    type = "geni_abac";
                                if (typeText.trim().toLowerCase().equals("sfa") ||
                                        typeText.trim().toLowerCase().equals("privilege")) {
                                    type = "sfa";
                                    version = "3";
                                }
                            }
                        }
                        if (withinAbacVersion) {
                            String typeText = streamReader.getText();
                            if (typeText != null) {
                                type = "geni_abac";
                                version = "1";
//                                version = typeText.trim();
                            }
                        }
                        break;
                    }
                    default:
                        break;
                }
            }
        } catch (XMLStreamException e) {
            throw new CredentialException("Exception while parsing credential XML", e);
        } finally {
            if (streamReader != null)
                try {
                    streamReader.close();
                } catch (XMLStreamException e) {
                    LOG.error("Exception closing streamReader is ignored", e);
                }
        }

        if (type == null || version == null)
            throw new CredentialException("Could not auto detect credential type.");

        if (type.equals("geni_abac") || type.equals("geni_abac"))
            return new AbacCredential(name, credentialXml, type, version);

        if (type.equals("sfa") || type.equals("geni_sfa"))
            return new SfaCredential(name, credentialXml, type, version);

        throw new CredentialException("Unsupported credential type: \"" + type + "\" (version=\"" + version + "\")");
    }

    public static AnyCredential create(String name, String credentialXml, String type, String version) throws CredentialException {
        if (type.equalsIgnoreCase("sfa") || type.equalsIgnoreCase("geni_sfa"))
            return new SfaCredential(name, credentialXml, type, version);

        if (type.equalsIgnoreCase("abac") || type.equalsIgnoreCase("geni_abac"))
            return new AbacCredential(name, credentialXml, type, version);

        return new AnyCredential(name, credentialXml, type, version);
    }

    public static AnyCredential createFromV3Map(String credName, Map ht) throws BadReplyGeniException, CredentialException {
        Object geniTypeO = ht.get("geni_type");
        Object geniVersionO = ht.get("geni_version");
        Object genivalueO = ht.get("geni_value");

        if (geniTypeO == null || geniVersionO == null || genivalueO == null)
            throw new BadReplyGeniException("The API specified a List of Maps sepcifying credentials. These must contain the keys geni_type, geni_version and geni_value." +
                    " At least one of these is missing. Map content: " + ht);

        if (!(geniTypeO instanceof String))
            throw new BadReplyGeniException("The API specified a credential, but the geni_type field of the Map is of type " + geniTypeO.getClass().getName() + " instead of String. value=" + geniTypeO);

        if (!(geniVersionO instanceof Integer) && !(geniVersionO instanceof String))
            throw new BadReplyGeniException("The API specified a credential, but the geni_version field of the Map is of type " + geniVersionO.getClass().getName() + " instead of String with Integer. value=" + geniVersionO);

        if (!(genivalueO instanceof String))
            throw new BadReplyGeniException("The API specified a credential, but the geni_value field of the Map is of type " + genivalueO.getClass().getName() + " instead of String. value=" + genivalueO);

        return create(credName, (String) genivalueO, (String) geniTypeO, geniVersionO.toString());
    }

    @Nonnull
    public static AnyCredential toSingleCredential(@Nonnull List<AnyCredential> credentialList) throws JFedException {
        if (credentialList.isEmpty())
            throw new JFedException("Requested a single credential, but none were provided");

        if (credentialList.size() > 1)
            LOG.warn("requested toSingleCredential, but {} credentials in list", +credentialList.size());

        return credentialList.get(0);
    }

    public static Date getExpireTime(List<AnyCredential> credentials) {
        for (AnyCredential cred : credentials)
            if (cred instanceof SfaCredential) {
                SfaCredential sfaCred = (SfaCredential) cred;
                if (sfaCred.getExpiresDate() != null)
                    return sfaCred.getExpiresDate();
            }

        return null;

        //TODO handle case of multiple different dates (perfer one, and warn?)
    }

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

    public String getType() {
        return type;
    }

    public String getVersion() {
        return version;
    }

    public Map getGeniV3Map() {
        Map<String, Object> res = new HashMap<>();
        res.put("geni_type", type);
        res.put("geni_version", version);
        res.put("geni_value", credentialXml);
        return res;
    }

    public String getName() {
        return name;
    }

    public void setName(String newName) {
        this.name = newName;
    }

    public String getExpires() {
        return null;
    }

    public Date getExpiresDate() {
        return null;
    }

    /**
     * @return null if unknown
     */
    public Boolean isTargetSubAuthority() {
        return null;
    }

    @Override
    public String toString() {
        return credentialXml;
    }

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

        AnyCredential that = (AnyCredential) 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;
    }

    public List<AnyCredential> toCredentialList() {
        return Collections.singletonList(this);
    }

    /**
     * validate using credentials.xsd
     */
    //does not work yet
//    public boolean validate() {
//        try {
//            Source schemaSource = new StreamSource(AnyCredential.class.getResourceAsStream("credential.xsd"));
//            Source schemaSource2 = new StreamSource(AnyCredential.class.getResourceAsStream("sig.xsd"));
//            Source schemaSource3 = new StreamSource(AnyCredential.class.getResourceAsStream("xml.xsd"));
//            Source credentialSource = new StreamSource(new StringReader(credentialXml));
//            SchemaFactory schemaFactory = SchemaFactory
//                    .newInstance(XMLConstants.W3C_XML_SCHEMA_NS_URI);
//            Schema schema = schemaFactory.newSchema(new Source[] {schemaSource, schemaSource2, schemaSource3});
//            Validator validator = schema.newValidator();
//            validator.validate(credentialSource);
//            LOG.debug("credential is valid");
//            return true;
//        } catch (SAXException e) {
//            LOG.info("credential is NOT valid. Reason: " + e.getLocalizedMessage(), e);
//            return false;
//        } catch (IOException e) {
//            LOG.error("IOException validating credential", e);
//            return false;
//        }
//    }
    protected void initDoc() throws CredentialException {
        DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
        dbf.setNamespaceAware(true);
        try {
            xmlDoc = dbf.newDocumentBuilder().parse(new InputSource(new StringReader(credentialXml)));
        } catch (SAXException | ParserConfigurationException | IOException e) {
            throw new CredentialException("Error parsing credential XML: " + e.getMessage(), e);
        }

        //It seems that DOM does NOT see xml:id as an attribute of type xs:id as it should! So we manually set it so for any <credential>
        Element rootEl = xmlDoc.getDocumentElement();
        NodeList credentialNl = rootEl.getElementsByTagName("credential");
        LOG.trace("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.trace("Marking <credential> Attribute as id: " + a);
                        credentialEl.setIdAttributeNode(a, true);
                    } else
                        LOG.trace("<credential> Attribute is not id: '" + a.getNamespaceURI() + "' : '" + a.getName() + "' -> " + a);
                }
            }
        }
    }

    public boolean check(TestbedInfoSource testbedInfoSource) throws CredentialException {
        JFedTrustStore jFedTrustStore = new JFedTrustStore();
        return check(jFedTrustStore.getTrustStore());
    }

    public boolean check(KeyStore trustStore) throws CredentialException {
        if (xmlDoc == null) initDoc();
        assert xmlDoc != null; //initDoc should have been called

        NodeList signatureNodeList = xmlDoc.getElementsByTagNameNS("http://www.w3.org/2000/09/xmldsig#", "Signature");
        if (signatureNodeList.getLength() == 0) {
            throw new CredentialException("Cannot find any Signature element: not a valid credential.");
            //return false;
        }
        LOG.debug("There are " + signatureNodeList.getLength() + " <signature> elements in <Signature> that need to be checked.");
        try {
            for (int i = 0; i < signatureNodeList.getLength(); i++) {
                Element sigElement = (Element) signatureNodeList.item(i);

                LOG.debug("Checking <signature> with apache santurio library");
                Init.init();

                ElementProxy.setDefaultPrefix(Constants.SignatureSpecNS, "");

                //Creates a XMLSignature from the element and uses the filename as
                //the baseURI. That URI is prepended to all relative URIs.
                XMLSignature signature =
                        new XMLSignature(sigElement, null);


                //Get the KeyInfo object, which might contain some clues as to what
                //key was used to create the signature. It might also contain the
                //full cert.
                KeyInfo ki = signature.getKeyInfo();

                ki.addStorageResolver(new StorageResolver(new KeyStoreResolver(trustStore)));

                //First try to see if it is an X509Cert
                X509Certificate cert =
                        signature.getKeyInfo().getX509Certificate();

                if (cert != null) {
                    //check if the signature is valid using the cert
                    boolean check = signature.checkSignatureValue(cert);
                    LOG.debug("X509Certificate Check: " + check);
                    if (!check)
                        return false;
                } else {
                    //Maybe it's a public key
                    PublicKey pk = signature.getKeyInfo().getPublicKey();
                    if (pk != null) {
                        //check if the signature is valid using the public key
                        boolean check = signature.checkSignatureValue(pk);
                        LOG.debug("PublicKey Check: " + check);
                        if (!check)
                            return false;
                    } else {
                        //No X509Cert or PublicKey could be found.
                        LOG.debug("Could not find Certificate or PublicKey");
                        return false;
                    }
                }
            }
        } catch (Exception e) {
            //TODO: do we really want to catch all errors as validation failure?
            LOG.error("Error during checkSignedCredential", e);
            return false;
        }

        //if anything failed, it would have returned false above.
        return true;
    }

    public List<X509Certificate> getSignerCertificates() throws CredentialException {
        List<X509Certificate> res = new ArrayList<>();

        if (xmlDoc == null) initDoc();
        assert xmlDoc != null; //initDoc should have been called

        NodeList signatureNodeList = xmlDoc.getElementsByTagNameNS("http://www.w3.org/2000/09/xmldsig#", "Signature");
        if (signatureNodeList.getLength() == 0) {
            throw new CredentialException("Cannot find any Signature element: not a valid credential.");
            //return false;
        }
        LOG.debug("There are " + signatureNodeList.getLength() + " <signature> elements in <Signature>");
        try {
            for (int i = 0; i < signatureNodeList.getLength(); i++) {
                Element sigElement = (Element) signatureNodeList.item(i);

                //<signatures>
                //  <Signature xmlns="http://www.w3.org/2000/09/xmldsig#" xml:id="Sig__0">>
                //     <KeyInfo>
                //        <X509Data>
                //          <X509Certificate>

                NodeList nl = sigElement.getElementsByTagName("KeyInfo");
                if (nl == null || nl.getLength() == 0) continue;
                if (!(nl.item(0) instanceof Element)) continue;
                Element nlElement = (Element) nl.item(0);
                NodeList xd = nlElement.getElementsByTagName("X509Data");
                if (xd == null || xd.getLength() == 0) continue;
                if (!(xd.item(0) instanceof Element)) continue;
                Element xdElement = (Element) xd.item(0);
                NodeList cert = xdElement.getElementsByTagName("X509Certificate");
                if (cert == null || cert.getLength() == 0) continue;
                if (!(cert.item(0) instanceof Element)) continue;
                Element certElement = (Element) cert.item(0);

                String certPem = certElement.getTextContent().trim();
                if (!certPem.startsWith("-----BEGIN CERTIFICATE-----")) {
                    certPem = "-----BEGIN CERTIFICATE-----\n" + certPem + "\n-----END CERTIFICATE-----\n";
                }
                X509Certificate x509Certificate = KeyUtil.pemToX509Certificate(certPem);
                if (x509Certificate == null) continue;
                res.add(x509Certificate);
            }
        } catch (Exception e) {
            LOG.error("Error searching signer certificates", e);
        }

        return res;
    }

    public boolean isSpeaksFor() {
        return false;
    }
}
