import {
    AddEquation,
    CustomBlending,
    GLSL3,
    OneFactor,
    OneMinusSrcAlphaFactor,
    RawShaderMaterial,
    Vector3
} from "three";

/*
 *
 * For ray projection using projection matrix : https://encreative.blogspot.com/2019/05/computing-ray-origin-and-direction-from.html
 */
const shader_vx = `

    in vec2 uv;
    in vec3 position;

    out vec2 vUv;
    // View-space position of this vertex. Interpolated per-fragment for the
    // tangent-space parallax computation in the fragment shader.
    out vec3 vViewPos;

    // The 4 atlas frames at the corners of the current grid cell are chosen
    // once per-impostor and forwarded as flat varyings so every fragment
    // samples the same set (without this, a perspective camera could flip
    // the chosen frames across the card and produce visible seams).
    //
    // We use bilinear interpolation over a full 2x2 cell, not a triangle
    // interpolation over 3 corners. Triangle interpolation has a
    // C1-discontinuous "diagonal switch" where one of the three contributing
    // frames swaps between (gridFloor+(1,0)) and (gridFloor+(0,1)); even
    // though the swap happens at zero weight (so the result is C0), the
    // weight DERIVATIVE jumps, which the eye perceives as a snap inside the
    // cell. Bilinear has no such switch — all 4 corners always contribute
    // with weights that are smooth products of fract(grid).
    flat out vec2 vGridFloor;
    flat out vec4 vWeights;       // (w00, w10, w01, w11)

    // Card's TBN basis in view space. The card is oriented so its NORMAL
    // points along the weighted blend of the 4 nearest baked view
    // directions, and TANGENT/BINORMAL line up with the *blended* bake
    // camera's right/up. Used by the fragment shader to transform the
    // view-space view-dir into the card's tangent frame.
    flat out vec3 vTangent;
    flat out vec3 vBinormal;
    flat out vec3 vNormal;

    // Per-frame card-UV -> texture-UV rotation matrices, one per cell
    // corner. Each frame was baked with its own camera right/up (from
    // three.js Camera.lookAt with up=(0,1,0)); near the octahedral pole
    // those right-axes can swing 90° or more between adjacent frames.
    // These matrices rotate the sample-UV into each frame's bake-camera
    // basis so the four corners line up consistently when blended.
    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;

    // ----- direction -> octahedral grid coord (range -1..+1) -----

    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)
    {
        if (uIsFullSphere) {
            return VecToSphereOct(v);
        } else {
            return VecToHemiSphereOct(v);
        }
    }

    // ----- octahedral grid coord (range 0..1) -> direction -----
    //
    // Inverse of VectorToGrid (after the 0.5/2.0 remap). Given the
    // normalised position of a frame in the atlas, return the world-space
    // direction the bake camera was sitting at when it captured that frame.

    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)
    {
        if (uIsFullSphere) {
            return OctaSphereDec(coord);
        } else {
            return OctaHemiSphereDec(coord);
        }
    }

    // Frame index (gridFloor + corner offset) -> baked view direction.
    // Frame coords at the grid border can sit slightly outside the
    // [0, uFrames-1] range (e.g. gridFloor + (1,1) at the corner); clamp so
    // we don't feed invalid UVs to the octahedral decoder.
    vec3 FrameToRay(vec2 frame, vec2 framesMinusOne)
    {
        vec2 f = clamp(frame / framesMinusOne, 0.0, 1.0);
        return normalize(GridToVector(f));
    }

    // Build the 2x2 matrix that maps card-centred UV (vUv - 0.5) into the
    // texture-centred UV of a single baked frame. The matrix is just the
    // (tangent_OS, binormal_OS) basis expressed in the frame's own
    // (bake_right, bake_up) basis — i.e. how a step along the card's local
    // axes shows up in the texture the frame was rendered into.
    //
    // bake_right/bake_up follow three.js Camera.lookAt with up=(0,1,0).
    // We mirror its polar nudge so our shader basis matches what the
    // baker actually used when D_frame is parallel to up.
    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)
        );
    }

    // Bilinear weights over the 4 corners of the current cell, in the
    // order (w00, w10, w01, w11) matching gridFloor + (0,0)/(1,0)/(0,1)/(1,1).
    vec4 BilinearWeights(vec2 frac_uv)
    {
        vec2 omuv = vec2(1.0) - frac_uv;
        return vec4(
            omuv.x * omuv.y,    // w00
            frac_uv.x * omuv.y, // w10
            omuv.x * frac_uv.y, // w01
            frac_uv.x * frac_uv.y // w11
        );
    }

    void main() {
        vUv = uv;

        // 1. Octahedral atlas lookup.
        //    cameraPos_OS is the camera position in the impostor's local
        //    space — the same coordinate frame the meshes were in during
        //    the bake. The encoded direction is what picks which atlas
        //    frames we'll sample.
        vec3 cameraPos_OS = (inverse(modelViewMatrix) * vec4(0.0, 0.0, 0.0, 1.0)).xyz;
        // Direction from the bounding-sphere CENTRE (uOffset) — not from
        // the object origin — to the camera. The bake camera is placed at
        // center + radius*D_frame and the atlas is indexed by that D_frame,
        // so the runtime lookup must also be measured from the centre,
        // otherwise any model whose centre isn't at the origin (i.e. most
        // GLTFs) gets a biased octahedral lookup and ends up sampling the
        // wrong atlas frames.
        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;
        // Clamp gridFloor so the +1 corners never reference past the last
        // valid frame index when grid is at the upper edge.
        vec2 gridFloor = min(floor(grid), framesMinusOne - 1.0);
        vec4 weights = BilinearWeights(grid - gridFloor);

        vGridFloor = gridFloor;
        vWeights = weights;

        // 2. Decode each of the 4 cell-corner frame indices back into its
        //    bake-time view direction, then blend with the bilinear weights.
        //    projectedRay is the "effective" baked view direction the card
        //    is showing — the direction the depth and colour textures we're
        //    blending were captured along.
        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
        );

        // 3. Build the card's TBN in object-local space.
        //    NORMAL    = projectedRay (the effective bake direction).
        //    TANGENT   = cross(up, NORMAL) — matches three.js Camera.lookAt
        //                (which is what the baker used), so position.x of
        //                the cutout shape lines up with the texture's X.
        //    BINORMAL  = cross(NORMAL, TANGENT) — matches the bake camera's
        //                up axis, so position.y lines up with texture Y.
        //
        //    The bake used world-up (0,1,0). The bake's world frame IS our
        //    object-local frame (the meshes lived there during the bake),
        //    so we use object-local (0,1,0) here too. The fallback covers
        //    the degenerate pole case where NORMAL is parallel to up.
        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);

        // 4. Position the card vertex in object-local space.
        //    position.xy is in [-0.5, 0.5] (the centred cutout shape), so
        //    multiplying by card_diameter places it on a card whose half-
        //    width and half-height are uRadius — tangentially spanning the
        //    bounding sphere. uOffset is the bounding sphere centre in
        //    object-local space (captured at bake time).
        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;

        // 5. TBN in view space — what the fragment shader needs to convert
        //    its view-space view-dir into tangent space. For typical
        //    impostor transforms (rigid + uniform scale) mat3(modelView) is
        //    orthogonal up to a uniform scalar, so renormalising after the
        //    multiply is sufficient; non-uniform scale would require the
        //    inverse-transpose.
        mat3 m3 = mat3(modelViewMatrix);
        vTangent  = normalize(m3 * tangent_OS);
        vBinormal = normalize(m3 * binormal_OS);
        vNormal   = normalize(m3 * normal_OS);

        // 6. Per-frame UV rotation matrices. ComputeFrameXform handles the
        //    polar-singularity nudge in three.js's lookAt internally.
        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;

    // Cell base + bilinear weights for the 4 corners (w00, w10, w01, w11),
    // computed per-impostor in the vertex shader. flat = no interpolation.
    flat in vec2 vGridFloor;
    flat in vec4 vWeights;

    // Card's TBN in view space (perpendicular to the effective bake
    // direction). Used to express the view direction in the card's tangent
    // frame for parallax.
    flat in vec3 vTangent;
    flat in vec3 vBinormal;
    flat in vec3 vNormal;

    // Per-frame 2x2 rotation matrices (a, b, c, d) that map card-centred
    // UV into each frame's bake-camera basis. Each cell corner has its own.
    flat in vec4 vFrameXform00;
    flat in vec4 vFrameXform10;
    flat in vec4 vFrameXform01;
    flat in vec4 vFrameXform11;

    out vec4 color_out;

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

    // Apply a per-frame 2x2 rotation matrix to a card UV, rotating around
    // the texture centre (0.5, 0.5). xform = (a, b, c, d) packs row-major
    // [a b ; c d].
    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;
    }

    // Sample the same card-UV from the 4 cell-corner atlas frames, applying
    // each frame's own UV rotation so the textures land in a consistent
    // orientation, then blend with the bilinear weights. Used for both
    // depth and colour.
    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);
        // Clamp inside the unit square AFTER the per-frame rotation: at
        // sharp angles the rotation can push a corner UV outside [0,1] and
        // bleed into the neighbouring atlas tile (which is a completely
        // different bake direction). Clamping keeps the sample inside the
        // frame it's supposed to belong to.
        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() {
        // View direction in the card's tangent space. The card is oriented
        // perpendicular to the effective bake direction, not view-aligned,
        // so we project the view-space view-dir onto the (T, B, N) basis we
        // built in the vertex shader.
        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)
        );

        // One-step approximate parallax: sample the blended depth at the
        // un-corrected UV, then shift the UV by the height that depth
        // implies.
        //
        // Our bake writes depth = 1 at the near plane (front of the bounding
        // sphere, closest to the bake camera), depth = 0 at the far plane,
        // depth = 0.5 at the bounding sphere centre — i.e. the card plane.
        // So (depth - 0.5) is the signed height of the surface above the
        // card plane, in units where ±0.5 = ±radius. The card spans 2*radius
        // mapped to UV [0,1], so that height is already in card-UV units.
        //
        // Standard parallax-mapping formula: shift base_uv by V.xy / V.z * h,
        // where V is the view direction in tangent space and h is the
        // surface height above the card plane.
        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;

        // Keep the parallax-shifted sample inside the current frame's tile.
        // Without this, large shifts at oblique angles bleed colour from the
        // neighbouring atlas frame (a completely different bake direction).
        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;
        }

        color_out = texel_color;
    }
`;

