//

import { Color, GLSL3, RawShaderMaterial, Vector2, Vector4 } from "three";

const shader_vx = `
    in vec3 position;
    in vec3 normal;
    in vec2 uv;
    in vec3 color;
    in vec4 tangent;

    out vec3 vNormal;
    out vec3 vTangent;
    out vec3 vBinormal;
    
    out vec2 vUv;
    out mat3 TBN;
    out vec3 vColor;
    out float depth;

    uniform mat4 modelViewMatrix;
    uniform mat4 projectionMatrix;
    uniform mat3 normalMatrix;
    uniform vec4 projection_params;
    

    void main() {

        vUv = uv;

        vColor = color;

        // get smooth normals
        vec4 mvPosition = modelViewMatrix * vec4( position, 1.0 );

        vec3 transformedNormal = normalMatrix * normal;
        
        vec3 T = normalize(vec3(modelViewMatrix * tangent) );
        vNormal = normalize( transformedNormal );
        vec3 N = vNormal;
        vec3 B = cross(N,T);

        // construct Tangent Bitangent Normal matrix
        TBN = mat3(T,B,N);
        

        gl_Position = projectionMatrix * mvPosition;


       depth = 1.0 + (mvPosition.z * projection_params.w);
    }
`;

const shader_fg = `
    precision highp float;
    precision highp int;

    layout(location = 0) out vec4 gColor;
    layout(location = 1) out vec4 gNormal;
    layout(location = 2) out vec4 gORM;

    uniform vec3 color;

    uniform sampler2D tDiffuse;
    uniform sampler2D tNormal;
    uniform sampler2D tRoughness;
    uniform sampler2D tMetalness;
    uniform sampler2D tOcclusion;
    uniform sampler2D tEmissive;
    
    uniform float uTextureLODBias;
    uniform vec2 uAtlasResolution;
    
    uniform vec2 repeat;

    in vec3 vNormal;
    in vec2 vUv;
    in mat3 TBN;
    in vec3 vColor;
    in float depth;
   
    vec4 texture_supersample_4(in sampler2D sampler, in vec2 uv, float lod_bias){
        vec2 half_pixel_size = vec2(0.5) / uAtlasResolution;
        
        // per pixel partial derivatives
        vec2 dx = dFdx(uv.xy);
        vec2 dy = dFdy(uv.xy);
        // rotated grid uv offsets
        vec2 uvOffsets = vec2(0.125, 0.375);
        vec4 offsetUV = vec4(0.0, 0.0, 0.0, uTextureLODBias);
        
        float bias = -1.0 + lod_bias;
        
        // supersampled using 2x2 rotated grid
        vec4 col;
        offsetUV.xy = uv.xy + uvOffsets.x * dx + uvOffsets.y * dy;
        col += texture(sampler, offsetUV.xy, bias);
        offsetUV.xy = uv.xy - uvOffsets.x * dx - uvOffsets.y * dy;
        col += texture(sampler, offsetUV.xy, bias);
        offsetUV.xy = uv.xy + uvOffsets.y * dx - uvOffsets.x * dy;
        col += texture(sampler, offsetUV.xy, bias);
        offsetUV.xy = uv.xy - uvOffsets.y * dx + uvOffsets.x * dy;
        col += texture(sampler, offsetUV.xy, bias);
        col *= 0.25;

        return col;
    }

    void main() {    
        // write color to G-Buffer
        vec4 diffuse_texture_sample = texture_supersample_4( tDiffuse, vUv, uTextureLODBias);
        
        if(diffuse_texture_sample.a < 0.5){
            // alpha masking
            discard;
        }
        
        gColor = diffuse_texture_sample;
        gColor.rgb *= vColor;
        gColor.rgb *= color;

        // write normals to G-Buffer
        // see https://learnopengl.com/Advanced-Lighting/Normal-Mapping
        vec3 normal_sample = texture_supersample_4(tNormal, vUv,uTextureLODBias).rgb*2.0 - 1.0;
        vec3 normal = normalize(TBN * normal_sample);
        
        gNormal = vec4( normal*0.5 + 0.5, depth);
        
        float roughness = texture_supersample_4(tRoughness, vUv,uTextureLODBias).g;
        float metalness = texture_supersample_4(tMetalness, vUv,uTextureLODBias).b;
        float occlusion = texture_supersample_4(tOcclusion, vUv,uTextureLODBias).r;
        vec3 emissive = texture_supersample_4(tEmissive, vUv,uTextureLODBias).rgb;

        gORM = vec4(occlusion, roughness, metalness, 1.0);
    }
`;

/**
 * PBR shader:
 *  - albedo + alpha
 *  - normal
 *  - roughness
 *  - metalness
 *  - ambient occlusion
 */
export class BakeShaderStandard extends RawShaderMaterial {
    constructor() {
        super({
            vertexShader: shader_vx,
            fragmentShader: shader_fg,
            glslVersion: GLSL3,
            uniforms: {
                color: {
                    value: new Color(1, 1, 1)
                },
                tDiffuse: {
                    value: null
                },
                tNormal: {
                    value: null
                },
                tRoughness: {
                    value: null
                },
                tMetalness: {
                    value: null
                },
                tOcclusion: {
                    value: null
                },
                tEmissive: {
                    value: null
                },
                /**
                 * @see https://docs.unity3d.com/2020.3/Documentation/Manual/SL-UnityShaderVariables.html
                 */
                projection_params: {
                    value: new Vector4(0, 0, 0, 0)
                },
                /**
                 * Bias texture LOD, -0.5 pushes MIP switching a half a LOD further away, and produces more crisp-looking results
                 */
                uTextureLODBias:{
                    value:-0.5
                },
                uAtlasResolution:{
                    value: new Vector2(1,1)
                }
            }
        });

        this.defaultAttributeValues.tangent = [
            1, 0, 0, 1
        ];
    }
}
