package be.iminds.ilabt.jfed.util;

import org.apache.xml.security.Init;
import org.apache.xml.security.exceptions.XMLSecurityException;
import org.apache.xml.security.transforms.Transforms;
import org.apache.xml.security.utils.Constants;
import org.apache.xml.security.utils.ElementProxy;
import org.slf4j.LoggerFactory;
import org.w3c.dom.*;
import org.w3c.dom.bootstrap.DOMImplementationRegistry;
import org.w3c.dom.ls.DOMImplementationLS;
import org.w3c.dom.ls.LSSerializer;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;

import javax.xml.crypto.MarshalException;
import javax.xml.crypto.dsig.*;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
import javax.xml.crypto.dsig.keyinfo.KeyValue;
import javax.xml.crypto.dsig.keyinfo.X509Data;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;
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.*;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import javax.xml.transform.stream.StreamSource;
import java.io.IOException;
import java.io.StringReader;
import java.io.StringWriter;
import java.security.InvalidAlgorithmParameterException;
import java.security.Key;
import java.security.KeyException;
import java.security.NoSuchAlgorithmException;
import java.security.cert.X509Certificate;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * XmlUtils
 */
public class XmlUtil {
    private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(XmlUtil.class);

    /* This new version uses only JDK classes, no apache xerces. Depends on java 1.6. */
    public static String formatXmlFromString(String unformattedXmlString) {
        try {
            final Document document = parseXmlFromStringMakeExceptionsRuntime(unformattedXmlString);

            DOMImplementationRegistry registry = DOMImplementationRegistry.newInstance();
            DOMImplementationLS impl = (DOMImplementationLS) registry.getDOMImplementation("LS");
            LSSerializer writer = impl.createLSSerializer();
            writer.getDomConfig().setParameter("format-pretty-print", Boolean.TRUE);
            writer.getDomConfig().setParameter("xml-declaration", false); // true -> the xml declaration is added

            return writer.writeToString(document);
        } catch (ClassNotFoundException | IllegalAccessException | InstantiationException e) {
            throw new RuntimeException(e);
        }
    }

    /* throws no errors, returns original string if something goes wrong */
    public static String formatXmlFromString_alwaysSafe(String unformattedXmlString) {
        try {
            final Document document = parseXmlFromStringMakeExceptionsRuntime(unformattedXmlString);

            DOMImplementationRegistry registry = DOMImplementationRegistry.newInstance();
            DOMImplementationLS impl = (DOMImplementationLS) registry.getDOMImplementation("LS");
            LSSerializer writer = impl.createLSSerializer();
            writer.getDomConfig().setParameter("format-pretty-print", Boolean.TRUE);
            writer.getDomConfig().setParameter("xml-declaration", false); // true -> the xml declaration is added

            return writer.writeToString(document);
        } catch (Exception e) {
            return unformattedXmlString;
        }
    }

