package shadertest;

import java.awt.BorderLayout;
import java.awt.GraphicsConfiguration;
import java.awt.event.WindowAdapter;
import java.awt.event.WindowEvent;
import java.awt.image.BufferedImage;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import javax.swing.JFrame;
import javax.swing.JPanel;
import javax.swing.Timer;
import javax.swing.WindowConstants;
import org.jogamp.java3d.Alpha;
import org.jogamp.java3d.Appearance;
import org.jogamp.java3d.Canvas3D;
import org.jogamp.java3d.GLSLShaderProgram;
import org.jogamp.java3d.GeometryArray;
import org.jogamp.java3d.GraphicsContext3D;
import org.jogamp.java3d.ImageComponent;
import org.jogamp.java3d.ImageComponent2D;
import org.jogamp.java3d.Raster;
import org.jogamp.java3d.Shader;
import org.jogamp.java3d.ShaderAppearance;
import org.jogamp.java3d.ShaderAttributeObject;
import org.jogamp.java3d.ShaderAttributeSet;
import org.jogamp.java3d.ShaderAttributeValue;
import org.jogamp.java3d.Shape3D;
import org.jogamp.java3d.SourceCodeShader;
import org.jogamp.java3d.Texture;
import org.jogamp.java3d.Texture2D;
import org.jogamp.java3d.TextureAttributes;
import org.jogamp.java3d.TextureUnitState;
import org.jogamp.java3d.Transform3D;
import org.jogamp.java3d.utils.geometry.GeometryInfo;
import org.jogamp.java3d.utils.geometry.Sphere;
import org.jogamp.java3d.utils.image.TextureLoader;
import org.jogamp.java3d.utils.shader.StringIO;
import org.jogamp.java3d.utils.universe.SimpleUniverse;
import org.jogamp.vecmath.Point3f;

public class CustomRenderingTest extends JFrame implements Runnable {
    int width = 800;
    int height = 800;
    Raster drawRaster;
    JPanel drawingPanel;
    Canvas3D canvas;
    GraphicsContext3D gc;
    Texture2D tex;
    Shape3D shape;
    Transform3D transformSpin;
    Alpha rotAlpha;
    AppearancePass1 appearancePass1;
    AppearancePass2 appearanceTemp;

    private final Timer timer;
    private final int delayInMilliSeconds = 100;
    private final List<Texture2D> texturePool = new ArrayList<Texture2D>();
    private final int poolSize = 50;
    private int texturePoolIndex = 0;

    public CustomRenderingTest() {
        this.timer = new Timer(delayInMilliSeconds, event -> render());
        this.addWindowListener(new WindowAdapter() {
            @Override
            public void windowClosing(WindowEvent e) {
                System.exit(0);
            }
        });
        
        fillTexturePool();
        
        initWindow();
        new Thread(this).start();
    }

    private void fillTexturePool() {
        for (int i = 0; i < poolSize; i++) {
            texturePool.add(new Texture2D(Texture.BASE_LEVEL, Texture.RGBA, 
                    width, height));
        }
    }

    @Override
    public void run() {
        Sphere s = new Sphere(0.8f);
        GeometryInfo info = new GeometryInfo((GeometryArray) s.getShape().getGeometry());
        BufferedImage bImg = new BufferedImage(width, height,
                BufferedImage.TYPE_4BYTE_ABGR);
        ImageComponent2D imageBuffer = new ImageComponent2D(ImageComponent.FORMAT_RGBA, bImg, true, true);
        imageBuffer.setCapability(ImageComponent2D.ALLOW_IMAGE_READ
                | ImageComponent2D.ALLOW_IMAGE_WRITE);
        drawRaster = new Raster(new Point3f(0.0f, 0.0f, 0.0f), Raster.RASTER_COLOR,
                0, 0, width, height, imageBuffer, null);
        drawRaster.setCapability(Raster.ALLOW_IMAGE_WRITE
                | Raster.ALLOW_IMAGE_READ | Raster.ALLOW_SIZE_READ
                | Raster.ALLOW_SIZE_WRITE);
        appearancePass1 = new AppearancePass1();
        shape = new Shape3D(info.getGeometryArray(), new AppearancePass2(bImg));
        tex = new Texture2D(Texture.BASE_LEVEL, Texture.RGBA, width, height);
        tex.setCapability(Texture.ALLOW_IMAGE_WRITE);
        transformSpin = new Transform3D();
        rotAlpha = new Alpha(-1, 6000);
        
//        while (true) {
//            render();
//            Thread.yield();
//        }

        timer.start();
    }
    

