package com.zcloud.basic.info.domain.utils;

import com.alibaba.cola.exception.BizException;
import org.springframework.web.multipart.MultipartFile;

import javax.imageio.IIOImage;
import javax.imageio.ImageIO;
import javax.imageio.ImageWriteParam;
import javax.imageio.ImageWriter;
import javax.imageio.stream.ImageInputStream;
import javax.imageio.stream.ImageOutputStream;
import java.awt.*;
import java.awt.geom.AffineTransform;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.util.Iterator;
import java.util.Locale;

public class ImageCompressUtil {

    private static final String JPG_SUFFIX = ".jpg";
    private static final String JPG_CONTENT_TYPE = "image/jpeg";
    private static final long DEFAULT_MAX_SIZE_KB = 100L;
    private static final float MIN_QUALITY = 0.05F;
    private static final float MAX_QUALITY = 1.0F;
    private static final int QUALITY_SEARCH_TIMES = 8;
    private static final double RESIZE_SCALE = 0.9D;
    private static final int EXIF_ORIENTATION_NORMAL = 1;
    private static final int EXIF_ORIENTATION_FLIP_HORIZONTAL = 2;
    private static final int EXIF_ORIENTATION_ROTATE_180 = 3;
    private static final int EXIF_ORIENTATION_FLIP_VERTICAL = 4;
    private static final int EXIF_ORIENTATION_TRANSPOSE = 5;
    private static final int EXIF_ORIENTATION_ROTATE_90 = 6;
    private static final int EXIF_ORIENTATION_TRANSVERSE = 7;
    private static final int EXIF_ORIENTATION_ROTATE_270 = 8;
    private static final int JPEG_SOI_MARKER = 0xD8;
    private static final int JPEG_EOI_MARKER = 0xD9;
    private static final int JPEG_SOS_MARKER = 0xDA;
    private static final int JPEG_APP1_MARKER = 0xE1;
    private static final int TIFF_HEADER_LENGTH = 8;
    private static final int IFD_ENTRY_LENGTH = 12;
    private static final int EXIF_ORIENTATION_TAG = 0x0112;
    private static final int TIFF_TYPE_SHORT = 3;
    private static final byte[] EXIF_HEADER = new byte[]{'E', 'x', 'i', 'f', 0, 0};

    private ImageCompressUtil() {
    }

    public static MultipartFile compressToJpg(MultipartFile file) throws IOException {
        return compressToJpg(file, DEFAULT_MAX_SIZE_KB);
    }

    public static MultipartFile compressToJpg(MultipartFile file, long maxSizeKb) {
        try {
            if (file == null || file.isEmpty()) {
                throw new IllegalArgumentException("file can not be null or empty");
            }
            if (maxSizeKb <= 0) {
                throw new IllegalArgumentException("maxSizeKb must be greater than 0");
            }

            long maxSizeBytes = maxSizeKb * 1024L;
            byte[] sourceBytes = file.getBytes();
            String formatName = getFormatName(sourceBytes);
            int orientation = getOrientation(sourceBytes);
            if (isTargetJpg(file, formatName, maxSizeBytes) && orientation == EXIF_ORIENTATION_NORMAL) {
                return file;
            }

            BufferedImage bufferedImage = ImageIO.read(new ByteArrayInputStream(sourceBytes));
            if (bufferedImage == null) {
                throw new IllegalArgumentException("file is not a supported image");
            }

            // 手机拍照图片经常依赖 EXIF 方向信息展示，压缩前先转正，避免保存后变成横图。
            BufferedImage orientedImage = applyOrientation(bufferedImage, orientation);
            BufferedImage rgbImage = toRgbImage(orientedImage);
            byte[] compressedBytes = compress(rgbImage, maxSizeBytes);
            String targetFileName = toJpgFileName(file.getOriginalFilename(), file.getName());
            return new ByteArrayMultipartFile(file.getName(), targetFileName, JPG_CONTENT_TYPE, compressedBytes);
        } catch (Exception e){
            e.printStackTrace();
            throw new BizException(e.getMessage());
        }

    }

