package com.ss.video.rtc.newdemo.common.utils;

import android.opengl.EGL14;
import android.opengl.GLES20;
import android.opengl.Matrix;

import com.ss.bytertc.engine.data.VideoPixelFormat;
import com.ss.bytertc.engine.video.IVideoProcessor;

import org.webrtc.GlShader;
import org.webrtc.GlUtil;

import java.nio.ByteBuffer;
import java.util.concurrent.atomic.AtomicBoolean;

import com.ss.bytertc.engine.utils.LogUtil;
import com.ss.bytertc.engine.video.VideoFrame;
import com.ss.bytertc.engine.video.builder.CpuBufferVideoFrameBuilder;
import com.ss.bytertc.engine.video.builder.GLTextureVideoFrameBuilder;

public class SimpleVideoProcessor extends IVideoProcessor {
    private static final String TAG = "SimpleVideoProcessor";
    ByteBuffer[] mYUVBuffers = new ByteBuffer[3];
    private AtomicBoolean mYUVBuffersAvailable = new AtomicBoolean(true);

    private GlColorEffect mTextureDrawer = null;
    FrameBuffer mFrameBuffer = null;
    private float[] TEX_MATRIX = new float[16];
    public static final int FORMAT_TEXTURE_2D = 10;
    public static final int FORMAT_TEXTURE_OES = 11;

    int mFrameCount = 0;

    static class FrameBuffer {
        int mFrameBufferId = 0;
        int mTextureId = 0;
        int mTexWidth = 0;
        int mTexHeight = 0;

        public FrameBuffer(int w, int h) {
            mTexWidth = w;
            mTexHeight = h;

            // fbo
            final int frameBuffers[] = new int[1];
            GLES20.glGenFramebuffers(1, frameBuffers, 0);
            mFrameBufferId = frameBuffers[0];
            GLES20.glBindFramebuffer(GLES20.GL_FRAMEBUFFER, mFrameBufferId);

            // texture
            mTextureId = GlUtil.generateTexture(GLES20.GL_TEXTURE_2D);

            GLES20.glActiveTexture(GLES20.GL_TEXTURE0);
            GLES20.glBindTexture(GLES20.GL_TEXTURE_2D, mTextureId);

            GLES20.glTexImage2D(GLES20.GL_TEXTURE_2D, 0, GLES20.GL_RGBA, w, h, 0, GLES20.GL_RGBA,
                    GLES20.GL_UNSIGNED_BYTE, null);
            GLES20.glFramebufferTexture2D(GLES20.GL_FRAMEBUFFER, GLES20.GL_COLOR_ATTACHMENT0,
                    GLES20.GL_TEXTURE_2D, mTextureId, 0);

            GLES20.glBindTexture(GLES20.GL_TEXTURE_2D, 0);
            GLES20.glBindFramebuffer(GLES20.GL_FRAMEBUFFER, 0);
            GlUtil.checkNoGLES2Error("new FrameBuffer");
        }

        public void dispose() {
            mTexWidth = 0;
            mTexHeight = 0;
            // delete texture
            GLES20.glDeleteTextures(1, new int[]{mTextureId}, 0);
            mTextureId = 0;

            // delete fbo
            GLES20.glDeleteFramebuffers(1, new int[] {mFrameBufferId}, 0);
            mFrameBufferId = 0;

            GlUtil.checkNoGLES2Error("releaseGl");
        }

        public void attachCurrent() {
            GLES20.glBindFramebuffer(GLES20.GL_FRAMEBUFFER, mFrameBufferId);
            GLES20.glViewport(0, 0, mTexWidth, mTexHeight);
            GlUtil.checkNoGLES2Error("attachCurrent");
        }

        public void detachCurrent() {
            GLES20.glBindFramebuffer(GLES20.GL_FRAMEBUFFER, 0);
            GlUtil.checkNoGLES2Error("detachCurrent");
        }
    }

