import bpy

def write(vert, frag):
    wrd = bpy.data.worlds['Arm']
    is_shadows = '_ShadowMap' in wrd.world_defs
    is_shadows_atlas = '_ShadowMapAtlas' in wrd.world_defs
    is_single_atlas = is_shadows_atlas and '_SingleAtlas' in wrd.world_defs

    frag.add_include_front('std/clusters.glsl')
    frag.add_uniform('vec2 cameraProj', link='_cameraPlaneProj')
    frag.add_uniform('vec2 cameraPlane', link='_cameraPlane')
    frag.add_uniform('vec4 lightsArray[maxLights * 2]', link='_lightsArray')
    frag.add_uniform('sampler2D clustersData', link='_clustersData')
    if is_shadows:
        frag.add_uniform('bool receiveShadow')
        frag.add_uniform('vec2 lightProj', link='_lightPlaneProj', included=True)
        if is_shadows_atlas:
            if not is_single_atlas:
                frag.add_uniform('sampler2DShadow shadowMapAtlasPoint', included=True)
            frag.add_uniform('vec4 pointLightDataArray[maxLightsCluster]', link='_pointLightsAtlasArray', included=True)
        else:
            frag.add_uniform('samplerCubeShadow shadowMapPoint[4]', included=True)
    vert.add_out('vec4 wvpposition')
    vert.write('wvpposition = gl_Position;')
    # wvpposition.z / wvpposition.w
    frag.write('float viewz = linearize(gl_FragCoord.z, cameraProj);')
    frag.write('int clusterI = getClusterI((wvpposition.xy / wvpposition.w) * 0.5 + 0.5, viewz, cameraPlane);')
    frag.write('int numLights = int(texelFetch(clustersData, ivec2(clusterI, 0), 0).r * 255);')

    frag.write('#ifdef HLSL')
    frag.write('viewz += texture(clustersData, vec2(0.0)).r * 1e-9;') # TODO: krafix bug, needs to generate sampler
    frag.write('#endif')

    if '_Spot' in wrd.world_defs:
        frag.add_uniform('vec4 lightsArraySpot[maxLights]', link='_lightsArraySpot')
        frag.write('int numSpots = int(texelFetch(clustersData, ivec2(clusterI, 1 + maxLightsCluster), 0).r * 255);')
        frag.write('int numPoints = numLights - numSpots;')
        if is_shadows:
            if is_shadows_atlas and not is_single_atlas:
                frag.add_uniform(f'sampler2DShadow shadowMapAtlasSpot', included=True)
            elif not is_shadows_atlas:
                frag.add_uniform('sampler2DShadow shadowMapSpot[4]', included=True)
            # FIXME: type is actually mat4, but otherwise it will not be set as floats when writing the shaders' json files
            frag.add_uniform('vec4 LWVPSpotArray[maxLightsCluster]', link='_biasLightWorldViewProjectionMatrixSpotArray', included=True)

    frag.write('for (int i = 0; i < min(numLights, maxLightsCluster); i++) {')
    frag.write('int li = int(texelFetch(clustersData, ivec2(clusterI, i + 1), 0).r * 255);')

    frag.write('direct += sampleLight(')
    frag.write('    wposition,')
    frag.write('    n,')
    frag.write('    vVec,')
    frag.write('    dotNV,')
    frag.write('    lightsArray[li * 2].xyz,') # lp
    frag.write('    lightsArray[li * 2 + 1].xyz,') # lightCol
    frag.write('    albedo,')
    frag.write('    roughness,')
    frag.write('    specular,')
    frag.write('    f0')
    if is_shadows:
        frag.write('    , li, lightsArray[li * 2].w, receiveShadow') # bias
    if '_Spot' in wrd.world_defs:
        frag.write('    , lightsArray[li * 2 + 1].w != 0.0')
        frag.write('    , lightsArray[li * 2 + 1].w') # cutoff
        frag.write('    , lightsArraySpot[li].w') # cutoff - exponent
        frag.write('    , lightsArraySpot[li].xyz') # spotDir
    if '_VoxelShadow' in wrd.world_defs and '_VoxelAOvar' in wrd.world_defs:
        frag.write('  , voxels, voxpos')
    frag.write(');')

    frag.write('}') # for numLights