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

// Per-pixel variant of ImpostorShaderV0. The atlas projection (octahedral
// grid lookup, cell selection, bilinear weights, card TBN and per-frame UV
// transforms) is computed once per FRAGMENT instead of once per VERTEX.
//
// In V0 those quantities are flat varyings — they don't depend on the
// individual vertex anyway (they're a function of cameraPos_OS, which is
// constant across the card), so V0 computes them once at the vertex stage
// and forwards them. This version does the same math in the fragment
// shader. Visual output matches V0; the difference is purely where the
// work lands on the GPU.
//
// Same uniforms as V0 so they're swap-in compatible.
const shader_vx = `

    in vec2 uv;
    in vec3 position;

    out vec2 vUv;
    out vec3 vViewPos;

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

        // Same card-positioning math as V0. We still need to do this in
        // the vertex shader — the gl_Position has to come from somewhere,
        // and card orientation depends on cameraPos_OS regardless of
        // whether we sample per-pixel later.
        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);

        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;

        vec4 mvPosition = modelViewMatrix * vec4(pos_OS, 1.0);
        vViewPos = mvPosition.xyz;
        gl_Position = projectionMatrix * mvPosition;
    }
`;
const shader_fg = `
    precision highp float;
    precision highp int;

    in vec2 vUv;
    in vec3 vViewPos;

    out vec4 color_out;

    uniform mat4 modelViewMatrix;

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

    uniform vec3 uOffset;
    uniform float uRadius;
    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
        );
    }

    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() {
        // ---- per-pixel projection ----
        // Recompute everything V0 flat-varies. cameraPos_OS is constant
        // per-card so the math gives identical results to V0 — this just
        // demonstrates running the projection in the fragment.
        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);

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

        mat3 m3 = mat3(modelViewMatrix);
        vec3 tangent_VS  = normalize(m3 * tangent_OS);
        vec3 binormal_VS = normalize(m3 * binormal_OS);
        vec3 normal_VS   = normalize(m3 * normal_OS);

        vec4 xform00 = ComputeFrameXform(ray00, tangent_OS, binormal_OS);
        vec4 xform10 = ComputeFrameXform(ray10, tangent_OS, binormal_OS);
        vec4 xform01 = ComputeFrameXform(ray01, tangent_OS, binormal_OS);
        vec4 xform11 = ComputeFrameXform(ray11, tangent_OS, binormal_OS);

        // ---- same sampling/parallax as V0 ----
        vec3 view_dir_view = normalize(-vViewPos);
        vec3 view_dir = vec3(
            dot(tangent_VS,  view_dir_view),
            dot(binormal_VS, view_dir_view),
            dot(normal_VS,   view_dir_view)
        );

        vec2 base_uv = vUv;
        float depth = blend_4_frames(
            tGeometry, base_uv,
            gridFloor, weights,
            xform00, xform10, xform01, xform11
        ).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,
            gridFloor, weights,
            xform00, xform10, xform01, xform11
        );

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

        color_out = texel_color;
    }
`;

export class ImpostorShaderPerPixelV0 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 = CustomBlending;
        this.blendEquation = AddEquation;
        this.blendSrc = OneFactor;
        this.blendDst = OneMinusSrcAlphaFactor;
    }
}
