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

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.lang.annotation.Annotation;
import java.lang.reflect.Type;
import java.nio.charset.StandardCharsets;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import javax.xml.transform.stream.StreamSource;

import jakarta.activation.DataHandler;
import jakarta.activation.DataSource;
import jakarta.ws.rs.WebApplicationException;
import jakarta.ws.rs.core.MediaType;
import jakarta.ws.rs.core.MultivaluedMap;
import jakarta.ws.rs.core.Response;
import jakarta.ws.rs.ext.Providers;
import jakarta.xml.bind.JAXBContext;
import jakarta.xml.bind.JAXBException;
import jakarta.xml.bind.Marshaller;
import jakarta.xml.bind.Unmarshaller;
import jakarta.xml.bind.attachment.AttachmentMarshaller;
import jakarta.xml.bind.attachment.AttachmentUnmarshaller;

import org.jboss.resteasy.plugins.providers.jaxb.AbstractJAXBProvider;
import org.jboss.resteasy.plugins.providers.multipart.i18n.LogMessages;
import org.jboss.resteasy.plugins.providers.multipart.i18n.Messages;

/**
 * A special JAXB Provider. It is not a real provider, it is only used as a
 * helper class inside {@link XopWithMultipartRelatedReader} and
 * {@link XopWithMultipartRelatedWriter}.
 *
 * @author Attila Kiraly
 * @version $Revision: 1 $
 */
public class XopWithMultipartRelatedJAXBProvider extends
        AbstractJAXBProvider<Object> {

    private static class XopAttachmentMarshaller extends AttachmentMarshaller {
        private final MultipartRelatedOutput xopPackage;

        private XopAttachmentMarshaller(final MultipartRelatedOutput xopPackage) {
            this.xopPackage = xopPackage;
        }

        @Override
        public String addMtomAttachment(DataHandler data,
                String elementNamespace, String elementLocalName) {
            return addBinary(data.getDataSource(), data.getContentType());
        }

        @Override
        public String addMtomAttachment(byte[] data, int offset, int length,
                String mimeType, String elementNamespace,
                String elementLocalName) {
            ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(
                    data, offset, length);
            return addBinary(byteArrayInputStream, mimeType);
        }

        protected String addBinary(Object object, String mimeType) {
            String addrSpec = ContentIDUtils.generateRFC822AddrSpec();
            String contentID = ContentIDUtils
                    .generateContentIDFromAddrSpec(addrSpec);
            xopPackage.addPart(object, MediaType.valueOf(mimeType), contentID,
                    "binary");
            return ContentIDUtils.generateCidFromAddrSpec(addrSpec);
        }

        @Override
        public String addSwaRefAttachment(DataHandler data) {
            throw new UnsupportedOperationException(Messages.MESSAGES.swaRefsNotSupported());
        }

        @Override
        public boolean isXOPPackage() {
            return true;
        }
    }

    private static class InputPartBackedDataSource implements DataSource {
        private final String cid;
        private final InputPart inputPart;

        private InputPartBackedDataSource(final String cid, final InputPart inputPart) {
            this.cid = cid;
            this.inputPart = inputPart;
        }

        public String getContentType() {
            return inputPart.getMediaType().toString();
        }

        public String getName() {
            return cid;
        }

        public InputStream getInputStream() throws IOException {
            return inputPart.getBody(InputStream.class, null);
        }

        public OutputStream getOutputStream() throws IOException {
            throw new IOException(Messages.MESSAGES.dataSourceRepresentsXopMessagePart());
        }
    }

    private static class XopAttachmentUnmarshaller extends
            AttachmentUnmarshaller {

        private final MultipartRelatedInput xopPackage;

        private XopAttachmentUnmarshaller(final MultipartRelatedInput xopPackage) {
            this.xopPackage = xopPackage;
        }

        @Override
        public byte[] getAttachmentAsByteArray(String cid) {
            InputPart inputPart = getInputPart(cid);
            try {
                return inputPart.getBody(byte[].class, null);
            } catch (IOException e) {
                throw new IllegalArgumentException(Messages.MESSAGES.exceptionWhileExtractionAttachment(cid), e);
            }
        }

        @Override
        public DataHandler getAttachmentAsDataHandler(final String cid) {
            final InputPart inputPart = getInputPart(cid);
            return new DataHandler(
                    new InputPartBackedDataSource(cid, inputPart));
        }

        protected InputPart getInputPart(String cid) {
            String contentID = ContentIDUtils.convertCidToContentID(cid);
            InputPart inputPart = xopPackage.getRelatedMap().get(contentID);
            if (inputPart == null)
                throw new IllegalArgumentException(Messages.MESSAGES.noAttachmentFound(cid, contentID));
            return inputPart;
        }

        @Override
        public boolean isXOPPackage() {
            return true;
        }
    }

    public XopWithMultipartRelatedJAXBProvider(final Providers providers) {
        super();
        this.providers = providers;
    }

    @Override
    protected boolean isReadWritable(Class<?> type, Type genericType,
            Annotation[] annotations, MediaType mediaType) {
        throw new UnsupportedOperationException(Messages.MESSAGES.notMeantForStandaloneUsage());
    }

    public Object readFrom(Class<Object> type, Type genericType,
            Annotation[] annotations, MediaType mediaType,
            MultivaluedMap<String, String> httpHeaders,
            InputStream entityStream, final MultipartRelatedInput xopPackage)
            throws IOException {
        try {
            LogMessages.LOGGER.debugf("Provider : %s,  Method : readFrom", getClass().getName());
            InputPart rootPart = xopPackage.getRootPart();
            JAXBContext jaxb = findJAXBContext(type, annotations, rootPart
                    .getMediaType(), true);
            Unmarshaller unmarshaller = jaxb.createUnmarshaller();
            unmarshaller
                    .setAttachmentUnmarshaller(new XopAttachmentUnmarshaller(
                            xopPackage));
            return unmarshaller.unmarshal(new StreamSource(rootPart.getBody(
                    InputStream.class, null)));
        } catch (JAXBException e) {
            Response response = Response.serverError().build();
            throw new WebApplicationException(e, response);
        }
    }

    public void writeTo(Object t, Class<?> type, Type genericType,
            Annotation[] annotations, MediaType mediaType,
            MultivaluedMap<String, Object> httpHeaders,
            final MultipartRelatedOutput xopPackage) throws IOException {
        try {
            LogMessages.LOGGER.debugf("Provider : %s,  Method : writeTo", getClass().getName());
            Map<String, String> mediaTypeParameters = new LinkedHashMap<String, String>();
            mediaTypeParameters.put("charset", StandardCharsets.UTF_8.name());
            mediaTypeParameters.put("type", "text/xml");

            MediaType xopRootMediaType = new MediaType("application",
                    "xop+xml", mediaTypeParameters);

            Marshaller marshaller = getMarshaller(type, annotations,
                    xopRootMediaType);
            marshaller.setAttachmentMarshaller(new XopAttachmentMarshaller(
                    xopPackage));
            ByteArrayOutputStream xml = new ByteArrayOutputStream();
            marshaller.marshal(t, xml);

            OutputPart outputPart = xopPackage.addPart(xml.toByteArray(),
                    xopRootMediaType, ContentIDUtils.generateContentID(), null);
            List<OutputPart> outputParts = xopPackage.getParts();
            outputParts.remove(outputPart);
            outputParts.add(0, outputPart);
        } catch (JAXBException e) {
            Response response = Response.serverError().build();
            throw new WebApplicationException(e, response);
        }
    }
}