    final GlColorEffect.ShaderCallbacks shaderCallback = new GlColorEffect.ShaderCallbacks();
    static class GlColorEffect extends GlGenericDrawer {
        private static final String FRAGMENT_SHADER = "uniform float greenV; \n"
                + "void main() {\n"
                + "  gl_FragColor = sample(tc);\n"
                + "  if (otc.t > 0.8 && otc.s > 0.5) {gl_FragColor = vec4(0,greenV,0,1);}\n"
                + "}\n";

        static class ShaderCallbacks implements GlGenericDrawer.ShaderCallbacks {
            private int greenVLocation;
            public float greenV = 1.0f;

            @Override
            public void onNewShader(GlShader shader) {
                greenVLocation = shader.getUniformLocation("greenV");
            }

            @Override
            public void onPrepareShader(GlShader shader, float[] texMatrix, int frameWidth, int frameHeight, 
                    int viewportWidth, int viewportHeight) {
                GLES20.glUniform1f(greenVLocation, greenV);
            }
        }

        public GlColorEffect(GlGenericDrawer.ShaderCallbacks callback) {
            super(FRAGMENT_SHADER, callback);
        }
    }

    public SimpleVideoProcessor() {
        Matrix.setIdentityM(TEX_MATRIX, 0);
    }

    @Override
    public VideoFrame processVideoFrame(VideoFrame frame) {
        try {
            mFrameCount++;
            switch (frame.getPixelFormat()) {
                case kVideoPixelFormatI420:
                    return processI420Frame(frame);
                case kVideoPixelFormatTextureOES:
                case kVideoPixelFormatTexture2D:
                    return processTextureFrame(frame);
            }
        } catch (Exception e) {
            LogUtil.e("SimpleVideoProcessor", "processFrame Exception" + e.toString());
            e.printStackTrace();
        }
        return null;
    }

    @Override
    public void onGLEnvInitiated() {
        // you can init Gl here or later in processTextureFrame.
        initGL();
    }

    @Override
    public void onGLEnvRelease() {
        // you need to release Gl here.
        releaseGl();
    }

    public VideoFrame processI420Frame(VideoFrame frame) {
        if (mYUVBuffersAvailable.get() == false) {
            // buffer is not available (not released), drop this frame.
            return null;
        }

        int width = frame.getWidth() ;
        int height = frame.getHeight();
        int strip_y = frame.getPlaneStride(0);
        int strip_u = frame.getPlaneStride(1);
        int strip_v =  frame.getPlaneStride(2);
    
        // allocate direct byte buffer.
        int chroma_wight = (width + 1) / 2;        
        int chroma_height = (height + 1) / 2;
        int[] planeLens = new int[]{strip_y * height, strip_u * chroma_height, strip_v * chroma_height};
        for(int i = 0; i < 3; i++) {
            ByteBuffer planeBuffer = mYUVBuffers[i];
            if (planeBuffer == null || planeBuffer.capacity() < planeLens[i]) {
                planeBuffer = ByteBuffer.allocateDirect(planeLens[i]);
                mYUVBuffers[i] = planeBuffer;
            }
        }

        // add color effect.
        int color = mFrameCount % 255;
        for (int j = 0; j < 3; j++) { // every plane
            ByteBuffer srcPlane = frame.getPlaneData(j).slice();
            srcPlane.rewind();
            srcPlane.limit(planeLens[j]);
            ByteBuffer plane = mYUVBuffers[j];
            plane.put(srcPlane); // copy srcPlane to direct buffer.
            int effectRows = (height + 7) / 8;
            int effectcols = (width + 1) / 2;
            if (j != 0) { // Y/U Plane
                effectRows = (chroma_height + 7) / 8;
                effectcols = (chroma_wight + 1) / 2;
            }
            // add color effect to effect region.
            for (int r = 0; r < effectRows; r++) {
                int pos = frame.getPlaneStride(j) * r + effectcols;;
                for (int c = 0; c < effectcols; c++) {
                    plane.put(pos + c, (byte) color);
                }
            }
            plane.rewind();
        }

        // we use a mYUVBuffers to optimize memory usage, don't forget to set it to available on setReleaseCallback.
        mYUVBuffersAvailable.set(false);
        // use CpuBufferVideoFrameBuilder to return a direct cpu buffer frame.
        CpuBufferVideoFrameBuilder builder = new CpuBufferVideoFrameBuilder(VideoPixelFormat.kVideoPixelFormatI420);
        builder.setWidth(width)
                .setHeight(height)
                .setRotation(frame.getRotation()) // set rotation back if rotation has not been changed.
                .setTimeStampUs(frame.getTimeStampUs())
                .setColorSpace(frame.getColorSpace())
                .setPlaneData(0, mYUVBuffers[0].slice())
                .setPlaneData(1, mYUVBuffers[1].slice())
                .setPlaneData(2, mYUVBuffers[2].slice())
                .setPlaneStride(0, strip_y)
                .setPlaneStride(1, strip_u)
                .setPlaneStride(2, strip_v)
                .setReleaseCallback(() -> {
                    // set mYUVBuffersAvailable to true to reuse mYUVBuffers.
                    mYUVBuffersAvailable.set(true);
                });

        return builder.build();
    }