    private void render() {
        // first, apply transforms
        double angle = rotAlpha.value() * 2.0 * Math.PI;
        transformSpin.rotY(angle);
        // then render the mask pass
        gc.clear();
        gc.setModelTransform(transformSpin);
        appearanceTemp = (AppearancePass2) shape.getAppearance(); 
        /* store the shape's individual appearance for later use 
        (but it must be a individually configured Pass2 appearance) */
        shape.setAppearance(appearancePass1);
        gc.draw(shape);
        /* after drawing the masks, read framebuffer to drawRaster and update the
        Pass2 shader's uniform with it */
        
        gc.readRaster(drawRaster);
        
        getTexture2DFromPool();
        
        //tex = new Texture2D(Texture.BASE_LEVEL, Texture.RGBA, width, height);
        tex.setImage(0, drawRaster.getImage());
        appearanceTemp.updateMaskBufferTexture(tex);
        // clear the screen to draw the second pass
        gc.clear();
        shape.setAppearance(appearanceTemp);
        gc.draw(shape);
        /* repeat, if another pass needed (like postFX that need more information than
        from the mask pass) */

        // last action = swap
        canvas.swap();
    }

    private void getTexture2DFromPool() {
        tex = texturePool.get(texturePoolIndex);
        
        texturePoolIndex++;
        if(texturePoolIndex == texturePool.size()) {
            texturePoolIndex = 0;
        }
    }

    private void initWindow() {
        drawingPanel = new JPanel();
        setDefaultCloseOperation(WindowConstants.DISPOSE_ON_CLOSE);
        setTitle("Immediate Multipass Rendering Test");
        drawingPanel.setLayout(new BorderLayout());
        drawingPanel.setPreferredSize(new java.awt.Dimension(width, height));
        getContentPane().add(drawingPanel, java.awt.BorderLayout.CENTER);
        pack();
        GraphicsConfiguration config = SimpleUniverse.getPreferredConfiguration();
        
        canvas = new Canvas3D(config);
        gc = canvas.getGraphicsContext3D();
        
        canvas.stopRenderer();
        SimpleUniverse universe = new SimpleUniverse(canvas);
        universe.getViewingPlatform().setNominalViewingTransform();
        drawingPanel.add(canvas, java.awt.BorderLayout.CENTER);
    }

    public static void main(String args[]) {
        System.setProperty("sun.awt.noerasebackground", "true");
        java.awt.EventQueue.invokeLater(() -> new CustomRenderingTest().setVisible(true));
    }
    
    private static class AppearancePass1 extends ShaderAppearance {

        private AppearancePass1() {
            super();
            
            final Shader vertexShader = loadShader("shaders/Pass1.vert",
                    Shader.SHADING_LANGUAGE_GLSL, Shader.SHADER_TYPE_VERTEX);
            final Shader fragmentShader = loadShader("shaders/Pass1.frag",
                    Shader.SHADING_LANGUAGE_GLSL, Shader.SHADER_TYPE_FRAGMENT);
            final Shader[] shaders = new Shader[]{vertexShader, fragmentShader};
            if (vertexShader == null || fragmentShader == null) {
                throw new RuntimeException("Unable to load Pass1 shader");
            }
            final GLSLShaderProgram glslShaderProgram = new GLSLShaderProgram();
            glslShaderProgram.setShaders(shaders);
            setShaderProgram(glslShaderProgram);
        }
    }

