package com.clarabot.picoture;

import com.twelvemonkeys.servlet.image.IIOProviderContextListener;
import io.javalin.Javalin;
import io.javalin.http.HttpCode;
import net.coobird.thumbnailator.ThumbnailParameter;
import net.coobird.thumbnailator.Thumbnails;
import net.coobird.thumbnailator.filters.ImageFilter;
import net.coobird.thumbnailator.geometry.Positions;
import net.coobird.thumbnailator.resizers.DefaultResizerFactory;
import net.coobird.thumbnailator.resizers.Resizers;
import net.coobird.thumbnailator.tasks.UnsupportedFormatException;
import net.coobird.thumbnailator.tasks.io.InputStreamImageSource;
import net.sourceforge.argparse4j.ArgumentParsers;
import net.sourceforge.argparse4j.impl.Arguments;
import net.sourceforge.argparse4j.inf.ArgumentParser;
import net.sourceforge.argparse4j.inf.ArgumentParserException;
import net.sourceforge.argparse4j.inf.Namespace;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;

import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.List;


// NOTE: The source of the Version class is generated by gradle during build.
//       This can be used to build just as is for development, for example.
//class Version {
//    public static String getVersion() { return "?"; }
//}

class InputPicture {
    public BufferedImage image;
    public List<ImageFilter> filters;

    public InputPicture(BufferedImage i, List<ImageFilter> f) {
        image = i;
        filters = f;
    }

//    public boolean isSwappingDimensions() {
//        if (filters != null) {
//            return this._isSwappingDimensions(filters);
//        }
//        return false;
//    }
//
//    private boolean _isSwappingDimensions(List<ImageFilter> fs) {
//        for (ImageFilter f : fs) {
//            if (f == SwapDimensions.getInstance()) {
//                return true;
//            }
//            if (f instanceof Pipeline) {
//                if (this._isSwappingDimensions(((Pipeline) f).getFilters())) {
//                    return true;
//                }
//            }
//        }
//        return false;
//    }
}

public class Picoture {
    private static final Logger logger = LoggerFactory.getLogger(Picoture.class);
    private static final int defaultThumbnailDimension = 96;
    private static final int defaultPreviewArea = 640000;
    private static final float defaultQuality = 0.77f;
    private static final int maximumAspectRatio = 5;

    private static class DrawWhiteCanvasFilter implements ImageFilter {
        @Override
        public BufferedImage apply(BufferedImage img) {
            final int w = img.getWidth();
            final int h = img.getHeight();
            BufferedImage newImage = new BufferedImage(w, h, BufferedImage.TYPE_INT_RGB);
            Graphics2D graphic = newImage.createGraphics();
            graphic.setColor(Color.white); // white canvas is the most common, so we add it manually
            graphic.fillRect(0, 0, w, h);
            graphic.drawRenderedImage(img, null);
            graphic.dispose();
            return newImage;
        }
    }

    private static InputPicture imageInputReader(String contentType, InputStream inputStream) throws Exception {
        BufferedImage thumbnail;
        List<ImageFilter> extraFilters = null;
        // Load input image by the relevant library selected by the content-type header.
        if ("application/pdf".equals(contentType)) {
            PDDocument pdfDoc = PDDocument.load(inputStream);
            PDFRenderer pdfRender = new PDFRenderer(pdfDoc);
            thumbnail = pdfRender.renderImageWithDPI(0, 150);
        } else {
            // For all other cases we try to load an image file with ImageIO.
            // We need to have some tricky code here to extract the exif orientation fixes for later.
            ThumbnailParameter prepareParam = new ThumbnailParameter(
                    new Dimension(16, 16), null, true, null, null, 1.0f, 0, null, DefaultResizerFactory.getInstance(), true, true
            );
            InputStreamImageSource imageSource = new InputStreamImageSource(inputStream);
            imageSource.setThumbnailParameter(prepareParam);
            thumbnail = imageSource.read(); // TODO: getNumImage && index not supported by thumbnailator yet for selection from ICO for example
            if (prepareParam.getImageFilters().size() > 0) {
                extraFilters = prepareParam.getImageFilters();
            }
        }
        return new InputPicture(thumbnail, extraFilters);
    }

    private static BufferedImage sanitizeAspectRatio(BufferedImage source) throws IOException {
        final int w = source.getWidth();
        final int wMax = w * maximumAspectRatio;
        final int h = source.getHeight();
        final int hMax = h * maximumAspectRatio;
        if (wMax < h) {
            // The width is much less than the height, we need to cut from the height.
            return Thumbnails.of(source).sourceRegion(0, (h - wMax) >> 1, w, wMax)
                    .crop(Positions.CENTER)
                    .resizer(Resizers.NULL)
                    .size(w, wMax)
                    .asBufferedImage();
        } else if (hMax < w) {
            // The height is much less than the width, we need to cut from the width.
            return Thumbnails.of(source).sourceRegion((w - hMax) >> 1, 0, hMax, h)
                    .crop(Positions.CENTER)
                    .resizer(Resizers.NULL)
                    .size(hMax, h)
                    .asBufferedImage();
        }
        return source;
    }

    private static Thumbnails.Builder<BufferedImage> sourceImageBuilder(BufferedImage source, List<ImageFilter> extraFilters) {
        Thumbnails.Builder<BufferedImage> builder = Thumbnails.of(source);
        if (source.getColorModel().hasAlpha()) {
            builder = builder.addFilter(new DrawWhiteCanvasFilter());
        }
        builder = builder.imageType(BufferedImage.TYPE_INT_ARGB); // workaround to issue 135 (alpha-palette)
        if (extraFilters != null) {
            builder = builder.addFilters(extraFilters);
        }
        return builder;
    }

