package eu.jonahbauer.raytracing.render;

import eu.jonahbauer.raytracing.render.canvas.Canvas;
import eu.jonahbauer.raytracing.render.canvas.Image;
import org.jetbrains.annotations.NotNull;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.zip.CRC32;
import java.util.zip.CheckedOutputStream;
import java.util.zip.DeflaterOutputStream;
import java.util.zip.InflaterInputStream;

public enum ImageFormat {
    PPM {
        @Override
        public void write(@NotNull Canvas image, @NotNull OutputStream out) throws IOException {
            try (var writer = new BufferedWriter(new OutputStreamWriter(out, StandardCharsets.US_ASCII))) {
                writer.write("P3\n");
                writer.write(String.valueOf(image.getWidth()));
                writer.write(" ");
                writer.write(String.valueOf(image.getHeight()));
                writer.write("\n255\n");

                var it = image.pixels().iterator();
                while (it.hasNext()) {
                    var color = it.next();
                    writer.write(String.valueOf(color.red()));
                    writer.write(" ");
                    writer.write(String.valueOf(color.green()));
                    writer.write(" ");
                    writer.write(String.valueOf(color.blue()));
                    writer.write("\n");
                }
            }
        }
    },
    PNG {
        private static final byte[] MAGIC = new byte[] { (byte) 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A };
        private static final int IHDR_LENGTH = 13;
        private static final int IHDR_TYPE = 0x49484452;
        private static final int IDAT_TYPE = 0x49444154;
        private static final int IEND_TYPE = 0x49454E44;
        private static final int IEND_CRC = 0xAE426082;

        @Override
        public void write(@NotNull Canvas image, @NotNull OutputStream out) throws IOException {
            try (var data = new NoCloseDataOutputStream(out); var _ = data.closeable()) {
                data.write(MAGIC);

                writeIHDR(image, data);
                writeIDAT(image, data);
                writeIEND(image, data);
            }
        }

        private void writeIHDR(@NotNull Canvas image, @NotNull DataOutputStream data) throws IOException {
            data.writeInt(IHDR_LENGTH);
            try (
                    var crc = new CheckedOutputStream(data, new CRC32());
                    var ihdr = new DataOutputStream(crc)
            ) {
                ihdr.writeInt(IHDR_TYPE);
                ihdr.writeInt(image.getWidth()); // image width
                ihdr.writeInt(image.getHeight()); // image height
                ihdr.writeByte(8); // bit depth
                ihdr.writeByte(2); // color type
                ihdr.writeByte(0); // compression method
                ihdr.writeByte(0); // filter method
                ihdr.writeByte(0); // interlace method
                ihdr.flush();
                data.writeInt((int) crc.getChecksum().getValue());
            }
        }

        private void writeIDAT(@NotNull Canvas image, @NotNull DataOutputStream data) throws IOException {
            try (
                    var baos = new ByteArrayOutputStream();
                    var crc = new CheckedOutputStream(baos, new CRC32());
                    var idat = new DataOutputStream(crc)
            ) {
                idat.writeInt(IDAT_TYPE);

                try (var deflate = new DataOutputStream(new DeflaterOutputStream(idat))) {
                    var pixels = image.pixels().iterator();
                    for (int i = 0; pixels.hasNext(); i = (i + 1) % image.getWidth()) {
                        if (i == 0) deflate.writeByte(0); // filter type
                        var pixel = pixels.next();
                        deflate.writeByte(pixel.red());
                        deflate.writeByte(pixel.green());
                        deflate.writeByte(pixel.blue());
                    }
                }

                var bytes = baos.toByteArray();
                data.writeInt(bytes.length - 4); // don't include type in length
                data.write(bytes);
                data.writeInt((int) crc.getChecksum().getValue());
            }
        }

        private void writeIEND(@NotNull Canvas image, @NotNull DataOutputStream data) throws IOException {
            data.writeInt(0);
            data.writeInt(IEND_TYPE);
            data.writeInt(IEND_CRC);
        }

        private static class NoCloseDataOutputStream extends DataOutputStream {
            public NoCloseDataOutputStream(OutputStream out) {
                super(out);
            }

            @Override
            public void close() {
                // do nothing
            }

            public Closeable closeable() {
                return super::close;
            }
        }
    },
    ;

    public void write(@NotNull Canvas image, @NotNull Path path) throws IOException {
        try (var out = Files.newOutputStream(path)) {
            write(image, out);
        }
    }

    public abstract void write(@NotNull Canvas image, @NotNull OutputStream out) throws IOException;
}
