package imageio;

import java.awt.Color;
import java.awt.Graphics2D;
import java.awt.image.BufferedImage;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;

import javax.imageio.IIOImage;
import javax.imageio.ImageIO;
import javax.imageio.ImageTypeSpecifier;
import javax.imageio.ImageWriteParam;
import javax.imageio.ImageWriter;
import javax.imageio.metadata.IIOMetadata;
import javax.imageio.stream.ImageOutputStream;

import com.sun.media.imageio.plugins.tiff.BaselineTIFFTagSet;
import com.sun.media.imageioimpl.plugins.tiff.TIFFField;
import com.sun.media.imageioimpl.plugins.tiff.TIFFIFD;
import com.sun.media.imageioimpl.plugins.tiff.TIFFImageMetadata;


/**
 * Test for TIFF JPEG compressor memory leak. 
 * @author Petr Synek
 */

public class TIFFJPEGMemoryLeak {
    
    public static void main(String[] args) {
        boolean success = true;
        TIFFJPEGMemoryLeak leakTest = new TIFFJPEGMemoryLeak();
        
        for (int i = 0; i < 3; i++) {
            System.out.println("Test round: " + i);
            success &= leakTest.test(i);
        }
        
        if (success) {
            System.out.println("\nTest OK");
        } else {
            System.out.println("\nTest failed");
        }
    }

    public boolean test(int testRoundIndex) {
        ImageWriter imageWriter = 
            (ImageWriter) ImageIO.getImageWritersByMIMEType(
                    "image/tiff").next();
        
        ImageWriteParam param = imageWriter.getDefaultWriteParam();
        if (param.canWriteCompressed()) {
            param.setCompressionMode(ImageWriteParam.MODE_EXPLICIT);
            param.setCompressionType("JPEG");
            param.setCompressionQuality(0.7f);
        }
        if (param.canWriteTiles()) {
            param.setTilingMode(ImageWriteParam.MODE_DISABLED);
        }
        
        BufferedImage image = null;
        ImageOutputStream ios = null;
        int dpi = 300;
        
        try {
            OutputStream outputStream = 
                new FileOutputStream("test" + testRoundIndex + ".tiff");
            
            ios = ImageIO.createImageOutputStream(outputStream);
            imageWriter.setOutput(ios);
            imageWriter.prepareWriteSequence(null);
            
            for (int pageNumber = 1; pageNumber <= 3; pageNumber++) {
                System.out.println(
                        "  writing page: " + testRoundIndex + ":" + pageNumber);

                image = getImage(testRoundIndex, pageNumber);
                
                IIOMetadata metadata = getMetadata(imageWriter, image, param, dpi);
                IIOImage iioImage = new IIOImage(image, null, metadata);
                imageWriter.writeToSequence(iioImage, param);
                
                if (image != null) {
                    image.flush();
                    image = null;
                }
                
                System.gc();
            }
        } catch (IOException ex) {
            ex.printStackTrace();
            return false;
        } catch (OutOfMemoryError e) {
            System.out.println("  OutOfMemory");
            return false;
        } finally {
            if (image != null) {
                try {
                    image.flush();
                    image = null;
                } catch (RuntimeException ex) {
                    ex.printStackTrace();
                }
            }
            
            if (ios != null) {
                try {
                    ios.flush();
                    ios.close();
                } catch (IOException ex) {
                    ex.printStackTrace();
                }
            }
            
            if (imageWriter != null) {
                try {
                    imageWriter.endWriteSequence();
                    imageWriter.reset();
                } catch (IOException ex) {
                    ex.printStackTrace();
                }
            }
        }
        
        return true;
    }
    
    private BufferedImage getImage(int testRoundIndex, int pageNumber) {
        int width = 1500;
        int height = 2000;
        
        BufferedImage image = new BufferedImage(
                width, height, BufferedImage.TYPE_INT_RGB);
        
        Graphics2D g2d = image.createGraphics();
        g2d.setBackground(Color.WHITE);
        g2d.setColor(Color.BLACK);
        
        g2d.clearRect(0, 0, width, height);
        g2d.drawString("Memory leak test page: " + testRoundIndex 
                       + ":" + pageNumber, 100, 100);
        return image;
    }
    
    private IIOMetadata getMetadata(ImageWriter imageWriter, 
            BufferedImage image, ImageWriteParam param, int dpi) 
    {
        TIFFImageMetadata tiffMetadata = 
            (TIFFImageMetadata) getIIOMetadata(image, imageWriter, param);
        TIFFIFD rootIFD = tiffMetadata.getRootIFD();
        BaselineTIFFTagSet base = BaselineTIFFTagSet.getInstance();
        
        // generate only 1 stripe
        rootIFD.addTIFFField(new TIFFField(base.getTag(278), image.getHeight()));
        
        // prepare array for DPI values (dpi/1)
        long dpiArray[][] = new long[1][2];
        dpiArray[0] = new long[] {dpi, 1};
        // set X DPI
        rootIFD.addTIFFField(new TIFFField(base.getTag(282), 5, 1, dpiArray));
        // set Y DPI
        rootIFD.addTIFFField(new TIFFField(base.getTag(283), 5, 1, dpiArray));
        
        return tiffMetadata;
    }
    
    private IIOMetadata getIIOMetadata(BufferedImage image, ImageWriter imageWriter, ImageWriteParam param) {
        ImageTypeSpecifier spec = ImageTypeSpecifier.createFromRenderedImage(image);
        IIOMetadata metadata = imageWriter.getDefaultImageMetadata(spec, param);
        return metadata;
    }
}