export class ImpostorShaderV0 extends RawShaderMaterial {
    constructor() {
        super({
            fragmentShader: shader_fg,
            vertexShader: shader_vx,
            uniforms: {
                /**
                 * RGB + Alpha
                 */
                tBase: {
                    value: null
                },
                /**
                 * Normal+Depth
                 */
                tGeometry: {
                    value: null
                },
                /**
                 * Material properties: Occlusion, Roughness, Metalness
                 * Alpha unused
                 */
                tMaterial: {
                    value: null
                },
                /**
                 * Number of frames
                 */
                uFrames: {
                    value: 0
                },
                /**
                 * Radius of bounding sphere of the impostor
                 */
                uRadius: {
                    value: 0
                },
                /**
                 * Impostor offset
                 */
                uOffset: {
                    value: new Vector3(0, 0, 0)
                },
                uIsFullSphere: {
                    value: false
                },
                /**
                 * Strength of the tangent-space parallax offset. The
                 * geometrically-correct value for a true sphere surface is
                 * about 1.0, but in practice 0.3–0.5 produces a more
                 * believable result because the bake only contains the
                 * silhouette and front-facing geometry — strong parallax
                 * starts revealing missing/occluded parts at oblique angles.
                 * Range: [0, 1].
                 */
                uDepthScale:{
                    value: 0.5
                }
            },
            glslVersion: GLSL3
        });

        // Save some effort by disabling blending
        this.blending = CustomBlending;
        this.blendEquation = AddEquation;
        this.blendSrc = OneFactor;
        this.blendDst = OneMinusSrcAlphaFactor;

    }
}