    private static boolean isTargetJpg(MultipartFile file, String formatName, long maxSizeBytes) {
        return file.getSize() <= maxSizeBytes
                && isJpgFormat(formatName)
                && hasJpgSuffix(file.getOriginalFilename());
    }

    private static boolean isJpgFormat(String formatName) {
        if (formatName == null) {
            return false;
        }
        String lowerFormatName = formatName.toLowerCase(Locale.ROOT);
        return "jpg".equals(lowerFormatName) || "jpeg".equals(lowerFormatName);
    }

    private static boolean hasJpgSuffix(String fileName) {
        if (fileName == null) {
            return false;
        }
        return fileName.toLowerCase(Locale.ROOT).endsWith(JPG_SUFFIX);
    }

    private static String getFormatName(byte[] sourceBytes) throws IOException {
        try (ImageInputStream imageInputStream = ImageIO.createImageInputStream(new ByteArrayInputStream(sourceBytes))) {
            Iterator<ImageWriter> writers = ImageIO.getImageWritersByFormatName("jpg");
            if (!writers.hasNext()) {
                throw new IOException("jpg writer not found");
            }
            Iterator<javax.imageio.ImageReader> readers = ImageIO.getImageReaders(imageInputStream);
            if (!readers.hasNext()) {
                return null;
            }
            javax.imageio.ImageReader imageReader = readers.next();
            try {
                return imageReader.getFormatName();
            } finally {
                imageReader.dispose();
            }
        }
    }

    private static int getOrientation(byte[] sourceBytes) {
        try {
            if (!isJpeg(sourceBytes)) {
                return EXIF_ORIENTATION_NORMAL;
            }
            return readJpegOrientation(sourceBytes);
        } catch (Exception e) {
            return EXIF_ORIENTATION_NORMAL;
        }
    }

    private static boolean isJpeg(byte[] sourceBytes) {
        return sourceBytes != null
                && sourceBytes.length > 3
                && (sourceBytes[0] & 0xFF) == 0xFF
                && (sourceBytes[1] & 0xFF) == JPEG_SOI_MARKER;
    }

    private static int readJpegOrientation(byte[] sourceBytes) {
        int offset = 2;
        while (offset + 4 <= sourceBytes.length) {
            if ((sourceBytes[offset] & 0xFF) != 0xFF) {
                return EXIF_ORIENTATION_NORMAL;
            }

            int marker = sourceBytes[offset + 1] & 0xFF;
            offset += 2;
            if (marker == JPEG_SOI_MARKER) {
                continue;
            }
            if (marker == JPEG_EOI_MARKER || marker == JPEG_SOS_MARKER) {
                break;
            }
            if (offset + 2 > sourceBytes.length) {
                return EXIF_ORIENTATION_NORMAL;
            }

            int segmentLength = readUnsignedShort(sourceBytes, offset, false);
            if (segmentLength < 2 || offset + segmentLength > sourceBytes.length) {
                return EXIF_ORIENTATION_NORMAL;
            }

            int segmentDataOffset = offset + 2;
            int segmentDataLength = segmentLength - 2;
            if (marker == JPEG_APP1_MARKER && hasExifHeader(sourceBytes, segmentDataOffset, segmentDataLength)) {
                return readExifOrientation(sourceBytes, segmentDataOffset, segmentDataLength);
            }
            offset += segmentLength;
        }
        return EXIF_ORIENTATION_NORMAL;
    }

    private static boolean hasExifHeader(byte[] sourceBytes, int offset, int length) {
        if (length < EXIF_HEADER.length || offset + EXIF_HEADER.length > sourceBytes.length) {
            return false;
        }
        for (int i = 0; i < EXIF_HEADER.length; i++) {
            if (sourceBytes[offset + i] != EXIF_HEADER[i]) {
                return false;
            }
        }
        return true;
    }

