package org.jboss.resteasy.plugins.providers.jaxb;

import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.lang.annotation.Annotation;
import java.lang.reflect.Array;
import java.lang.reflect.Type;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;

import javax.xml.namespace.QName;
import javax.xml.transform.sax.SAXSource;
import javax.xml.transform.stream.StreamSource;

import jakarta.ws.rs.Consumes;
import jakarta.ws.rs.Produces;
import jakarta.ws.rs.WebApplicationException;
import jakarta.ws.rs.core.Context;
import jakarta.ws.rs.core.MediaType;
import jakarta.ws.rs.core.MultivaluedMap;
import jakarta.ws.rs.ext.ContextResolver;
import jakarta.ws.rs.ext.MessageBodyReader;
import jakarta.ws.rs.ext.Provider;
import jakarta.ws.rs.ext.Providers;
import jakarta.xml.bind.JAXBContext;
import jakarta.xml.bind.JAXBElement;
import jakarta.xml.bind.JAXBException;
import jakarta.xml.bind.Marshaller;
import jakarta.xml.bind.Unmarshaller;
import jakarta.xml.bind.annotation.XmlRootElement;
import jakarta.xml.bind.annotation.XmlSeeAlso;
import jakarta.xml.bind.annotation.XmlType;

import org.jboss.resteasy.annotations.providers.jaxb.DoNotUseJAXBProvider;
import org.jboss.resteasy.annotations.providers.jaxb.Wrapped;
import org.jboss.resteasy.core.ResteasyContext;
import org.jboss.resteasy.core.messagebody.AsyncBufferedMessageBodyWriter;
import org.jboss.resteasy.plugins.providers.jaxb.i18n.LogMessages;
import org.jboss.resteasy.plugins.providers.jaxb.i18n.Messages;
import org.jboss.resteasy.spi.ResteasyConfiguration;
import org.jboss.resteasy.spi.util.FindAnnotation;
import org.jboss.resteasy.spi.util.Types;
import org.w3c.dom.Element;
import org.xml.sax.InputSource;

/**
 * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
 * @version $Revision: 1 $
 */
@Provider
@Produces({ "application/xml", "application/*+xml", "text/xml", "text/*+xml" })
@Consumes({ "application/xml", "application/*+xml", "text/xml", "text/*+xml" })
public class CollectionProvider implements MessageBodyReader<Object>, AsyncBufferedMessageBodyWriter<Object> {
    @Context
    protected Providers providers;

    private boolean disableExternalEntities = true;
    private boolean enableSecureProcessingFeature = true;
    private boolean disableDTDs = true;

    public CollectionProvider() {
        LogMessages.LOGGER.debugf("Provider : %s,  Method : CollectionProvider", getClass().getName());
        ResteasyConfiguration context = ResteasyContext.getContextData(ResteasyConfiguration.class);
        if (context != null) {
            String s = context.getParameter("resteasy.document.expand.entity.references");
            if (s != null) {
                setDisableExternalEntities(!Boolean.parseBoolean(s));
            }
            s = context.getParameter("resteasy.document.secure.processing.feature");
            if (s != null) {
                setEnableSecureProcessingFeature(Boolean.parseBoolean(s));
            }
            s = context.getParameter("resteasy.document.secure.disableDTDs");
            if (s != null) {
                setDisableDTDs(Boolean.parseBoolean(s));
            }
        }
    }

    protected JAXBContextFinder getFinder(MediaType type) {
        ContextResolver<JAXBContextFinder> resolver = providers.getContextResolver(JAXBContextFinder.class, type);
        if (resolver == null)
            return null;
        return resolver.getContext(null);
    }

    public boolean isReadable(Class<?> type, Type genericType, Annotation[] annotations, MediaType mediaType) {
        return isWrapped(type, genericType, annotations, mediaType);
    }

    protected boolean isWrapped(Class<?> type, Type genericType, Annotation[] annotations, MediaType mediaType) {
        if ((Collection.class.isAssignableFrom(type) || type.isArray()) && genericType != null) {
            Class<?> baseType = Types.getCollectionBaseType(type, genericType);
            if (baseType == null)
                return false;
            baseType = XmlAdapterWrapper.xmlAdapterValueType(baseType, annotations);
            return (baseType.isAnnotationPresent(XmlRootElement.class) || baseType.isAnnotationPresent(XmlType.class)
                    || baseType.isAnnotationPresent(XmlSeeAlso.class) || JAXBElement.class.equals(baseType))
                    && (FindAnnotation.findAnnotation(baseType, annotations, DoNotUseJAXBProvider.class) == null)
                    && !IgnoredMediaTypes.ignored(baseType, annotations, mediaType);
        }
        return false;
    }

