#version 330 core
in vec3 vWorldPos;
in vec3 vWorldNormal;
out vec4 FragColor;

uniform vec3 uCameraPos;
uniform float uPanelW;
uniform float uPanelH;
uniform float uPanelD;
uniform int uRows;
uniform int uCols;
uniform float uPitch;
uniform float uDotRadius;
uniform int uRenderMode; // 0=realistic, 1=dataViz
uniform bool uLightMode;

const float PI = 3.14159265359;

float saturate(float x) { return clamp(x, 0.0, 1.0); }

float hash12(vec2 p) {
    vec3 p3 = fract(vec3(p.xyx) * 0.1031);
    p3 += dot(p3, p3.yzx + 33.33);
    return fract((p3.x + p3.y) * p3.z);
}

float noise2d(vec2 p) {
    vec2 i = floor(p);
    vec2 f = fract(p);
    float a = hash12(i);
    float b = hash12(i + vec2(1.0, 0.0));
    float c = hash12(i + vec2(0.0, 1.0));
    float d = hash12(i + vec2(1.0, 1.0));
    vec2 u = f * f * (3.0 - 2.0 * f);
    return mix(a, b, u.x) + (c - a) * u.y * (1.0 - u.x) + (d - b) * u.x * u.y;
}

float fbm(vec2 p) {
    float v = 0.0;
    float a = 0.5;
    for (int i = 0; i < 4; ++i) {
        v += a * noise2d(p);
        p *= 2.0;
        a *= 0.5;
    }
    return v;
}

vec3 fresnelSchlick(float cosTheta, vec3 F0) {
    return F0 + (1.0 - F0) * pow(1.0 - cosTheta, 5.0);
}

float D_GGX(float NdotH, float roughness) {
    float a = max(0.04, roughness);
    float alpha = a * a;
    float alpha2 = alpha * alpha;
    float denom = (NdotH * NdotH) * (alpha2 - 1.0) + 1.0;
    return alpha2 / (PI * denom * denom + 1e-7);
}

float G_SchlickGGX(float NdotV, float roughness) {
    float r = roughness + 1.0;
    float k = (r * r) / 8.0;
    return NdotV / (NdotV * (1.0 - k) + k + 1e-7);
}

float G_Smith(float NdotV, float NdotL, float roughness) {
    float ggx1 = G_SchlickGGX(NdotV, roughness);
    float ggx2 = G_SchlickGGX(NdotL, roughness);
    return ggx1 * ggx2;
}

float D_GGX_Aniso(vec3 N, vec3 H, vec3 T, vec3 B, float ax, float ay) {
    float NdotH = saturate(dot(N, H));
    float TdotH = dot(T, H);
    float BdotH = dot(B, H);
    float ax2 = ax * ax;
    float ay2 = ay * ay;
    float denom = (TdotH * TdotH) / (ax2 + 1e-7) + (BdotH * BdotH) / (ay2 + 1e-7) + NdotH * NdotH;
    return 1.0 / (PI * ax * ay * denom * denom + 1e-7);
}

vec3 evalLight(
    vec3 N,
    vec3 V,
    vec3 L,
    vec3 lightColor,
    vec3 baseColor,
    float metallic,
    float roughness,
    float aniso,
    vec3 brushDir
) {
    float NdotL = saturate(dot(N, L));
    float NdotV = saturate(dot(N, V));
    if (NdotL <= 0.0 || NdotV <= 0.0) return vec3(0.0);

    vec3 H = normalize(V + L);
    float NdotH = saturate(dot(N, H));
    float VdotH = saturate(dot(V, H));

    vec3 F0 = mix(vec3(0.04), baseColor, metallic);
    vec3 F = fresnelSchlick(VdotH, F0);

    float D = D_GGX(NdotH, roughness);
    if (aniso > 0.001) {
        vec3 T = normalize(brushDir - N * dot(brushDir, N));
        vec3 B = normalize(cross(N, T));
        float alpha = max(0.04, roughness);
        float a = alpha * alpha;
        float ax = mix(a, a * 0.35, aniso);
        float ay = mix(a, a * 1.80, aniso);
        D = D_GGX_Aniso(N, H, T, B, ax, ay);
    }

    float G = G_Smith(NdotV, NdotL, roughness);
    vec3 spec = (D * G * F) / max(4.0 * NdotV * NdotL, 1e-6);

    vec3 kD = (vec3(1.0) - F) * (1.0 - metallic);
    vec3 diff = kD * baseColor / PI;

    return (diff + spec) * lightColor * NdotL;
}