    public synchronized VideoFrame processTextureFrame(VideoFrame frame) {
        // we can init gl resource in here if we don't init it in onGLEnvInitiated callback.
        if (mFrameBuffer == null || mFrameBuffer.mTexWidth != frame.getWidth()
                || mFrameBuffer.mTexHeight != frame.getHeight()) { // resize
            mFrameBuffer = new FrameBuffer(frame.getWidth(), frame.getHeight());
        }
        if (mTextureDrawer == null) { // create drawer
            mTextureDrawer = new GlColorEffect(shaderCallback);
        }

        float c = mFrameCount % 30;
        shaderCallback.greenV = (c+1) / 30; // update value of green channel.
    
        // draw color effect.
        mFrameBuffer.attachCurrent();
        float[] tex_mat = frame.getTextureMatrix();
        if (tex_mat == null) { // do matix transform. note: we should set maxtrix back if transform has not been changed.
            tex_mat = TEX_MATRIX;
        }
        int target = frame.getPixelFormat() == VideoPixelFormat.kVideoPixelFormatTextureOES ?
                FORMAT_TEXTURE_OES : FORMAT_TEXTURE_2D;        
        drawTexture(target, frame.getTextureID(), frame.getWidth(), frame.getHeight(), tex_mat);
        mFrameBuffer.detachCurrent();

        // use GLTextureVideoFrameBuilder to return a texture frame.
        GLTextureVideoFrameBuilder builder = new GLTextureVideoFrameBuilder(VideoPixelFormat.kVideoPixelFormatTexture2D);
        builder.setEGLContext(EGL14.eglGetCurrentContext())
                .setTextureID(mFrameBuffer.mTextureId)
                .setWidth(frame.getWidth())
                .setHeight(frame.getHeight())
                .setRotation(frame.getRotation()) // set rotation back if rotation has not been changed.
                .setTimeStampUs(frame.getTimeStampUs());

        return builder.build();
    }

    void initGL() {
    }

    void releaseGl() {
        if (mTextureDrawer != null) {
            mTextureDrawer.release();
            mTextureDrawer = null;
        }
        if (mFrameBuffer != null) {
            mFrameBuffer.dispose();
            mFrameBuffer = null;
        }
    }

    public void drawTexture(int target, int srcTextureId, int width, int height, float[] tex_mat) {
        switch (target) {
            case FORMAT_TEXTURE_2D:
                mTextureDrawer.drawRgb(srcTextureId, tex_mat, width, height, 0,0, width, height);
                break;
            case FORMAT_TEXTURE_OES:
                mTextureDrawer.drawOes(srcTextureId, tex_mat, width, height, 0,0, width, height);
                break;
            default:
                LogUtil.e(TAG, "draw texture type error");
        }
        GlUtil.checkNoGLES2Error("drawTexture");
    }
}