/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.sanselan.formats.jpeg.exifRewrite;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;

import org.apache.sanselan.ImageReadException;
import org.apache.sanselan.ImageWriteException;
import org.apache.sanselan.common.BinaryFileParser;
import org.apache.sanselan.common.byteSources.ByteSource;
import org.apache.sanselan.common.byteSources.ByteSourceArray;
import org.apache.sanselan.common.byteSources.ByteSourceFile;
import org.apache.sanselan.common.byteSources.ByteSourceInputStream;
import org.apache.sanselan.formats.jpeg.JpegConstants;
import org.apache.sanselan.formats.jpeg.JpegUtils;
import org.apache.sanselan.formats.tiff.write.TiffImageWriterBase;
import org.apache.sanselan.formats.tiff.write.TiffImageWriterLossless;
import org.apache.sanselan.formats.tiff.write.TiffImageWriterLossy;
import org.apache.sanselan.formats.tiff.write.TiffOutputSet;
import org.apache.sanselan.util.Debug;

/**
 * Interface for Exif write/update/remove functionality for Jpeg/JFIF images.
 * <p>
 * <p>
 * See the source of the ExifMetadataUpdateExample class for example usage.
 *
 * @see org.apache.sanselan.sampleUsage.WriteExifMetadataExample
 */
public class ExifRewriter extends BinaryFileParser implements JpegConstants
{
    /**
     * Constructor. to guess whether a file contains an image based on its file extension.
     */
    public ExifRewriter()
    {
        setByteOrder(BYTE_ORDER_NETWORK);
    }

    /**
     * Constructor.
     * <p>
     * @param  byteOrder  byte order of EXIF segment.  Optional.  See BinaryConstants class.
     *
     * @see org.apache.sanselan.common.BinaryConstants
     */
    public ExifRewriter(int byteOrder)
    {
        setByteOrder(byteOrder);
    }

    private static class JFIFPieces
    {
        public final List pieces;
        public final List exifPieces;

        public JFIFPieces(final List pieces, final List exifPieces)
        {
            this.pieces = pieces;
            this.exifPieces = exifPieces;
        }

    }

    private abstract static class JFIFPiece
    {
        protected abstract void write(OutputStream os) throws IOException;
    }

    private static class JFIFPieceSegment extends JFIFPiece
    {
        public final int marker;
        public final byte markerBytes[];
        public final byte markerLengthBytes[];
        public final byte segmentData[];

        public JFIFPieceSegment(final int marker, final byte[] markerBytes,
                final byte[] markerLengthBytes, final byte[] segmentData)
        {
            this.marker = marker;
            this.markerBytes = markerBytes;
            this.markerLengthBytes = markerLengthBytes;
            this.segmentData = segmentData;
        }

        protected void write(OutputStream os) throws IOException
        {
            os.write(markerBytes);
            os.write(markerLengthBytes);
            os.write(segmentData);
        }
    }

    private static class JFIFPieceSegmentExif extends JFIFPieceSegment
    {

        public JFIFPieceSegmentExif(final int marker, final byte[] markerBytes,
                final byte[] markerLengthBytes, final byte[] segmentData)
        {
            super(marker, markerBytes, markerLengthBytes, segmentData);
        }
    }

    private static class JFIFPieceImageData extends JFIFPiece
    {
        public final byte markerBytes[];
        public final byte imageData[];

        public JFIFPieceImageData(final byte[] markerBytes,
                final byte[] imageData)
        {
            super();
            this.markerBytes = markerBytes;
            this.imageData = imageData;
        }

        protected void write(OutputStream os) throws IOException
        {
            os.write(markerBytes);
            os.write(imageData);
        }
    }

