/*
 * JBoss, the OpenSource J2EE webOS
 *
 * Distributable under LGPL license.
 * See terms of license at gnu.org.
 */
package org.jboss.net.axis.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 org.jboss.logging.Logger;

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
 * supports the javax.xml.transform.Source type required by SAAJ
 * 
 * @author Scott.Stark@jboss.org
 * @version $Revision: 1.2.8.2 $
 */
public class XMLDataContentHandler implements DataContentHandler
{
   private static Logger log = Logger.getLogger(XMLDataContentHandler.class);

   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(IOException e)
      {
         throw e;
      }
      catch(Exception e)
      {
         log.error("Failed to write xml", e);
         throw new IOException("Failed to write xml");
      }
   }

   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);
      }
   }
}