    public Object getJAXBObject(JAXBContextFinder finder, MediaType mediaType, Class<?> clazz, Element element)
            throws JAXBException {
        JAXBContext ctx = finder.findCachedContext(clazz, mediaType, null);
        return ctx.createUnmarshaller().unmarshal(element);
    }

    @SuppressWarnings("unchecked")
    public Object readFrom(Class<Object> type, Type genericType, Annotation[] annotations, MediaType mediaType,
            MultivaluedMap<String, String> httpHeaders, InputStream entityStream) throws IOException, WebApplicationException {
        LogMessages.LOGGER.debugf("Provider : %s,  Method : readFrom", getClass().getName());
        JAXBContextFinder finder = getFinder(mediaType);
        if (finder == null) {
            throw new JAXBUnmarshalException(Messages.MESSAGES.unableToFindJAXBContext(mediaType));
        }
        Class baseType = Types.getCollectionBaseType(type, genericType);
        XmlAdapterWrapper xmlAdapter = XmlAdapterWrapper.getXmlAdapter(baseType, annotations);
        JaxbCollection col = null;
        try {
            JAXBElement<JaxbCollection> ele = null;

            if (needsSecurity()) {
                SAXSource source = null;
                if (getCharset(mediaType) == null) {
                    source = new SAXSource(new InputSource(new InputStreamReader(entityStream, StandardCharsets.UTF_8)));
                } else {
                    source = new SAXSource(new InputSource(entityStream));
                }
                JAXBContext ctx = finder.findCachedContext(JaxbCollection.class, mediaType, annotations);
                Unmarshaller unmarshaller = ctx.createUnmarshaller();
                unmarshaller = new SecureUnmarshaller(unmarshaller, disableExternalEntities, enableSecureProcessingFeature,
                        disableDTDs);
                ele = unmarshaller.unmarshal(source, JaxbCollection.class);
            } else {
                StreamSource source = null;
                if (getCharset(mediaType) == null) {
                    source = new StreamSource(new InputStreamReader(entityStream, StandardCharsets.UTF_8));
                } else {
                    source = new StreamSource(entityStream);
                }
                JAXBContext ctx = finder.findCachedContext(JaxbCollection.class, mediaType, annotations);
                ele = ctx.createUnmarshaller().unmarshal(source, JaxbCollection.class);
            }

            Wrapped wrapped = FindAnnotation.findAnnotation(annotations, Wrapped.class);
            if (wrapped != null) {
                if (!wrapped.element().equals(ele.getName().getLocalPart())) {
                    throw new JAXBUnmarshalException(Messages.MESSAGES.collectionWrappingFailedLocalPart(wrapped.element(),
                            ele.getName().getLocalPart()));
                }
                if (!wrapped.namespace().equals(ele.getName().getNamespaceURI())) {
                    throw new JAXBUnmarshalException(Messages.MESSAGES.collectionWrappingFailedNamespace(wrapped.namespace(),
                            ele.getName().getNamespaceURI()));
                }
            }

            col = ele.getValue();
        } catch (JAXBException e) {
            throw new JAXBUnmarshalException(e);
        }

        try {
            JAXBContext ctx = null;
            if (xmlAdapter != null) {
                Class<?> adaptedType = xmlAdapter.getValueType();
                ctx = finder.findCachedContext(adaptedType, mediaType, null);
            } else {
                ctx = finder.findCachedContext(baseType, mediaType, null);
            }
            Unmarshaller unmarshaller = ctx.createUnmarshaller();
            unmarshaller = AbstractJAXBProvider.decorateUnmarshaller(baseType, annotations, mediaType, unmarshaller);
            if (type.isArray()) {
                Object array = Array.newInstance(baseType, col.getValue().size());
                for (int i = 0; i < col.getValue().size(); i++) {
                    Element val = (Element) col.getValue().get(i);
                    Object o = unmarshaller.unmarshal(val);
                    if (xmlAdapter != null) {
                        try {
                            o = xmlAdapter.unmarshal(o);
                        } catch (Exception e) {
                            throw new JAXBUnmarshalException(e);
                        }
                    }
                    Array.set(array, i, o);
                }
                return array;
            } else {
                Collection outCol = null;
                if (type.isInterface()) {
                    if (List.class.isAssignableFrom(type))
                        outCol = new ArrayList();
                    else if (SortedSet.class.isAssignableFrom(type))
                        outCol = new TreeSet();
                    else if (Set.class.isAssignableFrom(type))
                        outCol = new HashSet();
                    else
                        outCol = new ArrayList();
                } else {
                    try {
                        outCol = (Collection) type.newInstance();
                    } catch (Exception e) {
                        throw new JAXBUnmarshalException(e);
                    }
                }
                for (Object obj : col.getValue()) {
                    Element val = (Element) obj;
                    Object o = unmarshaller.unmarshal(val);
                    if (xmlAdapter != null) {
                        try {
                            o = xmlAdapter.unmarshal(o);
                        } catch (Exception e) {
                            throw new JAXBUnmarshalException(e);
                        }
                    }
                    outCol.add(o);
                }
                return outCol;
            }
        } catch (JAXBException e) {
            throw new JAXBUnmarshalException(e);
        }
    }

