/***********************************************/
/*           Copyright (c) 2025 Belmu          */
/*             All Rights Reserved             */
/***********************************************/

uniform sampler2D albedoAtlas;
uniform sampler2D normalAtlas;
uniform sampler2D specularAtlas;

struct Material {
    int id;

    float albedo;
    float opacity;

    vec3 normal;
    vec3 tangentNormal;
    vec3 geometricNormal;

    float N;
    float K;

    float F0;
    vec2  alpha;
    float ao;
    float emission;
    float subsurface;
    float porosity;
};

float airIOR;

float f0ToIOR(float F0) {
    F0 = sqrt(F0) * 0.99999;
    return airIOR.r * ((1.0 + F0) / (1.0 - F0));
}

vec3 f0ToIOR(vec3 F0) {
    F0 = sqrt(F0) * 0.99999;
    return airIOR * ((1.0 + F0) / (1.0 - F0));
}

float iorToF0(float ior) {
    float a = (ior - airIOR.r) / (ior + airIOR.r);
    return a * a;
}

vec3 iorToF0(vec3 ior) {
    vec3 a = (ior - airIOR) / (ior + airIOR);
    return a * a;
}

#include "/include/materials/dispersion.glsl"
#include "/include/utility/spectral.glsl"

#include "/include/materials/hardcoded.glsl"
#include "/include/materials/gerstner.glsl"

#if RAIN_PUDDLES == 1
    #include "/include/materials/puddles.glsl"
#endif

mat3 getShadingTBNMatrix(vec3 hitNormal) {
    vec3 tangent = mix(vec3(hitNormal.z, 0.0, -hitNormal.x), vec3(1.0, 0.0, 0.0), step(0.5, abs(hitNormal.y)));
    return mat3(tangent, -cross(hitNormal, tangent), hitNormal);
}

ivec2 calculateTextureCoords(uint packedTileData, vec3 hitPos, vec3 hitNormal) {
    float tileSize  =  exp2((packedTileData & 0xf) + 4);
    vec2  tileIndex =  vec2(uvec2(packedTileData) >> uvec2(4, 16) & 0xfff);

    vec3 absNormal   = abs(hitNormal);
    vec2 fractCoords = fract(absNormal.y > absNormal.x && absNormal.y > absNormal.z ? hitPos.xz : vec2(hitPos.x * sign(hitNormal.z) - hitPos.z * sign(hitNormal.x), -hitPos.y));
    return ivec2(tileIndex * tileSize + fractCoords * tileSize);
}

vec3 getVoxelAlbedoRGB(uint packedData, ivec2 textureCoords) {
    int  blockId = unpackId(packedData);
    vec3 tint    = unpackTint(packedData);

    if (blockId == SPHERE_0 || blockId == SPHERE_1) return tint;

    vec3 albedo = texelFetch(albedoAtlas, textureCoords, 0).rgb * tint;

    #if TONEMAP == ACES
        albedo = srgbToAP1Albedo(albedo);
    #else
        albedo = srgbToLinear(albedo);
    #endif

    return albedo;
}

vec3 getMaterialNormal(uint packedData, ivec2 textureCoords, vec3 hitNormal) {
    int  blockId = unpackId(packedData);
    vec3 normal  = hitNormal;

    if (blockId == SPHERE_0 || blockId == SPHERE_1) return normal;
    
    vec4 normalTexture = texelFetch(normalAtlas, textureCoords, 0);

    if (all(greaterThan(normalTexture, vec4(EPS)))) {
        normal.xy = normalTexture.xy * 2.0 - 1.0;
        normal.z  = sqrt(1.0 - saturate(dot(normal.xy, normal.xy)));
        normal    = normalize(getShadingTBNMatrix(hitNormal) * normal);
    }

    return normal;
}

Material getVoxelMaterial(uint packedData, ivec2 textureCoords, vec3 hitPosition, vec3 hitNormal, float wavelength) {
    int  blockId = unpackId(packedData);
    vec3 tint    = unpackTint(packedData);

    switch (blockId) {
        case WATER:
            vec3 tangentNormal = getWaterNormals(hitPosition, WATER_OCTAVES);
            vec3 waveNormal    = normalize(getShadingTBNMatrix(hitNormal) * tangentNormal);
            return getWaterMaterial(wavelength, waveNormal, tangentNormal, hitNormal);

        #if SPHERES_MODE == 1
            case SPHERE_0: return getMetalMaterial(wavelength,       hitNormal, vec2(METAL_SPHERE_ROUGHNESS));
            case SPHERE_1: return getGlassMaterial(wavelength, tint, hitNormal, vec2(GLASS_SPHERE_ROUGHNESS));
        #endif
        default: break;
    }

    Material material;

    material.id = blockId;

    vec4 albedoTexture   = texelFetch(albedoAtlas  , textureCoords, 0);
    vec4 normalTexture   = texelFetch(normalAtlas  , textureCoords, 0);
    vec4 specularTexture = texelFetch(specularAtlas, textureCoords, 0);

    albedoTexture.rgb *= tint;

    material.opacity = albedoTexture.a;

    vec3 albedo = albedoTexture.rgb;

    albedo = srgbToLinear(albedo);

    #if WHITE_WORLD == 1
        albedo = vec3(1.0);
    #endif

    material.albedo = srgbToSpectrum(albedo, wavelength);

    material.F0         = specularTexture.g;
    material.ao 		= 1.0;
    material.alpha      = vec2(maxEps(pow2(saturate(1.0 - specularTexture.r))));
    material.emission   = specularTexture.a * 255.0 < 254.5 ? specularTexture.a : 0.0;
    //material.subsurface = specularTexture.b * 255.0 < 65.0  ? 0.0 : specularTexture.b;
    //material.porosity   = specularTexture.b * 255.0 > 64.0  ? 0.0 : specularTexture.b;

    #if HARDCODED_EMISSIVES == 1
        if (isEmissive(blockId)) material.emission = 1.0;
    #endif

    material.normal          = hitNormal;
    material.tangentNormal   = hitNormal;
    material.geometricNormal = hitNormal;
    
    if (all(greaterThan(normalTexture, vec4(EPS)))) {
        material.normal.xy     = normalTexture.xy * 2.0 - 1.0;
        material.normal.z      = sqrt(1.0 - saturate(dot(material.normal.xy, material.normal.xy)));
        material.tangentNormal = normalize(material.normal);
        material.normal        = normalize(getShadingTBNMatrix(hitNormal) * material.normal);

        #if MATERIAL_AO != 0
            material.ao = normalTexture.b;
        #endif
    }

    if (isGlass(blockId)) {
        material.N     = sellmeier(N_BK7HT, wavelength);
        material.F0    = iorToF0(material.N);
        material.alpha = vec2(1e-8);
        return material;
    }

    if (material.F0.r * maxFloat8 > 229.5) {
        mat2x3 hcm = getHardcodedMetal(material);

        material.N = srgbToSpectrum(hcm[0], wavelength);
        material.K = srgbToSpectrum(hcm[1], wavelength);
    } else {
        float ior = f0ToIOR(material.F0);

        material.N = ior;
        material.K = 0.0;
    }

    return material;
}