    private JFIFPieces analyzeJFIF(ByteSource byteSource)
            throws ImageReadException, IOException
    //            , ImageWriteException
    {
        final ArrayList pieces = new ArrayList();
        final List exifPieces = new ArrayList();

        JpegUtils.Visitor visitor = new JpegUtils.Visitor()
        {
            // return false to exit before reading image data.
            public boolean beginSOS()
            {
                return true;
            }

            public void visitSOS(int marker, byte markerBytes[],
                    byte imageData[])
            {
                pieces.add(new JFIFPieceImageData(markerBytes, imageData));
            }

            // return false to exit traversal.
            public boolean visitSegment(int marker, byte markerBytes[],
                    int markerLength, byte markerLengthBytes[],
                    byte segmentData[]) throws
            //                    ImageWriteException,
                    ImageReadException, IOException
            {
                if (marker != JPEG_APP1_Marker)
                {
                    pieces.add(new JFIFPieceSegment(marker, markerBytes,
                            markerLengthBytes, segmentData));
                }
                else if (!byteArrayHasPrefix(segmentData, EXIF_IDENTIFIER_CODE))
                {
                    pieces.add(new JFIFPieceSegment(marker, markerBytes,
                            markerLengthBytes, segmentData));
                }
                //                else if (exifSegmentArray[0] != null)
                //                {
                //                    // TODO: add support for multiple segments
                //                    throw new ImageReadException(
                //                            "More than one APP1 EXIF segment.");
                //                }
                else
                {
                    JFIFPiece piece = new JFIFPieceSegmentExif(marker,
                            markerBytes, markerLengthBytes, segmentData);
                    pieces.add(piece);
                    exifPieces.add(piece);
                }
                return true;
            }
        };

        new JpegUtils().traverseJFIF(byteSource, visitor);

        //        GenericSegment exifSegment = exifSegmentArray[0];
        //        if (exifSegments.size() < 1)
        //        {
        //            // TODO: add support for adding, not just replacing.
        //            throw new ImageReadException("No APP1 EXIF segment found.");
        //        }

        return new JFIFPieces(pieces, exifPieces);
    }

    /**
     * Reads a Jpeg image, removes all EXIF metadata (by removing the APP1 segment),
     * and writes the result to a stream.
     * <p>
     * @param  src  Image file.
     * @param  os  OutputStream to write the image to.
     *
     * @see java.io.File
     * @see java.io.OutputStream
     * @see java.io.File
     * @see java.io.OutputStream
     */
    public void removeExifMetadata(File src, OutputStream os)
            throws ImageReadException, IOException, ImageWriteException
    {
        ByteSource byteSource = new ByteSourceFile(src);
        removeExifMetadata(byteSource, os);
    }

    /**
     * Reads a Jpeg image, removes all EXIF metadata (by removing the APP1 segment),
     * and writes the result to a stream.
     * <p>
     * @param  src  Byte array containing Jpeg image data.
     * @param  os  OutputStream to write the image to.
     */
    public void removeExifMetadata(byte src[], OutputStream os)
            throws ImageReadException, IOException, ImageWriteException
    {
        ByteSource byteSource = new ByteSourceArray(src);
        removeExifMetadata(byteSource, os);
    }

    /**
     * Reads a Jpeg image, removes all EXIF metadata (by removing the APP1 segment),
     * and writes the result to a stream.
     * <p>
     * @param  src  InputStream containing Jpeg image data.
     * @param  os  OutputStream to write the image to.
     */
    public void removeExifMetadata(InputStream src, OutputStream os)
            throws ImageReadException, IOException, ImageWriteException
    {
        ByteSource byteSource = new ByteSourceInputStream(src, null);
        removeExifMetadata(byteSource, os);
    }

    /**
     * Reads a Jpeg image, removes all EXIF metadata (by removing the APP1 segment),
     * and writes the result to a stream.
     * <p>
     * @param  byteSource  ByteSource containing Jpeg image data.
     * @param  os  OutputStream to write the image to.
     */
    public void removeExifMetadata(ByteSource byteSource, OutputStream os)
            throws ImageReadException, IOException, ImageWriteException
    {
        JFIFPieces jfifPieces = analyzeJFIF(byteSource);
        List pieces = jfifPieces.pieces;

        //        Debug.debug("pieces", pieces);

        //        pieces.removeAll(jfifPieces.exifSegments);

        //        Debug.debug("pieces", pieces);

        writeSegmentsReplacingExif(os, pieces, null);
    }

