/*
 * 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.xmp;

import static org.apache.commons.imaging.common.BinaryFunctions.startsWith;

import java.io.DataOutputStream;
import java.io.IOException;
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.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.jpeg.iptc.IptcParser;

/**
 * Interface for Exif write/update/remove functionality for Jpeg/JFIF images.
 */
public class JpegRewriter extends BinaryFileParser {
    protected abstract static class JFIFPiece {
        @Override
        public String toString() {
            return "[" + this.getClass().getName() + "]";
        }

        protected abstract void write(OutputStream os) throws IOException;
    }

    static class JFIFPieceImageData extends JFIFPiece {
        private final byte[] markerBytes;
        private 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);
        }
    }

    protected static class JFIFPieces {
        public final List<JFIFPiece> pieces;
        public final List<JFIFPiece> segmentPieces;

        public JFIFPieces(final List<JFIFPiece> pieces, final List<JFIFPiece> segmentPieces) {
            this.pieces = pieces;
            this.segmentPieces = segmentPieces;
        }

    }

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

        public JFIFPieceSegment(final int marker, final byte[] segmentData) {
            this(marker, ByteConversions.toBytes((short) marker, JPEG_BYTE_ORDER), ByteConversions.toBytes((short) (segmentData.length + 2), JPEG_BYTE_ORDER),
                    segmentData);
        }

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

        public byte[] getSegmentData() {
            return segmentData.clone();
        }

        public boolean isApp1Segment() {
            return marker == JpegConstants.JPEG_APP1_MARKER;
        }

        public boolean isAppSegment() {
            return marker >= JpegConstants.JPEG_APP0_MARKER && marker <= JpegConstants.JPEG_APP15_MARKER;
        }

        public boolean isExifSegment() {
            if (marker != JpegConstants.JPEG_APP1_MARKER) {
                return false;
            }
            if (!startsWith(segmentData, JpegConstants.EXIF_IDENTIFIER_CODE)) {
                return false;
            }
            return true;
        }

        public boolean isPhotoshopApp13Segment() {
            if (marker != JpegConstants.JPEG_APP13_MARKER) {
                return false;
            }
            if (!new IptcParser().isPhotoshopJpegSegment(segmentData)) {
                return false;
            }
            return true;
        }

        public boolean isXmpSegment() {
            if (marker != JpegConstants.JPEG_APP1_MARKER) {
                return false;
            }
            if (!startsWith(segmentData, JpegConstants.XMP_IDENTIFIER)) {
                return false;
            }
            return true;
        }

        @Override
        public String toString() {
            return "[" + this.getClass().getName() + " (0x" + Integer.toHexString(marker) + ")]";
        }

        @Override
        protected void write(final OutputStream os) throws IOException {
            os.write(markerBytes);
            os.write(segmentLengthBytes);
            os.write(segmentData);
        }

    }

    private interface SegmentFilter {
        boolean filter(JFIFPieceSegment segment);
    }

    private static final ByteOrder JPEG_BYTE_ORDER = ByteOrder.BIG_ENDIAN;

    private static final SegmentFilter EXIF_SEGMENT_FILTER = JFIFPieceSegment::isExifSegment;

    private static final SegmentFilter XMP_SEGMENT_FILTER = JFIFPieceSegment::isXmpSegment;

    private static final SegmentFilter PHOTOSHOP_APP13_SEGMENT_FILTER = JFIFPieceSegment::isPhotoshopApp13Segment;

    /**
     * Constructs a new instance. to guess whether a file contains an image based on its file extension.
     */
    public JpegRewriter() {
        super(JPEG_BYTE_ORDER);
    }

    protected JFIFPieces analyzeJfif(final ByteSource byteSource) throws ImagingException, IOException {
        final List<JFIFPiece> pieces = new ArrayList<>();
        final List<JFIFPiece> segmentPieces = 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 segmentLength, final byte[] segmentLengthBytes,
                    final byte[] segmentData) throws ImagingException, IOException {
                final JFIFPiece piece = new JFIFPieceSegment(marker, markerBytes, segmentLengthBytes, segmentData);
                pieces.add(piece);
                segmentPieces.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);

        return new JFIFPieces(pieces, segmentPieces);
    }

    protected <T extends JFIFPiece> List<T> filterSegments(final List<T> segments, final SegmentFilter filter) {
        return filterSegments(segments, filter, false);
    }

    protected <T extends JFIFPiece> List<T> filterSegments(final List<T> segments, final SegmentFilter filter, final boolean reverse) {
        final List<T> result = new ArrayList<>();

        for (final T piece : segments) {
            if (piece instanceof JFIFPieceSegment) {
                if (filter.filter((JFIFPieceSegment) piece) == reverse) {
                    result.add(piece);
                }
            } else if (!reverse) {
                result.add(piece);
            }
        }

        return result;
    }

    protected <T extends JFIFPiece> List<T> findPhotoshopApp13Segments(final List<T> segments) {
        return filterSegments(segments, PHOTOSHOP_APP13_SEGMENT_FILTER, true);
    }

    protected <T extends JFIFPiece, U extends JFIFPiece> List<JFIFPiece> insertAfterLastAppSegments(final List<T> segments, final List<U> newSegments)
            throws ImagingException {
        int lastAppIndex = -1;
        for (int i = 0; i < segments.size(); i++) {
            final JFIFPiece piece = segments.get(i);
            if (!(piece instanceof JFIFPieceSegment)) {
                continue;
            }

            final JFIFPieceSegment segment = (JFIFPieceSegment) piece;
            if (segment.isAppSegment()) {
                lastAppIndex = i;
            }
        }

        final List<JFIFPiece> result = new ArrayList<>(segments);
        if (lastAppIndex == -1) {
            if (segments.isEmpty()) {
                throw new ImagingException("JPEG file has no APP segments.");
            }
            result.addAll(1, newSegments);
        } else {
            result.addAll(lastAppIndex + 1, newSegments);
        }

        return result;
    }

    protected <T extends JFIFPiece, U extends JFIFPiece> List<JFIFPiece> insertBeforeFirstAppSegments(final List<T> segments, final List<U> newSegments)
            throws ImagingException {
        int firstAppIndex = -1;
        for (int i = 0; i < segments.size(); i++) {
            final JFIFPiece piece = segments.get(i);
            if (!(piece instanceof JFIFPieceSegment)) {
                continue;
            }

            final JFIFPieceSegment segment = (JFIFPieceSegment) piece;
            if (segment.isAppSegment()) {
                if (firstAppIndex == -1) {
                    firstAppIndex = i;
                }
            }
        }

        final List<JFIFPiece> result = new ArrayList<>(segments);
        if (firstAppIndex == -1) {
            throw new ImagingException("JPEG file has no APP segments.");
        }
        result.addAll(firstAppIndex, newSegments);
        return result;
    }

    protected <T extends JFIFPiece> List<T> removeExifSegments(final List<T> segments) {
        return filterSegments(segments, EXIF_SEGMENT_FILTER);
    }

    protected <T extends JFIFPiece> List<T> removePhotoshopApp13Segments(final List<T> segments) {
        return filterSegments(segments, PHOTOSHOP_APP13_SEGMENT_FILTER);
    }

    protected <T extends JFIFPiece> List<T> removeXmpSegments(final List<T> segments) {
        return filterSegments(segments, XMP_SEGMENT_FILTER);
    }

    // private void writeSegment(OutputStream os, JFIFPieceSegment piece)
    // throws ImageWriteException, IOException
    // {
    // byte[] markerBytes = convertShortToByteArray(JPEG_APP1_MARKER,
    // JPEG_BYTE_ORDER);
    // if (piece.segmentData.length > 0xffff)
    // throw new JpegSegmentOverflowException("JPEG segment is too long: "
    // + piece.segmentData.length);
    // int segmentLength = piece.segmentData.length + 2;
    // byte[] segmentLengthBytes = convertShortToByteArray(segmentLength,
    // JPEG_BYTE_ORDER);
    //
    // os.write(markerBytes);
    // os.write(segmentLengthBytes);
    // os.write(piece.segmentData);
    // }

    protected void writeSegments(final OutputStream outputStream, final List<? extends JFIFPiece> segments) throws IOException {
        try (DataOutputStream os = new DataOutputStream(outputStream)) {
            JpegConstants.SOI.writeTo(os);

            for (final JFIFPiece piece : segments) {
                piece.write(os);
            }
        }
    }

}
