import { DoubleSide, ShaderMaterial } from "three";

export class VirtualTextureMaterial extends ShaderMaterial {
    constructor({
                    uniforms = {}
                }) {
        super({
            side: DoubleSide,
            uniforms,
            vertexShader: `
        
    out vec2 vUv;
    
    void main()
    {
        vUv = uv;
        
        vec4 mvPosition = modelViewMatrix * vec4( position, 1.0 );
        gl_Position = projectionMatrix * mvPosition;
    }
            `,
            fragmentShader: `
            
    precision highp usampler2D;

    uniform sampler2D u_page;
    uniform uvec2 u_page_resolution;
    uniform usampler2D u_mapping;
    
    uniform uint u_tile_resolution;
    uniform uint u_tile_padding;
    
    uniform uint u_texture_resolution;
    uniform uint u_max_mip_level;
    uniform uint u_mapping_texture_width;
    
    in vec2 vUv;
    
    float computeTextureLod(const in vec2 uv){
        float scale_x = float(u_texture_resolution) ;
        float scale_y = float(u_texture_resolution) ;
        
        vec2 dx = dFdx( uv * scale_x ) ;
        vec2 dy = dFdy( uv * scale_y ) ;
        
        float d = max( dot( dx, dx ), dot( dy, dy ) );
        
        return log2( d ) * 0.5;
    }
    
    uint split_by_2(const in uint a) {
        uint x = a;
        x = (x | x << 16) & 0x0000FFFFu;
        x = (x | x << 8) & 0x00FF00FFu;
        x = (x | x << 4) & 0x0F0F0F0Fu;
        x = (x | x << 2) & 0x33333333u;
        x = (x | x << 1) & 0x55555555u;
        return x;
    }
    
    uint compose_tile_address(uint mip, uint x, uint y){
        uint mip_resolution = 1u << mip;
         
        uint mip_mask = mip_resolution - 1u;
        uint index_offset = split_by_2(mip_mask);

        return index_offset + x + y * mip_resolution;
    }
    
    uvec2 PageTilePosition(
        const in uint tile_lod,
        const in uvec2 tile_coordinate, 
        out float f_mip_size
    ){
    
        uint tile_address = compose_tile_address(tile_lod, tile_coordinate.x , tile_coordinate.y);
        
        uint tile_mapping_x = tile_address % u_mapping_texture_width;
        uint tile_mapping_y = tile_address / u_mapping_texture_width;
    
        uint encoded_mapping = texelFetch(u_mapping, ivec2(tile_mapping_x, tile_mapping_y), 0 ).r;
        
        // decode mapping
        uint slot_index = encoded_mapping & 0x3FFFFFu;
        uint mip_level = ( encoded_mapping >> 22) & 0x3FFu;
        
        uint mip_size = 1u << mip_level; 
        
        f_mip_size = float(mip_size);
                        
        return uvec2(
            slot_index % u_page_resolution.x,
            slot_index / u_page_resolution.x
        );
    }
    
    vec2 PageTileSoft(in vec2 uv , in uvec2 virtual_tile_position, uint tile_lod){
    
        float f_mip_size;
        
        uvec2 tp00 = PageTilePosition( tile_lod, virtual_tile_position,  f_mip_size) ;
        
        vec2 tile_uv = fract(uv * f_mip_size);
        
        return vec2(tp00) + tile_uv;
    }
    
    vec4 VirtualFetchLod(const in vec2 uv, const in uint texture_lod_0){
        uint ui_texture_lod_0 = u_max_mip_level - texture_lod_0;
        uint tile_mip_size = 1u << ui_texture_lod_0;
                
        vec2 uvt = vec2(uv * float(tile_mip_size));
                                
        vec2 stuv = PageTileSoft(uv, uvec2(floor(uvt.x), floor(uvt.y)), ui_texture_lod_0);
        
        vec2 tile_uv = fract(stuv);
        vec2 tile_offset = stuv - tile_uv;
        
        float f_tile_resolution = float(u_tile_resolution);
        float slot_size = float(u_tile_resolution + u_tile_padding*2u);
        
        ivec2 texture_size = textureSize(u_page, 0);        
        
        vec2 texel = tile_offset * slot_size + float(u_tile_padding)  + tile_uv*(f_tile_resolution);
       
        vec2 page_uv =  texel / vec2(texture_size);
        
        return textureLod(u_page, page_uv, 0.0);
    }    
    
    vec4 VirtualFetch(const in vec2 uv){
    
        float texture_lod = computeTextureLod(uv);
        
        float lod_fract = fract(texture_lod);
        
        // return vec4(lod_fract, 1.0, 1.0, 1.0);
        
        float texture_lod_0 = float(texture_lod);
        float texture_lod_1 = ceil(texture_lod);
        
       
        vec4 lod_0 =  VirtualFetchLod(uv, uint(texture_lod_0));
        vec4 lod_1 =  VirtualFetchLod(uv, uint(texture_lod_1));
        
        vec4 mip_blend = mix(
            lod_0, lod_1, lod_fract
        );
        
        return mip_blend;
        
    }
    
    void main(){    
        
        pc_fragColor =  VirtualFetch(vUv);
        
    }
    
            `
        });
    }
}