    /**
     * Reads a Jpeg image, replaces the EXIF metadata and writes the result to a stream.
     * <p>
     * Note that this uses the "Lossless" approach - in order to preserve data embedded in the EXIF
     * segment that it can't parse (such as Maker Notes), this algorithm avoids overwriting
     * any part of the original segment that it couldn't parse.  This can cause the EXIF segment to
     * grow with each update, which is a serious issue, since all EXIF data must fit in a single APP1
     * segment of the Jpeg image.
     * <p>
     * @param  src  Image file.
     * @param  os  OutputStream to write the image to.
     * @param  outputSet  TiffOutputSet containing the EXIF data to write.
     */
    public void updateExifMetadataLossless(File src, OutputStream os,
            TiffOutputSet outputSet) throws ImageReadException, IOException,
            ImageWriteException
    {
        ByteSource byteSource = new ByteSourceFile(src);
        updateExifMetadataLossless(byteSource, os, outputSet);
    }

    /**
     * Reads a Jpeg image, replaces the EXIF metadata and writes the result to a stream.
     * <p>
     * Note that this uses the "Lossless" approach - in order to preserve data embedded in the EXIF
     * segment that it can't parse (such as Maker Notes), this algorithm avoids overwriting
     * any part of the original segment that it couldn't parse.  This can cause the EXIF segment to
     * grow with each update, which is a serious issue, since all EXIF data must fit in a single APP1
     * segment of the Jpeg image.
     * <p>
     * @param  src  Byte array containing Jpeg image data.
     * @param  os  OutputStream to write the image to.
     * @param  outputSet  TiffOutputSet containing the EXIF data to write.
     */
    public void updateExifMetadataLossless(byte src[], OutputStream os,
            TiffOutputSet outputSet) throws ImageReadException, IOException,
            ImageWriteException
    {
        ByteSource byteSource = new ByteSourceArray(src);
        updateExifMetadataLossless(byteSource, os, outputSet);
    }

    /**
     * Reads a Jpeg image, replaces the EXIF metadata and writes the result to a stream.
     * <p>
     * Note that this uses the "Lossless" approach - in order to preserve data embedded in the EXIF
     * segment that it can't parse (such as Maker Notes), this algorithm avoids overwriting
     * any part of the original segment that it couldn't parse.  This can cause the EXIF segment to
     * grow with each update, which is a serious issue, since all EXIF data must fit in a single APP1
     * segment of the Jpeg image.
     * <p>
     * @param  src  InputStream containing Jpeg image data.
     * @param  os  OutputStream to write the image to.
     * @param  outputSet  TiffOutputSet containing the EXIF data to write.
     */
    public void updateExifMetadataLossless(InputStream src, OutputStream os,
            TiffOutputSet outputSet) throws ImageReadException, IOException,
            ImageWriteException
    {
        ByteSource byteSource = new ByteSourceInputStream(src, null);
        updateExifMetadataLossless(byteSource, os, outputSet);
    }

