out vec4 FragColor;

in vec4 vertexColor;
in vec2 texCoords;
flat in ivec2 clutBase;
flat in ivec2 texpageBase;
flat in uint rawTexpage;
flat in int texMode;

uniform sampler2D u_vramTex;
uniform ivec4 u_texWindow;
uniform int u_forceSetSTP;
uniform int u_texPage;

int floatToU5(float f) {
    return int(floor(f * 31.0 + 0.5));
}

vec4 sampleVRAM(ivec2 coords) {
    coords &= ivec2(1023, 511); // Out-of-bounds VRAM accesses wrap
    return texelFetch(u_vramTex, coords, 0);
}

int sample16(ivec2 coords) {
    vec4 color = sampleVRAM(coords);
    int r = floatToU5(color.r);
    int g = floatToU5(color.g);
    int b = floatToU5(color.b);
    int msb = int(ceil(color.a)) << 15;
    return r | (g << 5) | (b << 10) | msb;
}

// Apply texture blending
// Formula for RGB8 colors: col1 * col2 / 128
vec4 texBlend(vec4 color1, vec4 color2) {
    vec4 ret = (color1 * color2) / (128.0 / 255.0);
    ret.a = 1.0;
    if (u_forceSetSTP == 0) {
        ret.a = color1.a;
    } else {
        ret.a = 1.0;
    }
    return ret;
}

// https://psx-spx.consoledev.net/graphicsprocessingunitgpu/#24bit-rgb-to-15bit-rgb-dithering-enabled-in-texpage-attribute
int[16] ditherLUT = int[16](
    -4, +0, -3, +1,
     2, -2, +3, -1,
    -3, +1, -4, +0,
    +3, -1, +2, -2
);

uvec3 doDither(uvec3 col)
{
  uvec2 c = uvec2(gl_FragCoord.xy) & uvec2(3u, 3u);
  int off = ditherLUT[int(c.y * 4u + c.x)];
  col = uvec3(clamp((ivec3(col) + off) >> 3, 0, 31)); // to 16bit
  return col;
}

vec4 convertTo16(vec4 c) {
    uvec3 vc = uvec3(c.rgb * vec3(255.0, 255.0, 255.0));
    if ((int(rawTexpage) & 0x0200) != 0) {
        vc = doDither(vc);
    } else {
        vc >>= 3;
    }
    vec4 col16 = vec4(vec3(vc), 1.0);
    col16.rgb /= 31.0;
    return vec4(col16.rgb, c.a);

    // true color
    // return vec4(c.rgb, c.a);
}

void setDepth(float a) {
    if (u_forceSetSTP == 1 || a > 0.0) {
        gl_FragDepth = gl_FragCoord.z;
        return;
    }
    gl_FragDepth = 0.0;
}

void main()
{
    if (texMode == 4) { // Untextured primitive
        FragColor = vertexColor;
        if (u_forceSetSTP != 0) {
            FragColor.a = 1.0;
        } else {
            FragColor.a = 0.0;
        }
        FragColor = convertTo16(FragColor);
        setDepth(FragColor.a);
        return;
    }

    // Fix up UVs and apply texture window
    ivec2 UV = ivec2(floor(texCoords + vec2(0.0001, 0.0001))) & ivec2(0xff);
    if (u_texWindow != ivec4(0, 0, 0, 0)) {
        UV = (UV & u_texWindow.xy) | u_texWindow.zw;
    }

    if (texMode == 0) { // 4bpp texture
        ivec2 texelCoord = ivec2(UV.x >> 2, UV.y) + texpageBase;

        int smpl = sample16(texelCoord);
        int shift = (UV.x & 3) << 2;
        int clutIndex = (smpl >> shift) & 0xf;

        ivec2 sampleCoords = ivec2(clutBase.x + clutIndex, clutBase.y);
        FragColor = texelFetch(u_vramTex, sampleCoords, 0);

        if (FragColor.rgb == vec3(0.0, 0.0, 0.0)) {
            discard;
        }
        FragColor = texBlend(FragColor, vertexColor);
        FragColor = convertTo16(FragColor);
        setDepth(FragColor.a);
    } else if (texMode == 1) { // 8bpp texture
        ivec2 texelCoord = ivec2(UV.x >> 1, UV.y) + texpageBase;

        int smpl = sample16(texelCoord);
        int shift = (UV.x & 1) << 3;
        int clutIndex = (smpl >> shift) & 0xff;

        ivec2 sampleCoords = ivec2(clutBase.x + clutIndex, clutBase.y);

        FragColor = texelFetch(u_vramTex, sampleCoords, 0);

        if (FragColor.rgb == vec3(0.0, 0.0, 0.0)) {
            discard;
        }
        FragColor = texBlend(FragColor, vertexColor);
        FragColor = convertTo16(FragColor);
        setDepth(FragColor.a);
    } else { // Texture depth 2 and 3 both indicate 16bpp textures
        ivec2 texelCoord = UV + texpageBase;
        FragColor = sampleVRAM(texelCoord);

        if (FragColor.rgb == vec3(0.0, 0.0, 0.0)) {
            discard;
        }
        FragColor = texBlend(FragColor, vertexColor);
        // FragColor = convertTo16(FragColor);
        setDepth(FragColor.a);
    }
}