float nearestDotDistanceXZ(vec2 xz) {
    if (uPitch <= 0.0 || uRows <= 0 || uCols <= 0) return 1e6;
    int colsM1 = max(uCols - 1, 0);
    int rowsM1 = max(uRows - 1, 0);
    float halfGridW = float(colsM1) * uPitch * 0.5;
    float halfGridD = float(rowsM1) * uPitch * 0.5;

    vec2 g = (xz + vec2(halfGridW, halfGridD)) / max(1e-6, uPitch);
    vec2 gi = floor(g + 0.5);
    gi = clamp(gi, vec2(0.0), vec2(float(colsM1), float(rowsM1)));

    vec2 c = gi * uPitch - vec2(halfGridW, halfGridD);
    return length(xz - c);
}

void main() {
    // panel 先用一个固定颜色（后续可以加光照/材质）
    vec3 N = normalize(vWorldNormal);
    vec3 V = normalize(uCameraPos - vWorldPos);

    float isTop = step(0.75, N.y);

    // ------------------------------------------------------------
    // Industrial engineering model: neutral matte gray panel (support layer only)
    // ------------------------------------------------------------
    vec3 topBase, sideBase;
    vec3 edgeCol, rimCol;
    if (uLightMode) {
        topBase = vec3(0.30, 0.31, 0.32);
        sideBase = vec3(0.27, 0.28, 0.29);

        edgeCol = vec3(0.020);
        rimCol = vec3(0.015);
    }
    else {
        topBase = vec3(0.78, 0.80, 0.84);
        sideBase = vec3(0.68, 0.70, 0.74);

        edgeCol = vec3(0.010, 0.12, 0.16);
        rimCol = vec3(0.12, 0.14, 0.18);
    }

    vec3 baseColor = mix(sideBase, topBase, isTop);

    // dataViz: flat/unlit, no lighting modulation (keep pure baseColor)
    if (uRenderMode == 1) {
        FragColor = vec4(clamp(baseColor, 0.0, 1.0), 1.0);
        return;
    }

    vec2 xz = vWorldPos.xz;
    float dotContact = 0.0;
    if (isTop > 0.5 && uDotRadius > 0.0) {
        float d = nearestDotDistanceXZ(xz);
        float w = max(0.002, uDotRadius * 0.22);
        dotContact = 1.0 - smoothstep(uDotRadius, uDotRadius + w, d);
    }

    // "Front light": make the light come from the camera direction (like a headlight/flashlight).
    vec3 L = V;
    // L = normalize(vec3(0.0, 0.15, -1.0));
    float diff = saturate(dot(N, L));
    float lighting = uLightMode ? (0.90 + 0.10 * diff) : (0.95 + 0.12 * diff);

    float hw = max(1e-6, uPanelW * 0.5);
    float hd = max(1e-6, uPanelD * 0.5);
    float edgeDist = min(hw - abs(vWorldPos.x), hd - abs(vWorldPos.z));
    float edgeW = max(0.002, min(hw, hd) * 0.012);
    float edgeLine = (1.0 - smoothstep(edgeW, edgeW * 2.5, edgeDist)) * isTop;

    float rim = pow(1.0 - saturate(dot(N, V)), 2.2) * isTop;
    float ao = 1.0 - dotContact * 0.08;

    vec3 col = baseColor * lighting * ao;
    col += edgeLine * edgeCol;
    col += rim * rimCol;

    // Slightly deepen the bottom face to read as thickness, but keep it subtle.
    float isBottom = step(0.75, -N.y);
    col *= mix(1.0, 0.92, isBottom);

    FragColor = vec4(clamp(col, 0.0, 1.0), 1.0);
}