    /**
     * Reads a Jpeg image, replaces the EXIF metadata and writes the result to a stream.
     * <p>
     * Note that this uses the "Lossless" approach - in order to preserve data embedded in the EXIF
     * segment that it can't parse (such as Maker Notes), this algorithm avoids overwriting
     * any part of the original segment that it couldn't parse.  This can cause the EXIF segment to
     * grow with each update, which is a serious issue, since all EXIF data must fit in a single APP1
     * segment of the Jpeg image.
     * <p>
     * @param  byteSource  ByteSource containing Jpeg image data.
     * @param  os  OutputStream to write the image to.
     * @param  outputSet  TiffOutputSet containing the EXIF data to write.
     */
    public void updateExifMetadataLossless(ByteSource byteSource,
            OutputStream os, TiffOutputSet outputSet)
            throws ImageReadException, IOException, ImageWriteException
    {
        //        List outputDirectories = outputSet.getDirectories();
        JFIFPieces jfifPieces = analyzeJFIF(byteSource);
        List pieces = jfifPieces.pieces;

        TiffImageWriterBase writer;
        // Just use first APP1 segment for now.
        // Multiple APP1 segments are rare and poorly supported.
        if (jfifPieces.exifPieces.size() > 0)
        {
            JFIFPieceSegment exifPiece = null;
            exifPiece = (JFIFPieceSegment) jfifPieces.exifPieces.get(0);

            byte exifBytes[] = exifPiece.segmentData;
            exifBytes = getByteArrayTail("trimmed exif bytes", exifBytes, 6);

            writer = new TiffImageWriterLossless(outputSet.byteOrder, exifBytes);

        }
        else
            writer = new TiffImageWriterLossy(outputSet.byteOrder);

        boolean includeEXIFPrefix = true;
        byte newBytes[] = writeExifSegment(writer, outputSet, includeEXIFPrefix);

        writeSegmentsReplacingExif(os, pieces, newBytes);
    }

    /**
     * Reads a Jpeg image, replaces the EXIF metadata and writes the result to a stream.
     * <p>
     * Note that this uses the "Lossy" approach - the algorithm overwrites the entire EXIF segment,
     * ignoring the possibility that it may be discarding data it couldn't parse (such as Maker Notes).
     * <p>
     * @param  src  Byte array containing Jpeg image data.
     * @param  os  OutputStream to write the image to.
     * @param  outputSet  TiffOutputSet containing the EXIF data to write.
     */
    public void updateExifMetadataLossy(byte src[], OutputStream os,
            TiffOutputSet outputSet) throws ImageReadException, IOException,
            ImageWriteException
    {
        ByteSource byteSource = new ByteSourceArray(src);
        updateExifMetadataLossy(byteSource, os, outputSet);
    }

    /**
     * Reads a Jpeg image, replaces the EXIF metadata and writes the result to a stream.
     * <p>
     * Note that this uses the "Lossy" approach - the algorithm overwrites the entire EXIF segment,
     * ignoring the possibility that it may be discarding data it couldn't parse (such as Maker Notes).
     * <p>
     * @param  src  InputStream containing Jpeg image data.
     * @param  os  OutputStream to write the image to.
     * @param  outputSet  TiffOutputSet containing the EXIF data to write.
     */
    public void updateExifMetadataLossy(InputStream src, OutputStream os,
            TiffOutputSet outputSet) throws ImageReadException, IOException,
            ImageWriteException
    {
        ByteSource byteSource = new ByteSourceInputStream(src, null);
        updateExifMetadataLossy(byteSource, os, outputSet);
    }

    /**
     * Reads a Jpeg image, replaces the EXIF metadata and writes the result to a stream.
     * <p>
     * Note that this uses the "Lossy" approach - the algorithm overwrites the entire EXIF segment,
     * ignoring the possibility that it may be discarding data it couldn't parse (such as Maker Notes).
     * <p>
     * @param  src  Image file.
     * @param  os  OutputStream to write the image to.
     * @param  outputSet  TiffOutputSet containing the EXIF data to write.
     */
    public void updateExifMetadataLossy(File src, OutputStream os,
            TiffOutputSet outputSet) throws ImageReadException, IOException,
            ImageWriteException
    {
        ByteSource byteSource = new ByteSourceFile(src);
        updateExifMetadataLossy(byteSource, os, outputSet);
    }

