#ifndef __IMPOSTOR_INC
#define __IMPOSTOR_INC

#include "UnityCG.cginc"
#include "vertex_deformation.hlsl"

SamplerState bilinear_clamp_s;
Texture2D _ImpostorAtlas;
int _GridResolution;
float _Cutoff;
float3 _ImpostorMainCameraPos;

#ifndef IMPOSTOR_SHADOW_PASS
float4 _Color;
float _DebugMode;
#endif

float2 HemiOctEncode(float3 N) {
    N.y = max(N.y, 1e-4);
    N = normalize(N);
    float3 p = hemi_octahedron_to_plane(N, 0, float3(1,0,0), float3(0,1,0), 1);
    return float2(p.x, p.z);
}

float3 HemiOctDecode(float2 uv) {
    return normalize(plane_to_hemi_octahedron(float3(uv.x, 0, uv.y), 0, float3(1,0,0), float3(0,1,0), 1));
}

void BillboardBasis(float3 fwd, out float3 right, out float3 up) {
    right = abs(fwd.y) > 0.999 ? float3(-1,0,0) : normalize(cross(float3(0,1,0), fwd));
    up = cross(fwd, right);
}

float2 GridFromDir(float3 viewDir, float gridRes) {
    float2 uv = HemiOctEncode(viewDir) * 0.5 + 0.5;
    return clamp(uv * (gridRes - 1), 0, gridRes - 1);
}

float3 DirFromCell(float2 cell, float gridRes) {
    float2 uv = cell / max(1.0, gridRes - 1) * 2.0 - 1.0;
    return HemiOctDecode(uv);
}

float4 SampleAtlas(float2 uv, float2 cell) {
    return _ImpostorAtlas.Sample(bilinear_clamp_s, (cell + uv) / _GridResolution);
}

// Compute blend weights for 3 frames in a grid cell triangle
// .xyz = weights for frames 1,2,3
// .w = which triangle (0 = lower-left, 1 = upper-right)
float4 QuadBlendWeights(float2 frac) {
    float4 res;
    res.x = min(1.0 - frac.x, 1.0 - frac.y);
    res.y = abs(frac.x - frac.y);
    res.z = min(frac.x, frac.y);
    res.w = ceil(frac.x - frac.y);
    return res;
}

// Compute UV on a virtual plane facing frameDir
float2 VirtualPlaneUV(float3 frameDir, float3 pivotToCam, float3 vertexToCam, float size) {
    float3 planeN = normalize(frameDir);
    float3 up = abs(planeN.y) > 0.999 ? float3(0,0,1) : float3(0,1,0);
    float3 planeX = normalize(cross(planeN, up));
    float3 planeY = normalize(cross(planeX, planeN));

    float projPivot = dot(planeN, pivotToCam);
    float projVertex = dot(planeN, vertexToCam);
    float ratio = projPivot / projVertex;

    float3 offset = vertexToCam * ratio - pivotToCam;

    float2 uv = float2(dot(planeX, offset), dot(planeY, offset));
    uv /= (2.0f * size);
    uv *= -1.0f;
    uv += 0.5;
    return uv;
}

struct appdata {
    float4 vertex : POSITION;
    UNITY_VERTEX_INPUT_INSTANCE_ID
#ifdef IMPOSTOR_SHADOW_PASS
    float3 normal : NORMAL;
#endif
};

struct v2f {
#ifdef IMPOSTOR_SHADOW_PASS
    V2F_SHADOW_CASTER;
#else
    float4 pos : SV_POSITION;
    UNITY_FOG_COORDS(7)
#endif
    float2 uv1 : TEXCOORD1;
    float2 uv2 : TEXCOORD2;
    float2 uv3 : TEXCOORD3;
    nointerpolation float2 cell1 : TEXCOORD4;
    nointerpolation float2 cell2 : TEXCOORD5;
    nointerpolation float2 cell3 : TEXCOORD6;
    float3 blendWeights : TEXCOORD8;
    UNITY_VERTEX_OUTPUT_STEREO
};