    private static Thumbnails.Builder<BufferedImage> resizeBuilderDims(BufferedImage source, int maxWidth, int maxHeight) {
        Thumbnails.Builder<BufferedImage> builder = Thumbnails.of(source);
        final int w = source.getWidth();
        final int h = source.getHeight();
        if (w > maxWidth || h > maxHeight) {
            builder = builder.size(
                    Math.min(w, maxWidth), Math.min(h, maxHeight)  // prevent upscaling the image
            );
        } else {
            builder = builder.resizer(Resizers.NULL).scale(1.0);
        }
        return builder;
    }

    private static byte[] intToByteArray(int data) {
        byte[] result = new byte[4];
        result[0] = (byte) ((data & 0xFF000000) >> 24);
        result[1] = (byte) ((data & 0x00FF0000) >> 16);
        result[2] = (byte) ((data & 0x0000FF00) >> 8);
        result[3] = (byte) (data & 0x000000FF);
        return result;
    }

    private static int encodeWidthHeight(int w, int h) {
        int o = 0;
        int m = 1;
        for (int i = 0;  i < 15; ++i) {
            o |= (((w & m) << 1) | (h & m)) << i;
            m <<= 1;
        }
        return o;
    }

    public static void main(String[] args) {
        System.setProperty("apple.awt.UIElement", "true");
        String name_version = "Picoture " + Version.getVersion();
        ArgumentParser parser = ArgumentParsers.newFor("Picoture").build()
                .defaultHelp(true)
                .description("Small worker-server to create pictures from input data.")
                .version(name_version);
        parser.addArgument("--version").action(Arguments.version());
        parser.addArgument("--host").setDefault("127.0.0.1").help("Specify the address to listen on");
        parser.addArgument("--port").type(Integer.class).choices(Arguments.range(1, 65535)).setDefault(7070).help("Specify the port to listen on");
        Namespace ns = null;
        try {
            ns = parser.parseArgs(args);
        } catch (ArgumentParserException e) {
            parser.handleError(e);
            System.exit(1);
        }
        Javalin app = Javalin.create(config -> {
            config.configureServletContextHandler(handler -> {
                handler.addEventListener(new IIOProviderContextListener());
            });
        }).start(ns.getString("host"), ns.getInt("port"));
        app.get("/version", ctx -> ctx.result(name_version));
        app.post("/pictures", ctx -> {
            final int pArea = ctx.queryParamAsClass("pArea", Integer.class).check(it -> it >= 40000 && it <= 4000000, "PREVIEW_AREA_NOT_IN_RANGE").getOrDefault(defaultPreviewArea);
            final int tMaxWidth = ctx.queryParamAsClass("tWidth", Integer.class).check(it -> it >= 16 && it <= 500, "THUMBNAIL_WIDTH_NOT_IN_RANGE").getOrDefault(defaultThumbnailDimension);
            final int tMaxHeight = ctx.queryParamAsClass("tHeight", Integer.class).check(it -> it >= 16 && it <= 500, "THUMBNAIL_HEIGHT_NOT_IN_RANGE").getOrDefault(defaultThumbnailDimension);
            final float quality = ctx.queryParamAsClass("quality", Float.class).check(it -> it >= 0.f && it <= 1.f, "QUALITY_NOT_IN_RANGE").getOrDefault(defaultQuality);
            try {
                // input
                InputPicture inputPicture = imageInputReader(ctx.contentType(), ctx.bodyAsInputStream());
                // preprocess the input (fill transparency, fix orientation, etc.)
                final BufferedImage preparedImage = sourceImageBuilder(inputPicture.image, inputPicture.filters)  // dimension may have changed due to exif orientation
                    .resizer(Resizers.NULL).scale(1.0)
                    .asBufferedImage();
                // transformations from preprocessed input to preview and from preview to thumbnail
                final float c = (float) preparedImage.getHeight() / preparedImage.getWidth();
                final int pWidth = (int) Math.round(Math.sqrt(pArea / c));
                final int pHeight = (int) Math.round(pWidth * c);
                final BufferedImage previewImage = resizeBuilderDims(preparedImage, pWidth, pHeight).asBufferedImage();
                // crop insane aspect ratio for thumbnail
                Thumbnails.Builder<BufferedImage> thumbnailBuilder = resizeBuilderDims(sanitizeAspectRatio(previewImage), tMaxWidth, tMaxHeight);
                // response output stream
                OutputStream outputStream = ctx.contentType("application/octet-stream").res.getOutputStream();
                // render thumbnail first into memory to know its exact size
                ByteArrayOutputStream thumbnailData = new ByteArrayOutputStream(131072);
                thumbnailBuilder.outputFormat("jpg").outputQuality(quality).toOutputStream(thumbnailData);
                // write the thumbnail to output with a size prefix
                outputStream.write(intToByteArray(thumbnailData.size()));
                thumbnailData.writeTo(outputStream);
                thumbnailData.close();
                // write the preview to output with the encoded width-height prefix
                outputStream.write(intToByteArray(encodeWidthHeight(previewImage.getWidth(), previewImage.getHeight())));
                Thumbnails.of(previewImage).resizer(Resizers.NULL).scale(1.0).outputFormat("jpg").outputQuality(quality).toOutputStream(outputStream);
            } catch (UnsupportedFormatException e) {
                ctx.status(HttpCode.UNSUPPORTED_MEDIA_TYPE);
            } catch (IOException e) {
                ctx.status(HttpCode.UNPROCESSABLE_ENTITY);
            }
        });
    }
}