    /**
     * Reads a Jpeg image, replaces the EXIF metadata and writes the result to a stream.
     * <p>
     * Note that this uses the "Lossy" approach - the algorithm overwrites the entire EXIF segment,
     * ignoring the possibility that it may be discarding data it couldn't parse (such as Maker Notes).
     * <p>
     * @param  byteSource  ByteSource containing Jpeg image data.
     * @param  os  OutputStream to write the image to.
     * @param  outputSet  TiffOutputSet containing the EXIF data to write.
     */
    public void updateExifMetadataLossy(ByteSource byteSource, OutputStream os,
            TiffOutputSet outputSet) throws ImageReadException, IOException,
            ImageWriteException
    {
        JFIFPieces jfifPieces = analyzeJFIF(byteSource);
        List pieces = jfifPieces.pieces;

        TiffImageWriterBase writer = new TiffImageWriterLossy(
                outputSet.byteOrder);

        boolean includeEXIFPrefix = true;
        byte newBytes[] = writeExifSegment(writer, outputSet, includeEXIFPrefix);

        writeSegmentsReplacingExif(os, pieces, newBytes);
    }

    private void writeSegmentsReplacingExif(OutputStream os, List segments,
            byte newBytes[]) throws ImageWriteException, IOException
    {
        int byteOrder = getByteOrder();

        try
        {
            os.write(SOI);

            boolean hasExif = false;

            for (int i = 0; i < segments.size(); i++)
            {
                JFIFPiece piece = (JFIFPiece) segments.get(i);
                if (piece instanceof JFIFPieceSegmentExif)
                    hasExif = true;
            }

            if (!hasExif && newBytes != null)
            {
                byte markerBytes[] = convertShortToByteArray(JPEG_APP1_Marker,
                        byteOrder);
                if (newBytes.length > 0xffff)
                    throw new ExifOverflowException(
                            "APP1 Segment is too long: " + newBytes.length);
                int markerLength = newBytes.length + 2;
                byte markerLengthBytes[] = convertShortToByteArray(
                        markerLength, byteOrder);

                int index = 0;
                JFIFPieceSegment firstSegment = (JFIFPieceSegment) segments
                        .get(index);
                if (firstSegment.marker == JFIFMarker)
                    index = 1;
                segments.add(0, new JFIFPieceSegmentExif(JPEG_APP1_Marker,
                        markerBytes, markerLengthBytes, newBytes));
            }

            boolean APP1Written = false;

            for (int i = 0; i < segments.size(); i++)
            {
                JFIFPiece piece = (JFIFPiece) segments.get(i);
                if (piece instanceof JFIFPieceSegmentExif)
                {
                    // only replace first APP1 segment; skips others.
                    if (APP1Written)
                        continue;
                    APP1Written = true;

                    if (newBytes == null)
                        continue;

                    byte markerBytes[] = convertShortToByteArray(
                            JPEG_APP1_Marker, byteOrder);
                    if (newBytes.length > 0xffff)
                        throw new ExifOverflowException(
                                "APP1 Segment is too long: " + newBytes.length);
                    int markerLength = newBytes.length + 2;
                    byte markerLengthBytes[] = convertShortToByteArray(
                            markerLength, byteOrder);

                    os.write(markerBytes);
                    os.write(markerLengthBytes);
                    os.write(newBytes);
                }
                else
                {
                    piece.write(os);
                }
            }
        }
        finally
        {
            try
            {
                os.close();
            }
            catch (Exception e)
            {
                Debug.debug(e);
            }
        }
    }

    public static class ExifOverflowException extends ImageWriteException
    {
        public ExifOverflowException(String s)
        {
            super(s);
        }
    }

    private byte[] writeExifSegment(TiffImageWriterBase writer,
            TiffOutputSet outputSet, boolean includeEXIFPrefix)
            throws IOException, ImageWriteException
    {
        ByteArrayOutputStream os = new ByteArrayOutputStream();

        if (includeEXIFPrefix)
        {
            os.write(EXIF_IDENTIFIER_CODE);
            os.write(0);
            os.write(0);
        }

        writer.write(os, outputSet);

        return os.toByteArray();
    }

}