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

// Shadow-pass depth material that matches ImpostorShaderV0's surface so the
// shadow caster silhouette agrees with what the colour pass draws. We need
// the full TBN + per-frame xform machinery so that parallax in the fragment
// reproduces the same UV displacement as V0 — otherwise the alpha discard
// would cut a different silhouette than the visible card.
//
// We DO write gl_FragDepth here. Reason: the caster card orients toward
// the light (modelViewMatrix in the shadow pass is light_view * model),
// while the visible card orients toward the camera. The two cards
// intersect at the bounding-sphere centre plane but diverge — without
// per-pixel depth the caster's silhouette would be a thin disc at that
// plane and any visible fragment "below" the centre (in light space)
// would read as self-shadowed. Projecting the parallaxed surface point
// through MVP and writing it as gl_FragDepth fixes the silhouette: the
// shadow map records the actual depth of the impostor surface at each
// (light_x, light_y), matching what the visible card produces.
// Cost: this disables early-Z in the shadow pass, but shadow passes are
// usually fillrate-light and the geometry is just a card.
//
// Output is RGBA-packed depth — that's the format three.js's PCFShadowMap
// reads, and it's what the default MeshDepthMaterial would write here.
const shader_vx = `

    in vec2 uv;
    in vec3 position;

    out vec2 vUv;
    out vec3 vViewPos;

    // Card position and normal in OBJECT space — the fragment shader uses
    // these to project the parallaxed surface point through MVP for the
    // gl_FragDepth write. Position interpolates between vertices; the
    // normal is the same for every card vertex so we mark it flat.
    out vec3 vCardPos_OS;
    flat out vec3 vCardNormal_OS;

    flat out vec2 vGridFloor;
    flat out vec4 vWeights;

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

    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;

        // For shadow casting, modelViewMatrix = light_view * model. The
        // camera in object space is therefore the LIGHT's position in
        // object space, so the impostor card naturally orients to face the
        // light — exactly what we want for a shadow caster.
        // For a directional light's ORTHOGRAPHIC shadow camera, three.js
        // positions the shadow camera at the light entity's world
        // position — which is independent of the light's actual ray
        // direction. Using "cameraPos_OS - uOffset" would index the atlas
        // by "direction to shadow camera position", which generally does
        // NOT equal the parallel light-ray direction (in our scene the
        // light sits at (30, 70, 30) but rays travel along (0.17, -1, 0.17),
        // which puts the camera position on the +X,+Z side while the
        // light arrives from -X,-Z — i.e. mirrored on X and Z). The
        // receiver (lit shader) uses the actual light direction for its
        // shadow coord, so we have to match it on the caster side too.
        //
        // Row 2 of modelViewMatrix encodes the camera's +Z axis in OS,
        // which for a three.js camera is "away from the looking direction"
        // — i.e. the direction OPPOSITE the way the light travels, which
        // is what we want for atlas indexing (the bake's D_frame =
        // direction-from-centre-to-bake-camera = -look-direction).
        vec3 pivotToCameraRay = normalize(vec3(
            modelViewMatrix[0][2],
            modelViewMatrix[1][2],
            modelViewMatrix[2][2]
        ));

        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);
        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;

        vCardPos_OS    = pos_OS;
        vCardNormal_OS = normal_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;

    in vec3 vCardPos_OS;
    flat in vec3 vCardNormal_OS;

    flat in vec2 vGridFloor;
    flat in vec4 vWeights;

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

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

    out vec4 color_out;

    // Needed in the fragment to re-project the parallax-shifted surface
    // point. three.js's WebGLRenderer sets both unconditionally for any
    // ShaderMaterial-derived material (RawShaderMaterial included).
    uniform mat4 modelViewMatrix;
    uniform mat4 projectionMatrix;

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

    // Mirrors three.js's packing.glsl chunk so the shadow map reader (PCF
    // sampler in shadowmap_pars_fragment) can unpack our output.
    const float PackUpscale = 256.0 / 255.0;
    const vec3 PackFactors = vec3(256.0 * 256.0 * 256.0, 256.0 * 256.0, 256.0);
    const float ShiftRight8 = 1.0 / 256.0;

    vec4 packDepthToRGBA(const in float v) {
        vec4 r = vec4(fract(v * PackFactors), v);
        r.yzw -= r.xyz * ShiftRight8;
        return r * PackUpscale;
    }

    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;
    }

    vec4 blend_4_frames(
        sampler2D tex, vec2 card_uv,
        vec2 gridFloor, vec4 w,
        vec4 x00, vec4 x10, vec4 x01, vec4 x11
    ) {
        vec2 frame_size = vec2(1.0 / uFrames);
        vec2 uv00 = clamp(apply_frame_xform(card_uv, x00), 0.0, 1.0);
        vec2 uv10 = clamp(apply_frame_xform(card_uv, x10), 0.0, 1.0);
        vec2 uv01 = clamp(apply_frame_xform(card_uv, x01), 0.0, 1.0);
        vec2 uv11 = clamp(apply_frame_xform(card_uv, x11), 0.0, 1.0);
        vec4 s00 = texture(tex, (gridFloor + vec2(0.0, 0.0) + uv00) * frame_size);
        vec4 s10 = texture(tex, (gridFloor + vec2(1.0, 0.0) + uv10) * frame_size);
        vec4 s01 = texture(tex, (gridFloor + vec2(0.0, 1.0) + uv01) * frame_size);
        vec4 s11 = texture(tex, (gridFloor + vec2(1.0, 1.0) + uv11) * frame_size);
        return s00 * w.x + s10 * w.y + s01 * w.z + s11 * w.w;
    }

    void main() {
        // Same parallax shift as V0 so the silhouette we test matches the
        // visible silhouette — otherwise the shadow would clip a slightly
        // different cutout than 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 base_uv = vUv;
        float depth = blend_4_frames(
            tGeometry, base_uv,
            vGridFloor, vWeights,
            vFrameXform00, vFrameXform10, vFrameXform01, vFrameXform11
        ).a;
        base_uv += (view_dir.xy / view_dir.z) * (depth - 0.5) * uDepthScale;
        base_uv = clamp(base_uv, 0.0, 1.0);

        vec4 texel_color = blend_4_frames(
            tBase, base_uv,
            vGridFloor, vWeights,
            vFrameXform00, vFrameXform10, vFrameXform01, vFrameXform11
        );

        if (texel_color.a <= 0.5) {
            discard;
        }

        // Re-sample depth at the parallax-shifted UV. We want the surface
        // depth at the post-parallax location, not the un-shifted one.
        // For a single-step parallax these usually round-trip to a
        // similar value, but the second sample keeps us self-consistent
        // when the parallax shift is non-trivial.
        float surface_depth = blend_4_frames(
            tGeometry, base_uv,
            vGridFloor, vWeights,
            vFrameXform00, vFrameXform10, vFrameXform01, vFrameXform11
        ).a;

        // Convention from BakeShaderStandard: depth=1 at the bounding
        // sphere's near plane (front, +radius along bake camera), 0.5 at
        // the centre (= the card plane), 0 at the far plane (-radius).
        // So (depth - 0.5) is the signed height along the card normal in
        // half-bounding-sphere units; multiply by 2*radius to get
        // object-space height.
        float height_OS = (surface_depth - 0.5) * 2.0 * uRadius;

        vec3 surface_OS = vCardPos_OS + height_OS * vCardNormal_OS;
        vec4 surface_clip = projectionMatrix * modelViewMatrix * vec4(surface_OS, 1.0);
        float surface_ndc_z = surface_clip.z / surface_clip.w;

        // gl_FragCoord.z / gl_FragDepth are in [0, 1] (assuming default
        // depth range). Map NDC z (-1..+1) to that range, then clamp so
        // out-of-frustum surface points don't write garbage into the
        // shadow map.
        gl_FragDepth = clamp(surface_ndc_z * 0.5 + 0.5, 0.0, 1.0);

        color_out = packDepthToRGBA(gl_FragDepth);
    }
`;

/**
 * Depth-pass companion to ImpostorShaderV0. Use this as
 * ShadedGeometry.depth_material so the impostor casts a shadow whose
 * silhouette matches the visible card.
 */
export class ImpostorShaderDepthV0 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
        });

        // Shadow pass writes a single RGBA depth value per fragment — no
        // need for blending.
        this.blending = NoBlending;
    }
}