    private static int readExifOrientation(byte[] sourceBytes, int exifOffset, int exifLength) {
        int tiffOffset = exifOffset + EXIF_HEADER.length;
        int tiffLimit = exifOffset + exifLength;
        if (tiffOffset + TIFF_HEADER_LENGTH > tiffLimit) {
            return EXIF_ORIENTATION_NORMAL;
        }

        Boolean littleEndian = resolveByteOrder(sourceBytes, tiffOffset);
        if (littleEndian == null) {
            return EXIF_ORIENTATION_NORMAL;
        }

        int firstIfdOffset = readUnsignedInt(sourceBytes, tiffOffset + 4, littleEndian);
        int ifdOffset = tiffOffset + firstIfdOffset;
        if (ifdOffset + 2 > tiffLimit) {
            return EXIF_ORIENTATION_NORMAL;
        }

        int entryCount = readUnsignedShort(sourceBytes, ifdOffset, littleEndian);
        int entryOffset = ifdOffset + 2;
        for (int i = 0; i < entryCount; i++) {
            int currentEntryOffset = entryOffset + i * IFD_ENTRY_LENGTH;
            if (currentEntryOffset + IFD_ENTRY_LENGTH > tiffLimit) {
                break;
            }

            int tag = readUnsignedShort(sourceBytes, currentEntryOffset, littleEndian);
            if (tag != EXIF_ORIENTATION_TAG) {
                continue;
            }

            int type = readUnsignedShort(sourceBytes, currentEntryOffset + 2, littleEndian);
            int count = readUnsignedInt(sourceBytes, currentEntryOffset + 4, littleEndian);
            if (type != TIFF_TYPE_SHORT || count < 1) {
                return EXIF_ORIENTATION_NORMAL;
            }

            int orientation = readUnsignedShort(sourceBytes, currentEntryOffset + 8, littleEndian);
            if (orientation < EXIF_ORIENTATION_NORMAL || orientation > EXIF_ORIENTATION_ROTATE_270) {
                return EXIF_ORIENTATION_NORMAL;
            }
            return orientation;
        }
        return EXIF_ORIENTATION_NORMAL;
    }

    private static Boolean resolveByteOrder(byte[] sourceBytes, int tiffOffset) {
        int first = sourceBytes[tiffOffset] & 0xFF;
        int second = sourceBytes[tiffOffset + 1] & 0xFF;
        if (first == 'I' && second == 'I') {
            return true;
        }
        if (first == 'M' && second == 'M') {
            return false;
        }
        return null;
    }

    private static int readUnsignedShort(byte[] sourceBytes, int offset, boolean littleEndian) {
        if (littleEndian) {
            return (sourceBytes[offset] & 0xFF) | ((sourceBytes[offset + 1] & 0xFF) << 8);
        }
        return ((sourceBytes[offset] & 0xFF) << 8) | (sourceBytes[offset + 1] & 0xFF);
    }

    private static int readUnsignedInt(byte[] sourceBytes, int offset, boolean littleEndian) {
        if (littleEndian) {
            return (sourceBytes[offset] & 0xFF)
                    | ((sourceBytes[offset + 1] & 0xFF) << 8)
                    | ((sourceBytes[offset + 2] & 0xFF) << 16)
                    | ((sourceBytes[offset + 3] & 0xFF) << 24);
        }
        return ((sourceBytes[offset] & 0xFF) << 24)
                | ((sourceBytes[offset + 1] & 0xFF) << 16)
                | ((sourceBytes[offset + 2] & 0xFF) << 8)
                | (sourceBytes[offset + 3] & 0xFF);
    }

    private static BufferedImage toRgbImage(BufferedImage sourceImage) {
        if (sourceImage.getType() == BufferedImage.TYPE_INT_RGB) {
            return sourceImage;
        }
        BufferedImage rgbImage = new BufferedImage(sourceImage.getWidth(), sourceImage.getHeight(), BufferedImage.TYPE_INT_RGB);
        Graphics2D graphics = rgbImage.createGraphics();
        try {
            graphics.setColor(Color.WHITE);
            graphics.fillRect(0, 0, sourceImage.getWidth(), sourceImage.getHeight());
            graphics.drawImage(sourceImage, 0, 0, null);
        } finally {
            graphics.dispose();
        }
        return rgbImage;
    }