    private static class AppearancePass2 extends ShaderAppearance {

        private AppearancePass2(BufferedImage pass1Masks) {
            super();

            final Shader testShaderVert = loadShader("shaders/Pass2.vert",
                    Shader.SHADING_LANGUAGE_GLSL, Shader.SHADER_TYPE_VERTEX);
            final Shader testShaderFrag = loadShader("shaders/Pass2.frag",
                    Shader.SHADING_LANGUAGE_GLSL, Shader.SHADER_TYPE_FRAGMENT);
            final Shader[] shaders = new Shader[]{testShaderVert, testShaderFrag};
            if (testShaderFrag == null) {
                throw new RuntimeException("Unable to load Pass2 shader");
            }
            final GLSLShaderProgram glslShaderProgram = new GLSLShaderProgram();
            glslShaderProgram.setShaders(shaders);
            setShaderProgram(glslShaderProgram);
            setCapability(Appearance.ALLOW_RENDERING_ATTRIBUTES_WRITE);
            setCapability(ShaderAppearance.ALLOW_TEXTURE_WRITE);
            final GLSLShaderProgram glslShader = new GLSLShaderProgram();
            glslShader.setShaders(shaders);
            setShaderProgram(glslShader);

            TextureLoader loader = new TextureLoader(pass1Masks,
                     TextureLoader.BY_REFERENCE | TextureLoader.Y_UP
                    | TextureLoader.ALLOW_NON_POWER_OF_TWO);
            Texture2D tex = (Texture2D) loader.getTexture();
            ImageComponent2D imageComp = (ImageComponent2D) tex.getImage(0);
            tex.setCapability(Texture.ALLOW_IMAGE_WRITE);
            tex.setBoundaryModeS(Texture.CLAMP);
            tex.setBoundaryModeT(Texture.CLAMP);
            tex.setBoundaryColor(1.0f, 1.0f, 1.0f, 1.0f);
            tex.setImage(0, imageComp);
            TextureAttributes texAttr = new TextureAttributes();
            texAttr.setTextureMode(TextureAttributes.REPLACE);
            TextureUnitState[] tus = new TextureUnitState[1];
            tus[0] = new TextureUnitState();
            tus[0].setCapability(TextureUnitState.ALLOW_STATE_WRITE
                    | TextureUnitState.ALLOW_STATE_READ);
            tus[0].setTexture(tex);
            tus[0].setTextureAttributes(texAttr);
            setTextureUnitState(tus);
            final String[] shaderAttrNames = new String[]{"masks", "screenWidth",
                "screenHeight"};
            glslShader.setShaderAttrNames(shaderAttrNames);
            ShaderAttributeSet shaderAttributeSet = new ShaderAttributeSet();
            shaderAttributeSet.put(new ShaderAttributeValue("masks", 0)); //TextureUnitState index 0
                shaderAttributeSet.put(new ShaderAttributeValue("screenWidth", 800));
            shaderAttributeSet.put(new ShaderAttributeValue("screenHeight", 800));
            /* width and height will be updated by buffer every frame, so initialization size is not
            important.*/
            setShaderAttributeSet(shaderAttributeSet);
        }

        public void updateMaskBufferTexture(Texture t) {
            TextureUnitState tus = this.getTextureUnitState(0);
            tus.setTexture(t);
            ((ShaderAttributeObject) getShaderAttributeSet().get("screenWidth")).setValue(t.getWidth());
            ((ShaderAttributeObject) getShaderAttributeSet().get("screenHeight")).setValue(t.getHeight());
        }
    }
    
    private static Shader loadShader(String shaderFileName, final int shaderLanguage,
            final int shaderType) {

        try {
            String program = StringIO.readFully(shaderFileName);
            return new SourceCodeShader(shaderLanguage, shaderType, program);
            
        } catch (IOException ex) {
            throw new RuntimeException(ex);
        }
    }
}