import { GLSL3, NoBlending, RawShaderMaterial, Vector3 } from "three";

// Debug-viz variant: outputs the parallax-corrected world-space surface
// normal as RGB, mapped (-1..+1) -> (0..1). Use this side-by-side with
// the lit impostor to verify the atlas normal decode is producing the
// same normals the underlying mesh would.
//
// Same uniforms / card geometry as ImpostorShaderV0. Per-corner the atlas
// stores normals in bake-camera view space (see BakeShaderStandard.js), so
// we have to undo each frame's view rotation before blending — blending
// raw encoded bytes mixes incompatible bases.
//
// We grab modelMatrix in the fragment by declaring it explicitly: three.js
// auto-injects it in the vertex prefix but not fragment. WebGLRenderer.js
// still calls p_uniforms.setValue('modelMatrix', ...) every draw, so the
// declaration is enough to make it bind.
const shader_vx = `

    in vec2 uv;
    in vec3 position;

    out vec2 vUv;
    out vec3 vViewPos;

    flat out vec2 vGridFloor;
    flat out vec4 vWeights;

    flat out vec3 vTangent;
    flat out vec3 vBinormal;
    flat out vec3 vNormal;

    // Object-space corner-frame ray directions; the fragment shader uses
    // each to rebuild that frame's bake basis and decode its atlas normal.
    flat out vec3 vRay00;
    flat out vec3 vRay10;
    flat out vec3 vRay01;
    flat out vec3 vRay11;

    flat out vec4 vFrameXform00;
    flat out vec4 vFrameXform10;
    flat out vec4 vFrameXform01;
    flat out vec4 vFrameXform11;

    uniform mat4 modelViewMatrix;
    uniform mat4 projectionMatrix;

    uniform vec3 uOffset;
    uniform float uRadius;
    uniform float uFrames;
    uniform bool uIsFullSphere;

    vec2 VecToSphereOct(vec3 pivotToCamera)
    {
        vec3 octant = sign(pivotToCamera);
        float sum = dot(pivotToCamera, octant);
        vec3 octahedron = pivotToCamera / sum;
        if (octahedron.y < 0.0) {
            vec3 absolute = abs(octahedron);
            octahedron.xz = octant.xz * vec2(1.0 - absolute.z, 1.0 - absolute.x);
        }
        return octahedron.xz;
    }

    vec2 VecToHemiSphereOct(vec3 v)
    {
        v.y = max(v.y, 0.001);
        v = normalize(v);
        vec3 octant = sign(v);
        float sum = dot(v, octant);
        vec3 octahedron = v / sum;
        return vec2(
            octahedron.x + octahedron.z,
            octahedron.z - octahedron.x
        );
    }

    vec2 VectorToGrid(vec3 v)
    {
        return uIsFullSphere ? VecToSphereOct(v) : VecToHemiSphereOct(v);
    }

    vec3 OctaSphereDec(vec2 coord)
    {
        coord = (coord - 0.5) * 2.0;
        vec3 p = vec3(coord.x, 0.0, coord.y);
        vec2 a = abs(p.xz);
        p.y = 1.0 - a.x - a.y;
        if (p.y < 0.0) {
            p.xz = sign(p.xz) * vec2(1.0 - a.y, 1.0 - a.x);
        }
        return p;
    }

    vec3 OctaHemiSphereDec(vec2 coord)
    {
        vec3 p = vec3(coord.x - coord.y, 0.0, -1.0 + coord.x + coord.y);
        vec2 a = abs(p.xz);
        p.y = 1.0 - a.x - a.y;
        return p;
    }

    vec3 GridToVector(vec2 coord)
    {
        return uIsFullSphere ? OctaSphereDec(coord) : OctaHemiSphereDec(coord);
    }

    vec3 FrameToRay(vec2 frame, vec2 framesMinusOne)
    {
        vec2 f = clamp(frame / framesMinusOne, 0.0, 1.0);
        return normalize(GridToVector(f));
    }

    vec4 ComputeFrameXform(vec3 D_frame, vec3 tangent_OS, vec3 binormal_OS)
    {
        vec3 D = abs(D_frame.y) > 0.99999
                 ? normalize(D_frame + vec3(0.0, 0.0, 0.0001))
                 : D_frame;
        vec3 bake_right = normalize(cross(vec3(0.0, 1.0, 0.0), D));
        vec3 bake_up    = cross(D, bake_right);
        return vec4(
            dot(tangent_OS,  bake_right),
            dot(binormal_OS, bake_right),
            dot(tangent_OS,  bake_up),
            dot(binormal_OS, bake_up)
        );
    }

    vec4 BilinearWeights(vec2 frac_uv)
    {
        vec2 omuv = vec2(1.0) - frac_uv;
        return vec4(
            omuv.x * omuv.y,
            frac_uv.x * omuv.y,
            omuv.x * frac_uv.y,
            frac_uv.x * frac_uv.y
        );
    }

    void main() {
        vUv = uv;

        vec3 cameraPos_OS = (inverse(modelViewMatrix) * vec4(0.0, 0.0, 0.0, 1.0)).xyz;
        // Atlas lookup must be relative to the bounding-sphere centre
        // (uOffset), since the bake's D_frame is measured from the centre,
        // not from the object origin.
        vec3 pivotToCameraRay = normalize(cameraPos_OS - uOffset);

        vec2 framesMinusOne = vec2(uFrames - 1.0);
        vec2 octahedral_uv = clamp(VectorToGrid(pivotToCameraRay) * 0.5 + 0.5, 0.0, 1.0);
        vec2 grid = octahedral_uv * framesMinusOne;
        vec2 gridFloor = min(floor(grid), framesMinusOne - 1.0);
        vec4 weights = BilinearWeights(grid - gridFloor);

        vGridFloor = gridFloor;
        vWeights = weights;

        vec3 ray00 = FrameToRay(gridFloor + vec2(0.0, 0.0), framesMinusOne);
        vec3 ray10 = FrameToRay(gridFloor + vec2(1.0, 0.0), framesMinusOne);
        vec3 ray01 = FrameToRay(gridFloor + vec2(0.0, 1.0), framesMinusOne);
        vec3 ray11 = FrameToRay(gridFloor + vec2(1.0, 1.0), framesMinusOne);

        vRay00 = ray00;
        vRay10 = ray10;
        vRay01 = ray01;
        vRay11 = ray11;

        vec3 projectedRay = normalize(
            ray00 * weights.x +
            ray10 * weights.y +
            ray01 * weights.z +
            ray11 * weights.w
        );

        vec3 normal_OS = projectedRay;
        vec3 up_OS = abs(normal_OS.y) > 0.999
                     ? vec3(0.0, 0.0, -1.0)
                     : vec3(0.0, 1.0,  0.0);
        vec3 tangent_OS  = normalize(cross(up_OS, normal_OS));
        vec3 binormal_OS = cross(normal_OS, tangent_OS);

        float card_diameter = uRadius * 2.0;
        vec3 pos_OS = uOffset
                    + position.x * card_diameter * tangent_OS
                    + position.y * card_diameter * binormal_OS;

        vec4 mvPosition = modelViewMatrix * vec4(pos_OS, 1.0);
        vViewPos = mvPosition.xyz;
        gl_Position = projectionMatrix * mvPosition;

        mat3 m3 = mat3(modelViewMatrix);
        vTangent  = normalize(m3 * tangent_OS);
        vBinormal = normalize(m3 * binormal_OS);
        vNormal   = normalize(m3 * normal_OS);

        vFrameXform00 = ComputeFrameXform(ray00, tangent_OS, binormal_OS);
        vFrameXform10 = ComputeFrameXform(ray10, tangent_OS, binormal_OS);
        vFrameXform01 = ComputeFrameXform(ray01, tangent_OS, binormal_OS);
        vFrameXform11 = ComputeFrameXform(ray11, tangent_OS, binormal_OS);
    }
`;
const shader_fg = `
    precision highp float;
    precision highp int;

    in vec2 vUv;
    in vec3 vViewPos;

    flat in vec2 vGridFloor;
    flat in vec4 vWeights;

    flat in vec3 vTangent;
    flat in vec3 vBinormal;
    flat in vec3 vNormal;

    flat in vec3 vRay00;
    flat in vec3 vRay10;
    flat in vec3 vRay01;
    flat in vec3 vRay11;

    flat in vec4 vFrameXform00;
    flat in vec4 vFrameXform10;
    flat in vec4 vFrameXform01;
    flat in vec4 vFrameXform11;

    out vec4 color_out;

    // three.js auto-injects modelMatrix only in the vertex prefix. Declare
    // it explicitly here so WebGLRenderer.js's unconditional setValue call
    // finds the location and binds object.matrixWorld.
    uniform mat4 modelMatrix;

    uniform sampler2D tBase;
    uniform sampler2D tGeometry;
    uniform float uFrames;
    uniform float uDepthScale;

    vec2 apply_frame_xform(vec2 card_uv, vec4 xform)
    {
        vec2 c = card_uv - 0.5;
        return vec2(
            xform.x * c.x + xform.y * c.y,
            xform.z * c.x + xform.w * c.y
        ) + 0.5;
    }

    // Decode encoded RGB (bake-view-space normal) back to OBJECT space
    // using the corner's D_frame direction. Mirrors BakeShaderStandard's
    // encode side.
    vec3 decode_normal_OS(vec3 encoded, vec3 D_frame_OS)
    {
        vec3 n_bv = encoded * 2.0 - 1.0;
        vec3 D = abs(D_frame_OS.y) > 0.99999
                 ? normalize(D_frame_OS + vec3(0.0, 0.0, 0.0001))
                 : D_frame_OS;
        vec3 bake_right = normalize(cross(vec3(0.0, 1.0, 0.0), D));
        vec3 bake_up    = cross(D, bake_right);
        return bake_right * n_bv.x + bake_up * n_bv.y + D * n_bv.z;
    }

    void main() {
        // Same parallax shift as V0 so the visualised normals correspond
        // to the surface positions you actually see in the colour pass.
        vec3 view_dir_view = normalize(-vViewPos);
        vec3 view_dir = vec3(
            dot(vTangent,  view_dir_view),
            dot(vBinormal, view_dir_view),
            dot(vNormal,   view_dir_view)
        );

        vec2 frame_size = vec2(1.0 / uFrames);
        vec2 base_uv = vUv;

        vec2 uv00 = clamp(apply_frame_xform(base_uv, vFrameXform00), 0.0, 1.0);
        vec2 uv10 = clamp(apply_frame_xform(base_uv, vFrameXform10), 0.0, 1.0);
        vec2 uv01 = clamp(apply_frame_xform(base_uv, vFrameXform01), 0.0, 1.0);
        vec2 uv11 = clamp(apply_frame_xform(base_uv, vFrameXform11), 0.0, 1.0);

        vec2 atlas_uv00 = (vGridFloor + vec2(0.0, 0.0) + uv00) * frame_size;
        vec2 atlas_uv10 = (vGridFloor + vec2(1.0, 0.0) + uv10) * frame_size;
        vec2 atlas_uv01 = (vGridFloor + vec2(0.0, 1.0) + uv01) * frame_size;
        vec2 atlas_uv11 = (vGridFloor + vec2(1.0, 1.0) + uv11) * frame_size;

        // Pass 1: parallax depth shift.
        float d00 = texture(tGeometry, atlas_uv00).a;
        float d10 = texture(tGeometry, atlas_uv10).a;
        float d01 = texture(tGeometry, atlas_uv01).a;
        float d11 = texture(tGeometry, atlas_uv11).a;
        float blended_depth =
            d00 * vWeights.x + d10 * vWeights.y +
            d01 * vWeights.z + d11 * vWeights.w;

        base_uv += (view_dir.xy / view_dir.z) * (blended_depth - 0.5) * uDepthScale;
        base_uv = clamp(base_uv, 0.0, 1.0);

        // Pass 2: resample at parallax-shifted UV.
        uv00 = clamp(apply_frame_xform(base_uv, vFrameXform00), 0.0, 1.0);
        uv10 = clamp(apply_frame_xform(base_uv, vFrameXform10), 0.0, 1.0);
        uv01 = clamp(apply_frame_xform(base_uv, vFrameXform01), 0.0, 1.0);
        uv11 = clamp(apply_frame_xform(base_uv, vFrameXform11), 0.0, 1.0);

        atlas_uv00 = (vGridFloor + vec2(0.0, 0.0) + uv00) * frame_size;
        atlas_uv10 = (vGridFloor + vec2(1.0, 0.0) + uv10) * frame_size;
        atlas_uv01 = (vGridFloor + vec2(0.0, 1.0) + uv01) * frame_size;
        atlas_uv11 = (vGridFloor + vec2(1.0, 1.0) + uv11) * frame_size;

        vec4 b00 = texture(tBase, atlas_uv00);
        vec4 b10 = texture(tBase, atlas_uv10);
        vec4 b01 = texture(tBase, atlas_uv01);
        vec4 b11 = texture(tBase, atlas_uv11);
        vec4 albedo =
            b00 * vWeights.x + b10 * vWeights.y +
            b01 * vWeights.z + b11 * vWeights.w;

        if (albedo.a <= 0.5) {
            // Honour the same cutout silhouette as V0 — otherwise the viz
            // would show garbage normals in the discarded margin.
            discard;
        }

        // Per-corner decode in each frame's own bake basis, then blend in
        // object space. (Blending the encoded RGB first would mix bases
        // and skew the normal toward whichever frames happened to have
        // their bake_right pointing similarly.)
        vec4 g00 = texture(tGeometry, atlas_uv00);
        vec4 g10 = texture(tGeometry, atlas_uv10);
        vec4 g01 = texture(tGeometry, atlas_uv01);
        vec4 g11 = texture(tGeometry, atlas_uv11);

        vec3 n_OS_00 = decode_normal_OS(g00.rgb, vRay00);
        vec3 n_OS_10 = decode_normal_OS(g10.rgb, vRay10);
        vec3 n_OS_01 = decode_normal_OS(g01.rgb, vRay01);
        vec3 n_OS_11 = decode_normal_OS(g11.rgb, vRay11);
        vec3 normal_OS = normalize(
            n_OS_00 * vWeights.x + n_OS_10 * vWeights.y +
            n_OS_01 * vWeights.z + n_OS_11 * vWeights.w
        );

        // Object -> world. Ignoring scale (mat3(modelMatrix) is fine for
        // rotation-only or uniform-scale impostors, which is the common
        // case). Non-uniform scale would need the inverse-transpose.
        vec3 normal_world = normalize(mat3(modelMatrix) * normal_OS);

        color_out = vec4(normal_world * 0.5 + 0.5, 1.0);
    }
`;

/**
 * Debug viz: outputs the reconstructed world-space surface normal as RGB
 * (n*0.5+0.5). Same atlas / parallax math as ImpostorShaderLitV0 — this
 * is the lit shader minus the lighting, so any wrongness visible here is
 * mirrored 1:1 in the lit output.
 *
 * Same uniforms as ImpostorShaderV0.
 */
export class ImpostorShaderNormalsV0 extends RawShaderMaterial {
    constructor() {
        super({
            fragmentShader: shader_fg,
            vertexShader: shader_vx,
            uniforms: {
                tBase: { value: null },
                tGeometry: { value: null },
                tMaterial: { value: null },
                uFrames: { value: 0 },
                uRadius: { value: 0 },
                uOffset: { value: new Vector3(0, 0, 0) },
                uIsFullSphere: { value: false },
                uDepthScale: { value: 0.5 }
            },
            glslVersion: GLSL3
        });

        this.blending = NoBlending;
    }
}
