#define PNANOVDB_HLSL
#define PNANOVDB_ADDRESS_32
#include "PNanoVDB.hlsl"

cbuffer CB0 : register(b1)
{
    float4x4 _InverseViewProjection;
    float3 _CameraPos;
    float _DensityScale;
    float3 _BBoxMin;
    float _StepSize;
    float3 _BBoxMax;
    uint _MaxSteps;
};

StructuredBuffer<uint> buf : register(t1);

struct VSInput
{
    float2 position : POSITION;
    float2 texcoord : TEXCOORD0;
};

struct PSInput
{
    float4 position : SV_POSITION;
    float2 texcoord : TEXCOORD0;
    float3 worldPos : TEXCOORD1;
};

struct NanoVolume
{
    pnanovdb_grid_handle_t grid;
    pnanovdb_grid_type_t grid_type;
    pnanovdb_readaccessor_t acc;
};

void initVolume(inout NanoVolume volume)
{
    pnanovdb_grid_handle_t grid;
    grid.address.byte_offset = 0;
    
    pnanovdb_grid_type_t grid_type = pnanovdb_buf_read_uint32(buf, PNANOVDB_GRID_OFF_GRID_TYPE);
    pnanovdb_tree_handle_t tree = pnanovdb_grid_get_tree(buf, grid);
    pnanovdb_root_handle_t root = pnanovdb_tree_get_root(buf, tree);
    pnanovdb_readaccessor_t acc;
    pnanovdb_readaccessor_init(acc, root);

    volume.grid = grid;
    volume.grid_type = grid_type;
    volume.acc = acc;
}

float get_value_coord(inout pnanovdb_readaccessor_t acc, float3 pos)
{
    pnanovdb_vec3_t p = pos;
    pnanovdb_coord_t ijk = pnanovdb_hdda_pos_to_ijk(p);
    pnanovdb_address_t address = pnanovdb_readaccessor_get_value_address(PNANOVDB_GRID_TYPE_FLOAT, buf, acc, ijk);
    return pnanovdb_read_float(buf, address);
}

uint get_dim_coord(inout pnanovdb_readaccessor_t acc, float3 pos)
{
    pnanovdb_vec3_t p = pos;
    pnanovdb_coord_t ijk = pnanovdb_hdda_pos_to_ijk(p);
    return pnanovdb_readaccessor_get_dim(PNANOVDB_GRID_TYPE_FLOAT, buf, acc, ijk);
}

bool get_hdda_hit(inout pnanovdb_readaccessor_t acc, inout float tmin, float3 origin, float3 direction, float tmax, out float valueAtHit)
{
    pnanovdb_vec3_t p_origin = origin;
    pnanovdb_vec3_t p_direction = direction;
    float thit;
    bool hit = pnanovdb_hdda_tree_marcher(
        PNANOVDB_GRID_TYPE_FLOAT,
        buf,
        acc,
        p_origin, tmin,
        p_direction, tmax,
        thit,
        valueAtHit
    );
    tmin = thit;
    return hit;
}

PSInput MainVS(VSInput input)
{
    PSInput output;
    output.position = float4(input.position, 0.0, 1.0);
    output.texcoord = input.texcoord;
    
    float4 worldPosH = mul(_InverseViewProjection, float4(input.position, 0.5, 1.0));
    output.worldPos = worldPosH.xyz / worldPosH.w;
    
    return output;
}

bool intersectBox(float3 origin, float3 dir, float3 boxMin, float3 boxMax, out float tmin, out float tmax)
{
    float3 invDir = 1.0 / dir;
    float3 t1 = (boxMin - origin) * invDir;
    float3 t2 = (boxMax - origin) * invDir;
    
    tmin = max(max(min(t1.x, t2.x), min(t1.y, t2.y)), min(t1.z, t2.z));
    tmax = min(min(max(t1.x, t2.x), max(t1.y, t2.y)), max(t1.z, t2.z));
    
    return tmax >= tmin && tmax > 0;
}

float4 MainPS(PSInput input) : SV_TARGET
{
    float3 rayDir = normalize(input.worldPos - _CameraPos);
    
    float tmin = 0.01;
    float tmax = 1000.0;
    
    NanoVolume volume;
    initVolume(volume);
    
    float3 color = float3(0, 0, 0);
    float transmittance = 1.0;
    
    float not_used;
    bool hit = get_hdda_hit(volume.acc, tmin, _CameraPos, rayDir, tmax, not_used);
    if (!hit) { return float4(0, 0, 0, 0); }
    
    float skip = 0;
    
    for (uint i = 0; i < _MaxSteps; i++) {
        if (tmin >= tmax || transmittance < 0.01) break;
        
        float3 worldPos = _CameraPos + rayDir * tmin;
        
        uint dim = get_dim_coord(volume.acc, worldPos);
        if (dim > 1) {
            float skip_step = 15.0;
            tmin += skip_step;
            skip = skip_step;
            continue;
        }
        
        float density = get_value_coord(volume.acc, worldPos) * _DensityScale;
        
        if (density < 0.001) {
            float skip_step = 5.0;
            tmin += skip_step;
            skip = skip_step;
            continue;
        }
        
        if (skip > 0) {
            tmin -= skip * 0.8;
            skip = 0;
            continue;
        }
        
        float3 S = density * float3(1, 1, 1);
        color += transmittance * S * _StepSize;
        transmittance *= exp(-density * _StepSize);
        
        tmin += _StepSize;
    }
    
    float alpha = 1.0 - transmittance;
    return float4(color, alpha);
}