    /* throws no errors, returns original string if something goes wrong */
    public static String elementToXmlString(Element el) {
        try {
            TransformerFactory transFactory = TransformerFactory.newInstance();
            Transformer transformer = transFactory.newTransformer();
            StringWriter buffer = new StringWriter();
            transformer.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION, "yes");
            transformer.transform(new DOMSource(el), new StreamResult(buffer));
            return buffer.toString();
        } catch (TransformerConfigurationException e) {
            LOG.error("TransformerConfigurationException in getCurrentRequestRspec()", e);
        } catch (TransformerException e) {
            LOG.error("TransformerException in getCurrentRequestRspec()", e);
        }
        return "error converting DOM element to Xml";
    }

    public static Document parseXmlFromStringMakeExceptionsRuntime(String in) {
        return parseXmlFromStringMakeExceptionsRuntime(in, false);
    }

    public static Document parseXmlFromStringMakeExceptionsRuntime(String in, boolean namespaceAware) {
        try {
            DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
            dbf.setNamespaceAware(namespaceAware);
            DocumentBuilder db = dbf.newDocumentBuilder();
            InputSource is = new InputSource(new StringReader(in));
            return db.parse(is);
        } catch (ParserConfigurationException | IOException | SAXException e) {
            throw new RuntimeException(e);
        }
    }

    public static Document parseXmlFromString(String in, boolean namespaceAware) throws SAXException, IOException {
        try {
            DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
            dbf.setNamespaceAware(namespaceAware);
            DocumentBuilder db = dbf.newDocumentBuilder();
            InputSource is = new InputSource(new StringReader(in));
            return db.parse(is);
        } catch (ParserConfigurationException e) {
            throw new RuntimeException(e);
        }
    }


    public static String signXml(Document doc, String elementToSignId,
                                 String elementToAddSignatureToName, X509Certificate cert, Key privateKey)
            throws TransformerException, XMLSecurityException {

        doc.normalizeDocument();

        if (cert.getPublicKey() instanceof RSAPublicKey && privateKey instanceof RSAPrivateKey) {
            assert Objects.equals(((RSAPublicKey) cert.getPublicKey()).getModulus(), ((RSAPrivateKey) privateKey).getModulus());
            if (LOG.isDebugEnabled()) {
                boolean same =
                        Objects.equals(((RSAPublicKey) cert.getPublicKey()).getModulus(), ((RSAPrivateKey) privateKey).getModulus());
                LOG.debug("signXml input cert and privateKey have same modulus");
            }
        } else
            LOG.debug("not Rsa private and/or public key: no check done. cert.getPublicKey() is " + cert.getPublicKey().getClass().getName() + " privateKey is " + privateKey.getClass().getName());


        //        return signXmlWithoutSanturio(doc, elementToSignId, elementToAddSignatureToName, cert, privateKey);
        return signXmlWithSanturio(doc, elementToSignId, elementToAddSignatureToName, cert, privateKey);
    }

    public static String signXmlWithoutSanturio(Document doc, String elementToSignId,
                                                String elementToAddSignatureToName, X509Certificate cert, Key privateKey)
            throws NoSuchAlgorithmException, InvalidAlgorithmParameterException, KeyException,
            MarshalException, XMLSignatureException, TransformerException {
        Element elementToAddSignatureTo = (Element) doc.getElementsByTagName(elementToAddSignatureToName).item(0);
        //        Element elementToSign = doc.getElementById(elementToSignId);
        assert elementToAddSignatureTo != null;
        //        assert elementToSign != null;

        XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");

        Reference ref = fac.newReference
                ("#" + elementToSignId, fac.newDigestMethod(DigestMethod.SHA1, null),
                        Collections.singletonList
                                (fac.newTransform
                                        (Transform.ENVELOPED, (TransformParameterSpec) null)),
                        null, null);

        // Create the SignedInfo
        SignedInfo si = fac.newSignedInfo
                (fac.newCanonicalizationMethod
                                (CanonicalizationMethod.INCLUSIVE,
                                        (C14NMethodParameterSpec) null),
                        fac.newSignatureMethod(SignatureMethod.RSA_SHA1, null),
                        Collections.singletonList(ref));


        // Create a KeyValue
        KeyInfoFactory kif = fac.getKeyInfoFactory();
        KeyValue kv = kif.newKeyValue(cert.getPublicKey());
        List<Object> x509DataList = new ArrayList<>();
        x509DataList.add(cert); //adds X509Certificate element
        //        x509DataList.add(""); //adds X509SubjectName element
        //can also use X509IssuerSerial, see newX509Data javadoc
        X509Data x509Data = kif.newX509Data(x509DataList);

        // Create a KeyInfo and add the KeyValue to it
        List<Object> kiList = new ArrayList<>();
        kiList.add(kv);
        kiList.add(x509Data);
        KeyInfo ki = kif.newKeyInfo(kiList);


        // Create a DOMSignContext and specify the PrivateKey and
        // location of the resulting XMLSignature's parent element
        DOMSignContext dsc = new DOMSignContext(privateKey, elementToAddSignatureTo);

        // Create the XMLSignature (but don't sign it yet)
        javax.xml.crypto.dsig.XMLSignature signature = fac.newXMLSignature(si, ki);

        // Marshal, generate (and sign) the enveloped signature
        signature.sign(dsc);


        TransformerFactory tf = TransformerFactory.newInstance();
        Transformer trans = tf.newTransformer();
        trans.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION, "no");
        trans.setOutputProperty(OutputKeys.ENCODING, "utf-8");
        trans.setOutputProperty(OutputKeys.INDENT, "no");
        final StringWriter stringWriter = new StringWriter();
        trans.transform(new DOMSource(doc), new StreamResult(stringWriter));
        return stringWriter.getBuffer().toString();

    }

    public static String signXmlWithSanturio(Document doc, String elementToSignId,
                                             String elementToAddSignatureToName, X509Certificate cert, Key privateKey)
            throws XMLSecurityException, TransformerException {
        Init.init();
        ElementProxy.setDefaultPrefix(Constants.SignatureSpecNS, "");

        // For testing:
        // Take the XML you get from here, and remove everything starting after <signatures>
        // and paste this (adapt Sig_ ....) and save as ..._raw
        //<Signature xmlns="http://www.w3.org/2000/09/xmldsig#" xml:id="Sig_cred-774800473">
        //<SignedInfo>
        //<CanonicalizationMethod Algorithm="http://www.w3.org/2001/10/xml-exc-c14n#"/>
        //<SignatureMethod Algorithm="http://www.w3.org/2000/09/xmldsig#rsa-sha1"/>
        //<Reference>
        //<Transforms>
        //<Transform Algorithm="http://www.w3.org/2000/09/xmldsig#enveloped-signature"/>
        //<Transform Algorithm="http://www.w3.org/2001/10/xml-exc-c14n#"/>
        //</Transforms>
        //<DigestMethod Algorithm="http://www.w3.org/2000/09/xmldsig#sha1"/>
        //<DigestValue />
        //</Reference>
        //</SignedInfo>
        //<SignatureValue />
        //<KeyInfo>
        //<X509Data />
        //</KeyInfo>
        //</Signature></signatures></signed-credential>
        // run on boss: xmlsec1 --sign --privkey-pem /tmp/wall2.pem,/tmp/wall2.pem --print-debug --output /tmp/speaksfordebug_xmlsec2.xml /tmp/speaksfordebug_raw
        // and verify with
        // /usr/testbed/sbin/verifygenicred -d /tmp/speaksfordebug3
        // file has to be in /tmp because of some hardcode stuff
        // that verifygenicred can also be tested on jFed generated credentials (should be OK on one of the first lines)

        final org.apache.xml.security.signature.XMLSignature sig =
                new org.apache.xml.security.signature.XMLSignature(doc, null, org.apache.xml.security.signature.XMLSignature.ALGO_ID_SIGNATURE_RSA, org.apache.xml.security.c14n.Canonicalizer.ALGO_ID_C14N_EXCL_OMIT_COMMENTS);
        //        final XMLSignature sig = new XMLSignature(elementToSign, null);
        Element elementToAddSignatureTo = (Element) doc.getElementsByTagName(elementToAddSignatureToName).item(0);

        Attr idAttr = doc.createAttribute("xml:id");
        String credentialId = "Sig_" + elementToSignId;
        idAttr.setValue(credentialId);
        sig.getElement().setAttributeNode(idAttr);
        sig.getElement().setIdAttributeNode(idAttr, true);


        elementToAddSignatureTo.appendChild(sig.getElement());
        final Transforms transforms = new Transforms(doc);
        transforms.addTransform(Transforms.TRANSFORM_ENVELOPED_SIGNATURE);
        transforms.addTransform(Transforms.TRANSFORM_C14N_EXCL_OMIT_COMMENTS);
        sig.addDocument("#" + elementToSignId, transforms, Constants.ALGO_ID_DIGEST_SHA1);

        sig.addKeyInfo(cert);
        sig.addKeyInfo(cert.getPublicKey());
        //sig.addKeyInfo(cert.getPublicKey());
        sig.sign(privateKey);


        final StringWriter stringWriter = new StringWriter();
        TransformerFactory transFactory = TransformerFactory.newInstance();
        Transformer transformer = transFactory.newTransformer();
        transformer.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION, "no");
        transformer.setOutputProperty(OutputKeys.ENCODING, "utf-8");
        transformer.setOutputProperty(OutputKeys.INDENT, "no");
        transformer.transform(new DOMSource(doc), new StreamResult(stringWriter));

        String res = stringWriter.getBuffer().toString();

        //a HACK, because emulab doesn;t like newlines for <SignatureValue> and <X509Certificate>
        res = res.replaceFirst(Pattern.quote("<SignatureValue>\n"), Matcher.quoteReplacement("<SignatureValue>"));
        res = res.replaceFirst(Pattern.quote("<X509Certificate>\n"), Matcher.quoteReplacement("<X509Certificate>"));
        res = res.replaceFirst(Pattern.quote("\n</SignatureValue>"), Matcher.quoteReplacement("</SignatureValue>"));
        res = res.replaceFirst(Pattern.quote("\n</X509Certificate>"), Matcher.quoteReplacement("</X509Certificate>"));

        return res;

        //meses up xsi:xmlns
        //        ByteArrayOutputStream os = new ByteArrayOutputStream();
        //        os.write(Canonicalizer.getInstance(Canonicalizer.ALGO_ID_C14N_WITH_COMMENTS).canonicalizeSubtree(doc));
        //        return os.toString("UTF-8");
    }


    public static void markXmlIdsAsXmlId(Document doc) {
        //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 attribute which is named xml:id
        NodeList allNodes = doc.getElementsByTagName("*");
        for (int j = 0; j < allNodes.getLength(); j++) {
            Node n = allNodes.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")) {
                        credentialEl.setIdAttributeNode(a, true);
                    }
                }
            }
        }
    }


    //TODO: these should specify when tey return null and when they throw an error. This should probably be part of the parameters.

    public static List<Element> getDomElementChildElements(Element e, String tagName) {
        assert e != null;
        assert tagName != null;

        NodeList resList = e.getElementsByTagName(tagName);
        List<Element> resElements = new ArrayList<>();
        for (int i = 0; i < resList.getLength(); i++) {
            org.w3c.dom.Node resNode = resList.item(i);
            assert resNode != null;
            assert resNode.getNodeType() == org.w3c.dom.Node.ELEMENT_NODE;
            Element resElement = (Element) resNode;
            resElements.add(resElement);
        }
        return resElements;
    }

    public static List<Element> getDomElementChildElementsNS(Element e, String namespace, String tagName) {
        assert e != null;
        assert tagName != null;

        NodeList resList = e.getElementsByTagNameNS(namespace, tagName);
        List<Element> resElements = new ArrayList<>();
        for (int i = 0; i < resList.getLength(); i++) {
            org.w3c.dom.Node resNode = resList.item(i);
            assert resNode != null;
            assert resNode.getNodeType() == org.w3c.dom.Node.ELEMENT_NODE;
            Element resElement = (Element) resNode;
            resElements.add(resElement);
        }
        return resElements;
    }

    public static Element getDomElementFirstChildElement(Element e, String tagName) {
        NodeList resList = e.getElementsByTagName(tagName);
        assert resList.getLength() == 1;
        org.w3c.dom.Node resNode = resList.item(0);
        assert resNode != null;
        assert resNode.getNodeType() == org.w3c.dom.Node.ELEMENT_NODE;
        return (Element) resNode;
    }

    public static Element getDomElementFirstChildElementNS(Element e, String namespace, String tagName) {
        NodeList resList = e.getElementsByTagNameNS(namespace, tagName);
        assert resList.getLength() == 1;
        org.w3c.dom.Node resNode = resList.item(0);
        assert resNode != null;
        assert resNode.getNodeType() == org.w3c.dom.Node.ELEMENT_NODE;
        return (Element) resNode;
    }

    public static String getDomElementFirstChildElementText(Element e, String tagName) {
        NodeList resList = e.getElementsByTagName(tagName);
        if (resList == null || resList.getLength() < 1) return null;
        assert resList.getLength() == 1;
        org.w3c.dom.Node resNode = resList.item(0);
        if (resNode == null) return null;
        if (resNode.getNodeType() != org.w3c.dom.Node.ELEMENT_NODE) return null;
//        assert resNode != null;
//        assert resNode.getNodeType() == org.w3c.dom.Node.ELEMENT_NODE;
        return resNode.getTextContent();
    }


    public static boolean couldBeXmlVeryQuickTest(String s) {
        return s.startsWith("<");
    }
    public static boolean isXmlLikeQuickTest(String s) {
        //if there are 10 valid XML events, it's most likely xml
        return isXmlQuickTestHelper(s, 10);
    }
    public static boolean isWellFormedXmlQuickTest(String s) {
        return isXmlQuickTestHelper(s, -1);
    }

    /**
     *
     * @param s String
     * @param eventLimit if this number of events are successfully parsed, the string is considered to be XML. Use 0 or -1 to require the entire string to be checked.
     */
    private static boolean isXmlQuickTestHelper(String s, int eventLimit) {
        if (!couldBeXmlVeryQuickTest(s)) return false;
        int eventCount = 0;
        XMLStreamReader reader = null;
        try {
            XMLInputFactory xif = XMLInputFactory.newFactory();
            StreamSource xml = new StreamSource(new StringReader(s));
            reader = xif.createXMLStreamReader(xml);

            reader.next();
            while (reader.hasNext()) {
                reader.next();
                eventCount++;

                if (eventLimit > 0 && eventCount >= eventLimit)
                    return true;
            }

            return true;
        } catch (XMLStreamException e) {
            //no need to log actual stacktrace etc.
            LOG.debug("isXmlQuickTest caught a parsing exception: not well formed XML. XMLStreamException=" + e);
            return false;
        } finally {
            if (reader != null)
                try {
                    reader.close();
                } catch (XMLStreamException e) {
                    //ignore
                }
        }
    }
}
