

package gov.va.med.cds.ars.jaxb;


//import com.sun.xml.internal.bind.marshaller.NamespacePrefixMapper;
import gov.va.med.cds.ars.requestresponse.generated.HTReportFilterType;
import gov.va.med.cds.util.StreamUtil;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.StringWriter;
import java.net.URL;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import javax.xml.bind.JAXBContext;
import javax.xml.bind.JAXBException;
import javax.xml.bind.Marshaller;
import javax.xml.bind.Unmarshaller;
import javax.xml.stream.XMLInputFactory;
import javax.xml.stream.XMLStreamReader;
import javax.xml.validation.Schema;
import javax.xml.validation.SchemaFactory;

//import org.apache.commons.pool.BaseKeyedPoolableObjectFactory;
import org.apache.commons.pool2.BaseKeyedPooledObjectFactory;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.commons.pool2.impl.GenericKeyedObjectPool;
import org.apache.commons.pool2.impl.GenericKeyedObjectPoolConfig;
import org.jetbrains.annotations.Nullable;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
import org.xml.sax.SAXException;

import com.sun.xml.bind.marshaller.NamespacePrefixMapper;




/**
 * Implementation which holds it's Schema instances in a map, and uses pooling
 * for JAXBContext, Marshaller and Unmarshaller instances. It is thread-safe and
 * re-entrant.
 */