    private static BufferedImage applyOrientation(BufferedImage sourceImage, int orientation) {
        if (orientation <= EXIF_ORIENTATION_NORMAL || orientation > EXIF_ORIENTATION_ROTATE_270) {
            return sourceImage;
        }

        int sourceWidth = sourceImage.getWidth();
        int sourceHeight = sourceImage.getHeight();
        boolean swapSize = orientation == EXIF_ORIENTATION_TRANSPOSE
                || orientation == EXIF_ORIENTATION_ROTATE_90
                || orientation == EXIF_ORIENTATION_TRANSVERSE
                || orientation == EXIF_ORIENTATION_ROTATE_270;
        int targetWidth = swapSize ? sourceHeight : sourceWidth;
        int targetHeight = swapSize ? sourceWidth : sourceHeight;
        BufferedImage targetImage = new BufferedImage(targetWidth, targetHeight, BufferedImage.TYPE_INT_RGB);
        Graphics2D graphics = targetImage.createGraphics();
        try {
            graphics.setColor(Color.WHITE);
            graphics.fillRect(0, 0, targetWidth, targetHeight);
            graphics.setRenderingHint(RenderingHints.KEY_INTERPOLATION, RenderingHints.VALUE_INTERPOLATION_BICUBIC);
            graphics.setRenderingHint(RenderingHints.KEY_RENDERING, RenderingHints.VALUE_RENDER_QUALITY);
            graphics.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
            graphics.drawImage(sourceImage, buildOrientationTransform(orientation, sourceWidth, sourceHeight), null);
        } finally {
            graphics.dispose();
        }
        return targetImage;
    }

    private static AffineTransform buildOrientationTransform(int orientation, int sourceWidth, int sourceHeight) {
        switch (orientation) {
            case EXIF_ORIENTATION_FLIP_HORIZONTAL:
                return new AffineTransform(-1, 0, 0, 1, sourceWidth, 0);
            case EXIF_ORIENTATION_ROTATE_180:
                return new AffineTransform(-1, 0, 0, -1, sourceWidth, sourceHeight);
            case EXIF_ORIENTATION_FLIP_VERTICAL:
                return new AffineTransform(1, 0, 0, -1, 0, sourceHeight);
            case EXIF_ORIENTATION_TRANSPOSE:
                return new AffineTransform(0, 1, 1, 0, 0, 0);
            case EXIF_ORIENTATION_ROTATE_90:
                return new AffineTransform(0, 1, -1, 0, sourceHeight, 0);
            case EXIF_ORIENTATION_TRANSVERSE:
                return new AffineTransform(0, -1, -1, 0, sourceHeight, sourceWidth);
            case EXIF_ORIENTATION_ROTATE_270:
                return new AffineTransform(0, -1, 1, 0, 0, sourceWidth);
            default:
                return new AffineTransform();
        }
    }

    private static byte[] compress(BufferedImage image, long maxSizeBytes) throws IOException {
        BufferedImage currentImage = image;
        while (true) {
            byte[] compressedBytes = searchBestQuality(currentImage, maxSizeBytes);
            if (compressedBytes != null) {
                return compressedBytes;
            }
            if (currentImage.getWidth() == 1 && currentImage.getHeight() == 1) {
                throw new IOException("can not compress image to target size");
            }
            currentImage = resize(currentImage, RESIZE_SCALE);
        }
    }

