/*
    * 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.commons.imaging.formats.jpeg.exif;
  
  import static org.apache.commons.imaging.common.BinaryFunctions.remainingBytes;
  import static org.apache.commons.imaging.common.BinaryFunctions.startsWith;
  
  import java.io.ByteArrayOutputStream;
  import java.io.DataOutputStream;
  import java.io.File;
  import java.io.IOException;
  import java.io.InputStream;
  import java.io.OutputStream;
  import java.nio.ByteOrder;
  import java.util.ArrayList;
  import java.util.List;
  
  import org.apache.commons.imaging.ImagingException;
  import org.apache.commons.imaging.ImagingOverflowException;
  import org.apache.commons.imaging.bytesource.ByteSource;
  import org.apache.commons.imaging.common.BinaryFileParser;
  import org.apache.commons.imaging.common.ByteConversions;
  import org.apache.commons.imaging.formats.jpeg.JpegConstants;
  import org.apache.commons.imaging.formats.jpeg.JpegUtils;
  import org.apache.commons.imaging.formats.tiff.write.AbstractTiffImageWriter;
  import org.apache.commons.imaging.formats.tiff.write.TiffImageWriterLossless;
  import org.apache.commons.imaging.formats.tiff.write.TiffImageWriterLossy;
  import org.apache.commons.imaging.formats.tiff.write.TiffOutputSet;
  
  /**
   * Interface for Exif write/update/remove functionality for Jpeg/JFIF images.
   *
   * <p>
   * See the source of the ExifMetadataUpdateExample class for example usage.
   * </p>
   *
   * @see <a href=
   *      "https://svn.apache.org/repos/asf/commons/proper/imaging/trunk/src/test/java/org/apache/commons/imaging/examples/WriteExifMetadataExample.java">
   *      org.apache.commons.imaging.examples.WriteExifMetadataExample</a>
   */
  public class ExifRewriter extends BinaryFileParser {
  
      private abstract static class JFIFPiece {
          protected abstract void write(OutputStream os) throws IOException;
      }
  
      private static final class JFIFPieceImageData extends JFIFPiece {
          public final byte[] markerBytes;
          public final byte[] imageData;
  
          JFIFPieceImageData(final byte[] markerBytes, final byte[] imageData) {
              this.markerBytes = markerBytes;
              this.imageData = imageData;
          }
  
          @Override
          protected void write(final OutputStream os) throws IOException {
              os.write(markerBytes);
              os.write(imageData);
          }
      }
  
      private static final class JFIFPieces {
          public final List<JFIFPiece> pieces;
          public final List<JFIFPiece> exifPieces;
  
          JFIFPieces(final List<JFIFPiece> pieces, final List<JFIFPiece> exifPieces) {
              this.pieces = pieces;
              this.exifPieces = exifPieces;
          }
  
      }
  
      private static class JFIFPieceSegment extends JFIFPiece {
          public final int marker;
          public final byte[] markerBytes;
          public final byte[] markerLengthBytes;
          public final byte[] segmentData;
  
          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;
          }
 
         @Override
         protected void write(final OutputStream os) throws IOException {
             os.write(markerBytes);
             os.write(markerLengthBytes);
             os.write(segmentData);
         }
     }
 
     private static final class JFIFPieceSegmentExif extends JFIFPieceSegment {
 
         JFIFPieceSegmentExif(final int marker, final byte[] markerBytes, final byte[] markerLengthBytes, final byte[] segmentData) {
             super(marker, markerBytes, markerLengthBytes, segmentData);
         }
     }
 
     /**
      * Constructs a new instance. to guess whether a file contains an image based on its file extension.
      */
     public ExifRewriter() {
         this(ByteOrder.BIG_ENDIAN);
     }
 
     /**
      * Constructs a new instance.
      * <p>
      *
      * @param byteOrder byte order of EXIF segment.
      */
     public ExifRewriter(final ByteOrder byteOrder) {
         super(byteOrder);
     }
 
     private JFIFPieces analyzeJfif(final ByteSource byteSource) throws ImagingException, IOException {
         final List<JFIFPiece> pieces = new ArrayList<>();
         final List<JFIFPiece> exifPieces = new ArrayList<>();
 
         final JpegUtils.Visitor visitor = new JpegUtils.Visitor() {
             // return false to exit before reading image data.
             @Override
             public boolean beginSos() {
                 return true;
             }
 
             // return false to exit traversal.
             @Override
             public boolean visitSegment(final int marker, final byte[] markerBytes, final int markerLength, final byte[] markerLengthBytes,
                     final byte[] segmentData) throws
             // ImageWriteException,
             ImagingException, IOException {
                 if (marker != JpegConstants.JPEG_APP1_MARKER || !startsWith(segmentData, JpegConstants.EXIF_IDENTIFIER_CODE)) {
                     pieces.add(new JFIFPieceSegment(marker, markerBytes, markerLengthBytes, segmentData));
                 } else {
                     final JFIFPiece piece = new JFIFPieceSegmentExif(marker, markerBytes, markerLengthBytes, segmentData);
                     pieces.add(piece);
                     exifPieces.add(piece);
                 }
                 return true;
             }
 
             @Override
             public void visitSos(final int marker, final byte[] markerBytes, final byte[] imageData) {
                 pieces.add(new JFIFPieceImageData(markerBytes, imageData));
             }
         };
 
         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.
      *
      * @param src Byte array containing JPEG image data.
      * @param os  OutputStream to write the image to.
      * @throws ImagingException if it fails to read the JFIF segments
      * @throws IOException      if it fails to read the image data
      * @throws ImagingException if it fails to write the updated data
      */
     public void removeExifMetadata(final byte[] src, final OutputStream os) throws ImagingException, IOException, ImagingException {
         final ByteSource byteSource = ByteSource.array(src);
         removeExifMetadata(byteSource, os);
     }
 
     /**
      * Reads a JPEG image, removes all EXIF metadata (by removing the APP1 segment), and writes the result to a stream.
      *
      * @param byteSource ByteSource containing JPEG image data.
      * @param os         OutputStream to write the image to.
      * @throws ImagingException if it fails to read the JFIF segments
      * @throws IOException      if it fails to read the image data
      * @throws ImagingException if it fails to write the updated data
      */
     public void removeExifMetadata(final ByteSource byteSource, final OutputStream os) throws ImagingException, IOException, ImagingException {
         final JFIFPieces jfifPieces = analyzeJfif(byteSource);
         final List<JFIFPiece> pieces = jfifPieces.pieces;
 
         // Debug.debug("pieces", pieces);
 
         // pieces.removeAll(jfifPieces.exifSegments);
 
         // Debug.debug("pieces", pieces);
 
         writeSegmentsReplacingExif(os, pieces, null);
     }
 
     /**
      * 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.
      *
      * @throws ImagingException if it fails to read the JFIF segments
      * @throws IOException      if it fails to read the image data
      * @throws ImagingException if it fails to write the updated data
      * @see java.io.File
      * @see java.io.OutputStream
      * @see java.io.File
      * @see java.io.OutputStream
      */
     public void removeExifMetadata(final File src, final OutputStream os) throws ImagingException, IOException, ImagingException {
         final ByteSource byteSource = ByteSource.file(src);
         removeExifMetadata(byteSource, os);
     }
 
     /**
      * Reads a JPEG image, removes all EXIF metadata (by removing the APP1 segment), and writes the result to a stream.
      *
      * @param src InputStream containing JPEG image data.
      * @param os  OutputStream to write the image to.
      * @throws ImagingException if it fails to read the JFIF segments
      * @throws IOException      if it fails to read the image data
      * @throws ImagingException if it fails to write the updated data
      */
     public void removeExifMetadata(final InputStream src, final OutputStream os) throws ImagingException, IOException, ImagingException {
         final ByteSource byteSource = ByteSource.inputStream(src, null);
         removeExifMetadata(byteSource, os);
     }
 
     /**
      * 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.
      * @throws ImagingException if it fails to read the JFIF segments
      * @throws IOException      if it fails to read the image data
      * @throws ImagingException if it fails to write the updated data
      */
     public void updateExifMetadataLossless(final byte[] src, final OutputStream os, final TiffOutputSet outputSet)
             throws ImagingException, IOException, ImagingException {
         final ByteSource byteSource = ByteSource.array(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 byteSource ByteSource containing JPEG image data.
      * @param os         OutputStream to write the image to.
      * @param outputSet  TiffOutputSet containing the EXIF data to write.
      * @throws ImagingException if it fails to read the JFIF segments
      * @throws IOException      if it fails to read the image data
      * @throws ImagingException if it fails to write the updated data
      */
     public void updateExifMetadataLossless(final ByteSource byteSource, final OutputStream os, final TiffOutputSet outputSet)
             throws ImagingException, IOException, ImagingException {
         // List outputDirectories = outputSet.getDirectories();
         final JFIFPieces jfifPieces = analyzeJfif(byteSource);
         final List<JFIFPiece> pieces = jfifPieces.pieces;
 
         AbstractTiffImageWriter writer;
         // Just use first APP1 segment for now.
         // Multiple APP1 segments are rare and poorly supported.
         if (!jfifPieces.exifPieces.isEmpty()) {
             final JFIFPieceSegment exifPiece = (JFIFPieceSegment) jfifPieces.exifPieces.get(0);
 
             byte[] exifBytes = exifPiece.segmentData;
             exifBytes = remainingBytes("trimmed exif bytes", exifBytes, 6);
 
             writer = new TiffImageWriterLossless(outputSet.byteOrder, exifBytes);
 
         } else {
             writer = new TiffImageWriterLossy(outputSet.byteOrder);
         }
 
         final boolean includeEXIFPrefix = true;
         final 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 "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.
      * @throws ImagingException if it fails to read the JFIF segments
      * @throws IOException      if it fails to read the image data
      * @throws ImagingException if it fails to write the updated data
      */
     public void updateExifMetadataLossless(final File src, final OutputStream os, final TiffOutputSet outputSet)
             throws ImagingException, IOException, ImagingException {
         final ByteSource byteSource = ByteSource.file(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.
      * @throws ImagingException if it fails to read the JFIF segments
      * @throws IOException      if it fails to read the image data
      * @throws ImagingException if it fails to write the updated data
      */
     public void updateExifMetadataLossless(final InputStream src, final OutputStream os, final TiffOutputSet outputSet)
             throws ImagingException, IOException, ImagingException {
         final ByteSource byteSource = ByteSource.inputStream(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 "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.
      * @throws ImagingException if it fails to read the JFIF segments
      * @throws IOException      if it fails to read the image data
      * @throws ImagingException if it fails to write the updated data
      */
     public void updateExifMetadataLossy(final byte[] src, final OutputStream os, final TiffOutputSet outputSet)
             throws ImagingException, IOException, ImagingException {
         final ByteSource byteSource = ByteSource.array(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.
      * @throws ImagingException if it fails to read the JFIF segments
      * @throws IOException      if it fails to read the image data
      * @throws ImagingException if it fails to write the updated data
      */
     public void updateExifMetadataLossy(final ByteSource byteSource, final OutputStream os, final TiffOutputSet outputSet)
             throws ImagingException, IOException, ImagingException {
         final JFIFPieces jfifPieces = analyzeJfif(byteSource);
         final List<JFIFPiece> pieces = jfifPieces.pieces;
 
         final AbstractTiffImageWriter writer = new TiffImageWriterLossy(outputSet.byteOrder);
 
         final boolean includeEXIFPrefix = true;
         final 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       Image file.
      * @param os        OutputStream to write the image to.
      * @param outputSet TiffOutputSet containing the EXIF data to write.
      * @throws ImagingException if it fails to read the JFIF segments
      * @throws IOException      if it fails to read the image data
      * @throws ImagingException if it fails to write the updated data
      */
     public void updateExifMetadataLossy(final File src, final OutputStream os, final TiffOutputSet outputSet)
             throws ImagingException, IOException, ImagingException {
         final ByteSource byteSource = ByteSource.file(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.
      * @throws ImagingException if it fails to read the JFIF segments
      * @throws IOException      if it fails to read the image data
      * @throws ImagingException if it fails to write the updated data
      */
     public void updateExifMetadataLossy(final InputStream src, final OutputStream os, final TiffOutputSet outputSet)
             throws ImagingException, IOException, ImagingException {
         final ByteSource byteSource = ByteSource.inputStream(src, null);
         updateExifMetadataLossy(byteSource, os, outputSet);
     }
 
     private byte[] writeExifSegment(final AbstractTiffImageWriter writer, final TiffOutputSet outputSet, final boolean includeEXIFPrefix)
             throws IOException, ImagingException {
         final ByteArrayOutputStream os = new ByteArrayOutputStream();
 
         if (includeEXIFPrefix) {
             JpegConstants.EXIF_IDENTIFIER_CODE.writeTo(os);
             os.write(0);
             os.write(0);
         }
 
         writer.write(os, outputSet);
 
         return os.toByteArray();
     }
 
     private void writeSegmentsReplacingExif(final OutputStream outputStream, final List<JFIFPiece> segments, final byte[] newBytes)
             throws ImagingException, IOException {
 
         try (DataOutputStream os = new DataOutputStream(outputStream)) {
             JpegConstants.SOI.writeTo(os);
 
             boolean hasExif = false;
 
             for (final JFIFPiece piece : segments) {
                 if (piece instanceof JFIFPieceSegmentExif) {
                     hasExif = true;
                     break;
                 }
             }
 
             if (!hasExif && newBytes != null) {
                 final byte[] markerBytes = ByteConversions.toBytes((short) JpegConstants.JPEG_APP1_MARKER, getByteOrder());
                 if (newBytes.length > 0xffff) {
                     throw new ImagingOverflowException("APP1 Segment is too long: " + newBytes.length);
                 }
                 final int markerLength = newBytes.length + 2;
                 final byte[] markerLengthBytes = ByteConversions.toBytes((short) markerLength, getByteOrder());
 
                 int index = 0;
                 final JFIFPieceSegment firstSegment = (JFIFPieceSegment) segments.get(index);
                 if (firstSegment.marker == JpegConstants.JFIF_MARKER) {
                     index = 1;
                 }
                 segments.add(index, new JFIFPieceSegmentExif(JpegConstants.JPEG_APP1_MARKER, markerBytes, markerLengthBytes, newBytes));
             }
 
             boolean APP1Written = false;
 
             for (final JFIFPiece piece : segments) {
                 if (piece instanceof JFIFPieceSegmentExif) {
                     // only replace first APP1 segment; skips others.
                     if (APP1Written) {
                         continue;
                     }
                     APP1Written = true;
 
                     if (newBytes == null) {
                         continue;
                     }
 
                     final byte[] markerBytes = ByteConversions.toBytes((short) JpegConstants.JPEG_APP1_MARKER, getByteOrder());
                     if (newBytes.length > 0xffff) {
                         throw new ImagingOverflowException("APP1 Segment is too long: " + newBytes.length);
                     }
                     final int markerLength = newBytes.length + 2;
                     final byte[] markerLengthBytes = ByteConversions.toBytes((short) markerLength, getByteOrder());
 
                     os.write(markerBytes);
                     os.write(markerLengthBytes);
                     os.write(newBytes);
                 } else {
                     piece.write(os);
                 }
             }
         }
     }
 
 }