public class PooledJaxbHelper
    implements
        JaxbHelper
{
    private static final String W3C_XML_SCHEMA_NS_URI = "http://www.w3.org/2001/XMLSchema";   
    private static Map<String, Schema> schemaMap = new ConcurrentHashMap<String, Schema>();
    private static Map<Class, JAXBContext> jaxbContextMap = new ConcurrentHashMap<Class, JAXBContext>();
//    private static GenericKeyedObjectPool<Class, JAXBContext> jaxbContextPool = new GenericKeyedObjectPool<Class, JAXBContext>(
//                    new JaxbContextFactory(), new CustomPoolConfig() );
    private static GenericKeyedObjectPool<PoolKey, Marshaller> marshallerPool = new GenericKeyedObjectPool<PoolKey, Marshaller>(
                    new MarshallerFactory(), new CustomPoolConfig() );
    private static GenericKeyedObjectPool<PoolKey, Unmarshaller> unmarshallerPool = new GenericKeyedObjectPool<PoolKey, Unmarshaller>(
                    new UnmarshallerFactory(), new CustomPoolConfig() );


    private static class PoolKey
    {
        private Class clazz;
        private String schemaLocation;


        private PoolKey( Class aClass, @Nullable String schemaLocation )
        {
            this.clazz = aClass;
            this.schemaLocation = schemaLocation;
        }


        @Override
        public boolean equals( Object o )
        {
            if ( this == o )
            {
                return true;
            }
            
            if ( o == null || getClass() != o.getClass() )
            {
                return false;
            }

            PoolKey poolKey = ( PoolKey )o;

            return clazz.equals( poolKey.clazz )
                        && !( schemaLocation != null ? !schemaLocation.equals( poolKey.schemaLocation ) : poolKey.schemaLocation != null );

        }


        @Override
        public int hashCode( )
        {
            int result = clazz.hashCode();
            result = 31 * result + ( schemaLocation != null ? schemaLocation.hashCode() : 0 );
            
            return result;
        }


        public Class getClazz( )
        {
            return clazz;
        }


        @Nullable
        public String getSchemaLocation( )
        {
            return schemaLocation;
        }
    }

    
    private static class MarshallerFactory
        extends
        BaseKeyedPooledObjectFactory<PoolKey, Marshaller>
    {
//        static final String NAMESPACE_PREFIX_MAPPER = "com.sun.xml.internal.bind.namespacePrefixMapper";
        static final String NAMESPACE_PREFIX_MAPPER = "com.sun.xml.bind.namespacePrefixMapper";
        
        @Override
        public Marshaller create(PoolKey key) throws Exception{
        	
        	 JAXBContext jaxbContext = getJAXBContext( key.getClazz() );

             Marshaller marshaller = jaxbContext.createMarshaller();
//NOTE: removed explicit use of NamespacePrefixMapper - instead edited package-info.Java to apply the same preferred prefix logic - the Marshaller was inconsistently ignoring preferred prefixes causing errors in processing
//             NamespacePrefixMapper prefixMapper = new HtReportFilterNamespaceMapper();
//             marshaller.setProperty( NAMESPACE_PREFIX_MAPPER, prefixMapper );

             if ( key.getSchemaLocation() != null )
             {
                 Schema schema = getSchema( key );
                 marshaller.setSchema( schema );
             }

             return marshaller;
        }
        
        @Override
        public PooledObject<Marshaller> wrap(Marshaller marshaller){
        	
        	return new DefaultPooledObject<Marshaller>(marshaller);
         	
        }
        
        @Override
        public PooledObject<Marshaller> makeObject(PoolKey key) throws Exception{
        	
        	Marshaller marshaller = create(key);
        	return wrap(marshaller);
        
        }
 
    }

    
    private static class UnmarshallerFactory
        extends
        BaseKeyedPooledObjectFactory<PoolKey, Unmarshaller>
    {
      
        @Override
        public Unmarshaller create(PoolKey key) throws Exception{
        	
//            JAXBContext jaxbContext = jaxbContextPool.borrowObject( key.getClazz() );
            JAXBContext jaxbContext = getJAXBContext( key.getClazz() );

            Unmarshaller unmarshaller = jaxbContext.createUnmarshaller();

            if ( key.getSchemaLocation() != null )
            {
                Schema schema = getSchema( key );
                unmarshaller.setSchema( schema );
            }

//            jaxbContextPool.returnObject( key.getClazz(), jaxbContext );

            return unmarshaller;
        }
        
        @Override
        public PooledObject<Unmarshaller> wrap(Unmarshaller unmarshaller){
        	
        	return new DefaultPooledObject<Unmarshaller>(unmarshaller);
         	
        }
        
        @Override
        public PooledObject<Unmarshaller> makeObject(PoolKey key) throws Exception{
        	
        	Unmarshaller unmarshaller = create(key);
        	return wrap(unmarshaller);
        
        }
        
    }

    
//    private static class JaxbContextFactory
//        extends
//            BaseKeyedPoolableObjectFactory<Class, JAXBContext>
//    {
//        @Override
//        public JAXBContext makeObject( Class clazz )
//            throws Exception
//        {
//            return JAXBContext.newInstance( clazz );
//        }
//    }

    
    private static JAXBContext getJAXBContext( Class aClass )
        throws JAXBException
    {
        JAXBContext jaxbContext = jaxbContextMap.get( aClass );

        if ( jaxbContext == null )
        {
            jaxbContext = JAXBContext.newInstance( aClass );
            jaxbContextMap.put( aClass, jaxbContext );
        }
        
        return jaxbContext;
    }
    
    
    private static class CustomPoolConfig
        extends
            GenericKeyedObjectPoolConfig
    {
        {
            setMaxIdlePerKey(3);
            //maxActive= 10;//moving from common-pools 1.x to 2.x could not find this flag or use anywhere ... not sure how this will affect pooling under load
            setMaxTotal(100);
            setMinIdlePerKey(1);
            setBlockWhenExhausted(true);//moving from common-pools 1.x - 1.x used WhenExhaustedGrow - but that was removed in 2.x, so hopefully blockiong when exhausted provides equal performnace
            setTimeBetweenEvictionRunsMillis(1000L * 60L * 10L);
            setNumTestsPerEvictionRun(50);
            setMinEvictableIdleTimeMillis(1000L * 60L * 5L);
        }
    }


    private static Schema getSchema( PoolKey poolKey )
        throws JAXBException,
               SAXException,
               IOException
    {
        Schema schema = schemaMap.get( poolKey.getSchemaLocation() );

        if ( schema == null )
        {
            SchemaFactory schemaFactory = SchemaFactory.newInstance( W3C_XML_SCHEMA_NS_URI );
            PathMatchingResourcePatternResolver resolver = new PathMatchingResourcePatternResolver();
            Resource resource = resolver.getResource( poolKey.getSchemaLocation() );
            URL fileURL = resource.getURL();
            schema = schemaFactory.newSchema( fileURL );
            schemaMap.put( poolKey.getSchemaLocation(), schema );
        }
        return schema;
    }


    @Override
    public <T> String marshal( T instance, @Nullable String schemaLocation )
        throws Exception
    {
        StringWriter result = new StringWriter();

        PoolKey poolKey = new PoolKey( instance.getClass(), schemaLocation );
        Marshaller marshaller = marshallerPool.borrowObject( poolKey );

        try
        {
            marshaller.marshal( instance, result );

            marshallerPool.returnObject( poolKey, marshaller );

            return result.toString();
        }
        catch ( Exception e )
        {
            marshallerPool.invalidateObject( poolKey, marshaller );
            throw new RuntimeException( e );
        }
    }


    @Override
    public <T> T unmarshal( String xml, Class<T> clazz, @Nullable String schemaLocation )
        throws Exception
    {
        T result;

        PoolKey poolKey = new PoolKey( clazz, schemaLocation );
        Unmarshaller unmarshaller = unmarshallerPool.borrowObject( poolKey );
        
        /*
         * Based on Fortify security scan we need to address XML External Entity Injection
         */
        XMLInputFactory xmlInputFactory = XMLInputFactory.newFactory();
        xmlInputFactory.setProperty( XMLInputFactory.IS_SUPPORTING_EXTERNAL_ENTITIES, false );
        xmlInputFactory.setProperty( XMLInputFactory.SUPPORT_DTD, false );
       // StreamSource xmlStream = new StreamSource( new StringReader( xml ) );
        InputStream xmlStream = new ByteArrayInputStream( xml.getBytes() );
        XMLStreamReader xmlStreamReader = xmlInputFactory.createXMLStreamReader( xmlStream );


        try
        {
            // noinspection unchecked
          //  result = ( T )unmarshaller.unmarshal( new StringReader( xml ) );
        	
        	 result = ( T )unmarshaller.unmarshal( xmlStreamReader );

            unmarshallerPool.returnObject( poolKey, unmarshaller );

            return result;
        }
        catch ( Exception e )
        {
            unmarshallerPool.invalidateObject( poolKey, unmarshaller );
            throw new RuntimeException( e );
        }
    }
    
    
    public static void main( String[] args )
        throws Exception
    {
        String SCHEMA_LOCATION = "classpath:main/resources/HTReportFilter.xsd";
        String filterInstance = StreamUtil.resourceToString( new FileSystemResource(
                        "./src/test/resources/filters/DMPFilter" ) );
        
        PooledJaxbHelper pooledJaxbHelper = new PooledJaxbHelper();
        HTReportFilterType htFilter = pooledJaxbHelper.unmarshal( filterInstance, HTReportFilterType.class, SCHEMA_LOCATION );
        String sampleXml = pooledJaxbHelper.marshal( htFilter, SCHEMA_LOCATION );
//        String sampleXml = pooledJaxbHelper.marshal( htFilter, null );
        System.out.println( "XML: " + sampleXml );
    }
}