v2f vert(appdata v) {
    v2f o;
    UNITY_SETUP_INSTANCE_ID(v);
    UNITY_INITIALIZE_VERTEX_OUTPUT_STEREO(o);

    float3 center = mul(unity_ObjectToWorld, float4(0,0,0,1)).xyz;
    float3 scale = float3(
        length(unity_ObjectToWorld._m00_m10_m20),
        length(unity_ObjectToWorld._m01_m11_m21),
        length(unity_ObjectToWorld._m02_m12_m22));

#ifdef IMPOSTOR_SHADOW_PASS
    float3 camPos = _ImpostorMainCameraPos;
#else
    float3 camPos = _WorldSpaceCameraPos;
#endif

    // View direction in object space
    float3 viewOS = normalize(mul((float3x3)unity_WorldToObject, normalize(camPos - center)));

    // Get continuous grid position and find the 3 frames
    float gridRes = (float)_GridResolution;
    float2 grid = GridFromDir(viewOS, gridRes);
    float2 gridFloor = floor(grid);
    float2 gridFrac = frac(grid);

    // Compute blend weights and determine which triangle
    float4 weights = QuadBlendWeights(gridFrac);
    o.blendWeights = weights.xyz;

    // Frame 1: base corner
    o.cell1 = gridFloor;

    // Frame 2: depends on which triangle (lower-left vs upper-right)
    // weights.w = 0: lower-left triangle -> (0,1) offset
    // weights.w = 1: upper-right triangle -> (1,0) offset
    o.cell2 = clamp(gridFloor + lerp(float2(0,1), float2(1,0), weights.w), 0, gridRes - 1);

    // Frame 3: diagonal corner
    o.cell3 = clamp(gridFloor + float2(1,1), 0, gridRes - 1);

    // Get directions for each frame
    float3 dir1 = DirFromCell(o.cell1, gridRes);
    float3 dir2 = DirFromCell(o.cell2, gridRes);
    float3 dir3 = DirFromCell(o.cell3, gridRes);

    // Billboard facing the camera direction (not snapped)
    float3 viewWS = normalize(camPos - center);
    float3 right, up;
    BillboardBasis(viewWS, right, up);
    float3 worldPos = center + v.vertex.x * right * scale.x + v.vertex.y * up * scale.y;

#ifdef IMPOSTOR_SHADOW_PASS
    v.vertex = mul(unity_WorldToObject, float4(worldPos, 1));
    v.normal = -viewWS;
    TRANSFER_SHADOW_CASTER_NORMALOFFSET(o)
#else
    o.pos = mul(UNITY_MATRIX_VP, float4(worldPos, 1));
    UNITY_TRANSFER_FOG(o, o.pos);
#endif

    // Compute virtual plane UVs for all 3 frames
    // pivotToCam and vertexToCam in object space
    float3 pivotToCamOS = mul((float3x3)unity_WorldToObject, camPos - center);
    float3 vertexPosOS = mul((float3x3)unity_WorldToObject, worldPos - center);
    float3 vertexToCamOS = pivotToCamOS - vertexPosOS;

    // size = 0.5 because the mesh is a unit quad (-0.5 to 0.5)
    o.uv1 = VirtualPlaneUV(dir1, pivotToCamOS, vertexToCamOS, 0.5);
    o.uv2 = VirtualPlaneUV(dir2, pivotToCamOS, vertexToCamOS, 0.5);
    o.uv3 = VirtualPlaneUV(dir3, pivotToCamOS, vertexToCamOS, 0.5);

    return o;
}

float4 BlendFrames(v2f i) {
    float4 s1 = SampleAtlas(i.uv1, i.cell1);
    float4 s2 = SampleAtlas(i.uv2, i.cell2);
    float4 s3 = SampleAtlas(i.uv3, i.cell3);
    return s1 * i.blendWeights.x + s2 * i.blendWeights.y + s3 * i.blendWeights.z;
}

#ifdef IMPOSTOR_SHADOW_PASS
float4 frag(v2f i) : SV_Target {
    float4 col = BlendFrames(i);
    clip(col.a - _Cutoff);
    SHADOW_CASTER_FRAGMENT(i)
}
#else
float4 frag(v2f i) : SV_Target {
    float4 col = BlendFrames(i);

    if (_DebugMode > 0.5)
        return float4(i.blendWeights.xy, 0, 1);

    col *= _Color;
    clip(col.a - _Cutoff);
    UNITY_APPLY_FOG(i.fogCoord, col);
    return col;
}
#endif

#endif