package org.keycloak.protocol.cas.representations;

import org.keycloak.dom.saml.v1.assertion.*;
import org.keycloak.dom.saml.v1.protocol.SAML11ResponseType;
import org.keycloak.dom.saml.v1.protocol.SAML11StatusCodeType;
import org.keycloak.dom.saml.v1.protocol.SAML11StatusType;
import org.keycloak.protocol.cas.utils.CASValidationException;
import org.keycloak.saml.common.exceptions.ProcessingException;
import org.keycloak.saml.processing.core.saml.v1.SAML11Constants;
import org.keycloak.saml.processing.core.saml.v1.writers.SAML11ResponseWriter;
import org.keycloak.services.validation.Validation;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;

import javax.xml.datatype.DatatypeConfigurationException;
import javax.xml.datatype.DatatypeFactory;
import javax.xml.datatype.XMLGregorianCalendar;
import javax.xml.namespace.QName;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.stream.XMLOutputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamWriter;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMResult;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import java.io.StringWriter;
import java.net.URI;
import java.time.ZoneOffset;
import java.time.ZonedDateTime;
import java.util.*;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class SamlResponseHelper {
    private final static DatatypeFactory factory;

    static {
        try {
            factory = DatatypeFactory.newInstance();
        } catch (DatatypeConfigurationException e) {
            throw new RuntimeException(e);
        }
    }

    public static SAML11ResponseType errorResponse(CASValidationException ex) {
        ZonedDateTime nowZoned = ZonedDateTime.now(ZoneOffset.UTC);
        XMLGregorianCalendar now = factory.newXMLGregorianCalendar(GregorianCalendar.from(nowZoned));

        return applyTo(new SAML11ResponseType("_" + UUID.randomUUID().toString(), now), obj -> {
            obj.setStatus(applyTo(new SAML11StatusType(), status -> {
                status.setStatusCode(new SAML11StatusCodeType(QName.valueOf("samlp:RequestDenied")));
                status.setStatusMessage(ex.getErrorDescription());
            }));
        });
    }

    public static SAML11ResponseType successResponse(String issuer, String username, Map<String, Object> attributes) {
        ZonedDateTime nowZoned = ZonedDateTime.now(ZoneOffset.UTC);
        XMLGregorianCalendar now = factory.newXMLGregorianCalendar(GregorianCalendar.from(nowZoned));

        return applyTo(new SAML11ResponseType("_" + UUID.randomUUID().toString(), now),
                obj -> {
                    obj.setStatus(applyTo(new SAML11StatusType(), status -> status.setStatusCode(SAML11StatusCodeType.SUCCESS)));
                    obj.add(applyTo(new SAML11AssertionType("_" + UUID.randomUUID().toString(), now), assertion -> {
                        assertion.setIssuer(issuer);
                        assertion.setConditions(applyTo(new SAML11ConditionsType(), conditions -> {
                            conditions.setNotBefore(now);
                            conditions.setNotOnOrAfter(factory.newXMLGregorianCalendar(GregorianCalendar.from(nowZoned.plusMinutes(5))));
                        }));
                        assertion.add(applyTo(new SAML11AuthenticationStatementType(
                                URI.create(SAMLCASConstants.AUTH_METHOD_PASSWORD),
                                now
                        ), stmt -> stmt.setSubject(toSubject(username))));
                        assertion.addAllStatements(toAttributes(username, attributes));
                    }));
                }
        );
    }

    private static List<SAML11StatementAbstractType> toAttributes(String username, Map<String, Object> attributes) {
        List<SAML11AttributeType> converted = attributeElements(attributes);
        if (converted.isEmpty()) {
            return Collections.emptyList();
        }
        return Collections.singletonList(applyTo(
                new SAML11AttributeStatementType(),
                attrs -> {
                    attrs.setSubject(toSubject(username));
                    attrs.addAllAttributes(converted);
                })
        );
    }

    private static List<SAML11AttributeType> attributeElements(Map<String, Object> attributes) {
        return attributes.entrySet().stream().flatMap(e ->
                toAttribute(e.getKey(), e.getValue())
        ).filter(a -> !a.get().isEmpty()).collect(Collectors.toList());
    }

    private static Stream<SAML11AttributeType> toAttribute(String name, Object value) {
        if (name == null || value == null) {
            return Stream.empty();
        }

        if (value instanceof Collection) {
            return Stream.of(samlAttribute(name, listString((Collection<?>) value)));
        }
        return Stream.of(samlAttribute(name, Collections.singletonList(value.toString())));
    }

    private static SAML11AttributeType samlAttribute(String name, List<Object> listString) {
        return applyTo(
                new SAML11AttributeType(name, URI.create("http://www.ja-sig.org/products/cas/")),
                attr -> attr.addAll(listString)
        );
    }

    private static List<Object> listString(Collection<?> value) {
        return value.stream().map(Object::toString).collect(Collectors.toList());
    }

    private static SAML11SubjectType toSubject(String username) {
        return applyTo(
                new SAML11SubjectType(),
                subject -> subject.setChoice(
                        new SAML11SubjectType.SAML11SubjectTypeChoice(
                                applyTo(
                                        new SAML11NameIdentifierType(username),
                                        ctype -> ctype.setFormat(nameIdFormat(username))
                                )
                        )
                )
        );
    }

    private static URI nameIdFormat(String username) {
        return URI.create(Validation.isEmailValid(username) ?
                SAMLCASConstants.FORMAT_EMAIL_ADDRESS :
                SAMLCASConstants.FORMAT_UNSPECIFIED
        );
    }

    private static <A> A applyTo(A input, Consumer<A> setter) {
        setter.accept(input);
        return input;
    }

    public static String soap(SAML11ResponseType response) {
        try {
            Document result = toDOM(response);

            Document doc = wrapSoap(result.getDocumentElement());
            return toString(doc);

        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public static Document toDOM(SAML11ResponseType response) throws ParserConfigurationException, XMLStreamException, ProcessingException {
        DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
        dbf.setNamespaceAware(true);

        XMLOutputFactory factory = XMLOutputFactory.newFactory();

        Document doc = dbf.newDocumentBuilder().newDocument();
        DOMResult result = new DOMResult(doc);
        XMLStreamWriter xmlWriter = factory.createXMLStreamWriter(result);
        SAML11ResponseWriter writer = new SAML11ResponseWriter(xmlWriter);
        writer.write(response);
        return doc;
    }

    private static Document wrapSoap(Node node) throws ParserConfigurationException {
        DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
        dbf.setNamespaceAware(true);
        Document doc = dbf.newDocumentBuilder().newDocument();

        Element envelope = doc.createElementNS("http://schemas.xmlsoap.org/soap/envelope/", "soap:Envelope");
        envelope.appendChild(doc.createElementNS("http://schemas.xmlsoap.org/soap/envelope/", "soap:Header"));
        Element body = doc.createElementNS("http://schemas.xmlsoap.org/soap/envelope/", "soap:Body");

        Node imported = doc.importNode(node, true);

        body.appendChild(imported);
        doc.appendChild(body);
        envelope.appendChild(body);
        doc.appendChild(envelope);
        return doc;
    }

    public static String toString(Document document) throws TransformerException {
        // Output the Document
        TransformerFactory tf = TransformerFactory.newInstance();
        Transformer t = tf.newTransformer();
        DOMSource source = new DOMSource(document);
        StringWriter writer = new StringWriter();
        StreamResult result = new StreamResult(writer);
        t.transform(source, result);
        return writer.toString();
    }
}