    public boolean isWriteable(Class<?> type, Type genericType, Annotation[] annotations, MediaType mediaType) {
        return isWrapped(type, genericType, annotations, mediaType);
    }

    public long getSize(Object entry, Class<?> type, Type genericType, Annotation[] annotations, MediaType mediaType) {
        return -1;
    }

    public void writeTo(Object entry, Class<?> type, Type genericType, Annotation[] annotations, MediaType mediaType,
            MultivaluedMap<String, Object> httpHeaders, OutputStream entityStream) throws IOException, WebApplicationException {
        LogMessages.LOGGER.debugf("Provider : %s,  Method : writeTo", getClass().getName());
        JAXBContextFinder finder = getFinder(mediaType);
        if (finder == null) {
            throw new JAXBMarshalException(Messages.MESSAGES.unableToFindJAXBContext(mediaType));
        }
        Class baseType = Types.getCollectionBaseType(type, genericType);
        XmlAdapterWrapper xmlAdapter = XmlAdapterWrapper.getXmlAdapter(baseType, annotations);
        if (xmlAdapter != null) {
            baseType = xmlAdapter.getValueType();
        }

        try {
            JAXBContext ctx = finder.findCacheContext(mediaType, annotations, JaxbCollection.class, baseType);
            JaxbCollection col = new JaxbCollection();
            if (type.isArray()) {
                Object[] array = (Object[]) entry;
                for (Object obj : array) {
                    if (xmlAdapter != null) {
                        try {
                            obj = xmlAdapter.marshal(obj);
                        } catch (Exception e) {
                            throw new JAXBUnmarshalException(e);
                        }
                    }
                    col.getValue().add(obj);
                }
            } else {
                Collection collection = (Collection) entry;
                for (Object obj : collection) {
                    if (xmlAdapter != null) {
                        try {
                            obj = xmlAdapter.marshal(obj);
                        } catch (Exception e) {
                            throw new JAXBUnmarshalException(e);
                        }
                    }
                    col.getValue().add(obj);
                }
            }

            String element = "collection";
            String namespaceURI = "";
            String prefix = "";

            Wrapped wrapped = FindAnnotation.findAnnotation(annotations, Wrapped.class);
            if (wrapped != null) {
                element = wrapped.element();
                namespaceURI = wrapped.namespace();
                prefix = wrapped.prefix();
            }

            JAXBElement<JaxbCollection> collection = new JAXBElement<JaxbCollection>(new QName(namespaceURI, element, prefix),
                    JaxbCollection.class, col);
            Marshaller marshaller = ctx.createMarshaller();
            AbstractJAXBProvider.decorateMarshaller(baseType, annotations, mediaType, marshaller);
            marshaller.marshal(collection, entityStream);
        } catch (JAXBException e) {
            throw new JAXBMarshalException(e);
        }
    }

    public boolean isDisableExternalEntities() {
        return disableExternalEntities;
    }

    public void setDisableExternalEntities(boolean disableExternalEntities) {
        this.disableExternalEntities = disableExternalEntities;
    }

    public boolean isEnableSecureProcessingFeature() {
        return enableSecureProcessingFeature;
    }

    public void setEnableSecureProcessingFeature(boolean enableSecureProcessingFeature) {
        this.enableSecureProcessingFeature = enableSecureProcessingFeature;
    }

    public boolean isDisableDTDs() {
        return disableDTDs;
    }

    public void setDisableDTDs(boolean disableDTDs) {
        this.disableDTDs = disableDTDs;
    }

    public static String getCharset(final MediaType mediaType) {
        if (mediaType != null) {
            return mediaType.getParameters().get("charset");
        }
        return null;
    }

    protected boolean needsSecurity() {
        return true;
    }
}