    private static byte[] searchBestQuality(BufferedImage image, long maxSizeBytes) throws IOException {
        byte[] maxQualityBytes = writeJpg(image, MAX_QUALITY);
        if (maxQualityBytes.length <= maxSizeBytes) {
            return maxQualityBytes;
        }

        byte[] minQualityBytes = writeJpg(image, MIN_QUALITY);
        if (minQualityBytes.length > maxSizeBytes) {
            return null;
        }

        byte[] bestBytes = minQualityBytes;
        float low = MIN_QUALITY;
        float high = MAX_QUALITY;
        for (int i = 0; i < QUALITY_SEARCH_TIMES; i++) {
            float mid = (low + high) / 2;
            byte[] currentBytes = writeJpg(image, mid);
            if (currentBytes.length <= maxSizeBytes) {
                bestBytes = currentBytes;
                low = mid;
            } else {
                high = mid;
            }
        }
        return bestBytes;
    }

    private static byte[] writeJpg(BufferedImage image, float quality) throws IOException {
        Iterator<ImageWriter> writers = ImageIO.getImageWritersByFormatName("jpg");
        if (!writers.hasNext()) {
            throw new IOException("jpg writer not found");
        }

        ImageWriter imageWriter = writers.next();
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try (ImageOutputStream imageOutputStream = ImageIO.createImageOutputStream(outputStream)) {
            imageWriter.setOutput(imageOutputStream);
            ImageWriteParam writeParam = imageWriter.getDefaultWriteParam();
            if (writeParam.canWriteCompressed()) {
                writeParam.setCompressionMode(ImageWriteParam.MODE_EXPLICIT);
                writeParam.setCompressionQuality(quality);
            }
            imageWriter.write(null, new IIOImage(image, null, null), writeParam);
        } finally {
            imageWriter.dispose();
        }
        return outputStream.toByteArray();
    }

    private static BufferedImage resize(BufferedImage sourceImage, double scale) {
        int targetWidth = Math.max(1, (int) Math.round(sourceImage.getWidth() * scale));
        int targetHeight = Math.max(1, (int) Math.round(sourceImage.getHeight() * scale));
        BufferedImage resizedImage = new BufferedImage(targetWidth, targetHeight, BufferedImage.TYPE_INT_RGB);
        Graphics2D graphics = resizedImage.createGraphics();
        try {
            graphics.setRenderingHint(RenderingHints.KEY_INTERPOLATION, RenderingHints.VALUE_INTERPOLATION_BICUBIC);
            graphics.setRenderingHint(RenderingHints.KEY_RENDERING, RenderingHints.VALUE_RENDER_QUALITY);
            graphics.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
            graphics.drawImage(sourceImage, 0, 0, targetWidth, targetHeight, null);
        } finally {
            graphics.dispose();
        }
        return resizedImage;
    }

    private static String toJpgFileName(String originalFileName, String fallbackName) {
        String fileName = originalFileName;
        if (fileName == null || fileName.trim().isEmpty()) {
            fileName = fallbackName;
        }
        if (fileName == null || fileName.trim().isEmpty()) {
            fileName = "image";
        }
        int dotIndex = fileName.lastIndexOf('.');
        if (dotIndex > -1) {
            fileName = fileName.substring(0, dotIndex);
        }
        return fileName + JPG_SUFFIX;
    }

    private static class ByteArrayMultipartFile implements MultipartFile {

        private final String name;
        private final String originalFilename;
        private final String contentType;
        private final byte[] content;

        private ByteArrayMultipartFile(String name, String originalFilename, String contentType, byte[] content) {
            this.name = name;
            this.originalFilename = originalFilename;
            this.contentType = contentType;
            this.content = content;
        }

        @Override
        public String getName() {
            return name;
        }

        @Override
        public String getOriginalFilename() {
            return originalFilename;
        }

        @Override
        public String getContentType() {
            return contentType;
        }

        @Override
        public boolean isEmpty() {
            return content.length == 0;
        }

        @Override
        public long getSize() {
            return content.length;
        }

        @Override
        public byte[] getBytes() {
            return content;
        }

        @Override
        public InputStream getInputStream() {
            return new ByteArrayInputStream(content);
        }

        @Override
        public void transferTo(File dest) throws IOException, IllegalStateException {
            Files.write(dest.toPath(), content);
        }
    }
}
