/*

 * JBoss, the OpenSource J2EE webOS

 *

 * Distributable under LGPL license.

 * See terms of license at gnu.org.

 */

package org.jboss.webservice.jaf;


import org.dom4j.Document;
import org.dom4j.io.DOMReader;
import org.dom4j.io.SAXReader;
import org.dom4j.io.XMLWriter;
import org.xml.sax.EntityResolver;
import org.xml.sax.InputSource;

import javax.activation.ActivationDataFlavor;
import javax.activation.DataContentHandler;
import javax.activation.DataSource;
import javax.xml.transform.Source;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.sax.SAXSource;
import javax.xml.transform.stream.StreamSource;
import java.awt.datatransfer.DataFlavor;
import java.awt.datatransfer.UnsupportedFlavorException;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.OutputStream;


/**
 * A JAF DataContentHandler for text/xml and application/xml that
 * <p/>
 * supports the javax.xml.transform.Source type required by SAAJ
 *
 * @author Scott.Stark@jboss.org
 * @version $Revision: 1.3 $
 */

public class XMLDataContentHandler implements DataContentHandler

{

   static NullEntityResolver resolver = new NullEntityResolver();


   public Object getContent(DataSource dataSource)

           throws IOException

   {

      StreamSource stream = new StreamSource(dataSource.getInputStream());

      return stream;

   }


   public Object getTransferData(DataFlavor dataFlavor, DataSource dataSource)

           throws UnsupportedFlavorException, IOException

   {

      Object data = null;

      // Handle text/xml & application/xml

      String mimeType = dataFlavor.getMimeType();

      if (mimeType.startsWith("text/xml")

              || mimeType.startsWith("application/xml"))

      {

         data = new StreamSource(dataSource.getInputStream());

      }

      return data;

   }


   public DataFlavor[] getTransferDataFlavors()

   {

      DataFlavor flavors[] = {

         new ActivationDataFlavor(Source.class, "text/xml", "XML"),

         new ActivationDataFlavor(Source.class, "application/xml", "XML")

      };

      return flavors;

   }


   public void writeTo(Object obj, String mimeType, OutputStream os)

           throws IOException

   {

      if (mimeType.startsWith("text/xml") == false

              && mimeType.startsWith("application/xml") == false)

      {

         String msg = "Only text/xml, application/xml are supported, mimeType=" + mimeType;

         throw new IOException(msg);

      }


      try

      {

         if (obj instanceof StreamSource)

         {

            StreamSource ss = (StreamSource)obj;

            XMLWriter writer = new XMLWriter(os);

            SAXReader reader = new SAXReader(false);

            reader.setIncludeExternalDTDDeclarations(false);

            reader.setEntityResolver(resolver);

            Document doc = reader.read(ss.getInputStream());

            writer.write(doc);

         }

         else if (obj instanceof DOMSource)

         {

            DOMSource ds = (DOMSource)obj;

            org.w3c.dom.Document doc = ds.getNode().getOwnerDocument();

            DOMReader reader = new DOMReader();

            Document doc2 = reader.read(doc);

            XMLWriter writer = new XMLWriter(os);

            writer.write(doc2);

         }

         else if (obj instanceof SAXSource)

         {

            SAXSource ss = (SAXSource)obj;

            XMLWriter writer = new XMLWriter(os);

            SAXReader reader = new SAXReader(false);

            reader.setIncludeExternalDTDDeclarations(false);

            reader.setEntityResolver(resolver);

            Document doc = reader.read(ss.getInputSource());

            writer.write(doc);

         }

      }

      catch (Exception e)

      {

         if (e instanceof IOException)

            throw (IOException)e;

         else

         {

            IOException ioe = new IOException("Failed to write xml");

            ioe.initCause(e);

         }

      }

   }


   static class NullEntityResolver implements EntityResolver

   {

      static String dtd = "";

      static ByteArrayInputStream bais = new ByteArrayInputStream(dtd.getBytes());

      public InputSource resolveEntity(String publicId, String systemId)

      {

         return new InputSource(bais);

      }

   }

}