enable wgpu_ray_query;

#define_import_path bevy_solari::scene_bindings

#import bevy_pbr::lighting::perceptualRoughnessToRoughness
#import bevy_pbr::pbr_functions::calculate_tbn_mikktspace

struct InstanceGeometryIds {
    vertex_buffer_id: u32,
    vertex_buffer_offset: u32,
    index_buffer_id: u32,
    index_buffer_offset: u32,
    triangle_count: u32,
}

struct VertexBuffer { vertices: array<PackedVertex> }

struct IndexBuffer { indices: array<u32> }

struct PackedVertex {
    a: vec4<f32>,
    b: vec4<f32>,
    tangent: vec4<f32>,
}

struct Vertex {
    position: vec3<f32>,
    normal: vec3<f32>,
    uv: vec2<f32>,
    tangent: vec4<f32>,
}

fn unpack_vertex(packed: PackedVertex) -> Vertex {
    var vertex: Vertex;
    vertex.position = packed.a.xyz;
    vertex.normal = vec3(packed.a.w, packed.b.xy);
    vertex.uv = packed.b.zw;
    vertex.tangent = packed.tangent;
    return vertex;
}

struct SurfaceUniform {
    shader_id: u32,
    flags: u32,
    alpha_cutoff: f32,
    abi_version: u32,
    params: array<vec4<f32>, 16>,
    uv_transforms: array<vec4<f32>, 8>,
}

struct Material {
    texture_ids: vec4<u32>,
    surface_texture_ids_0: vec4<u32>,
    surface_texture_ids_1: vec4<u32>,
    material_info: vec4<u32>,
    base_color: vec4<f32>,
    emissive_roughness: vec4<f32>,
    pbr: vec4<f32>,
    surface_uniform: SurfaceUniform,
}

const TEXTURE_MAP_NONE = 0xFFFFFFFFu;
const MATERIAL_FLAG_ALPHA_CUTOUT = 1u << 0u;
const MATERIAL_FLAG_SURFACE = 1u << 1u;
const BLACKSITE_SURFACE_UNLIT = 1u;

const MIRROR_ROUGHNESS_THRESHOLD = 0.001f;

struct LightSource {
    kind: u32, // 1 bit for kind, 31 bits for extra data
    id: u32,
}

const LIGHT_SOURCE_KIND_EMISSIVE_MESH = 0u;
const LIGHT_SOURCE_KIND_DIRECTIONAL = 1u;

struct DirectionalLight {
    direction_to_light: vec3<f32>,
    cos_theta_max: f32,
    luminance: vec3<f32>,
    inverse_pdf: f32,
}

const LIGHT_NOT_PRESENT_THIS_FRAME = 0xFFFFFFFFu;

@group(0) @binding(0) var<storage> vertex_buffers: binding_array<VertexBuffer>;
@group(0) @binding(1) var<storage> index_buffers: binding_array<IndexBuffer>;
@group(0) @binding(2) var textures: binding_array<texture_2d<f32>>;
@group(0) @binding(3) var samplers: binding_array<sampler>;
@group(0) @binding(4) var<storage> materials: array<Material>;
@group(0) @binding(5) var tlas: acceleration_structure;
@group(0) @binding(6) var<storage> transforms: array<mat4x4<f32>>; // TODO: Use mat3x4<f32>?
@group(0) @binding(7) var<storage> previous_frame_transforms: array<mat4x4<f32>>; // TODO: Use mat3x4<f32>?
@group(0) @binding(8) var<storage> geometry_ids: array<InstanceGeometryIds>;
@group(0) @binding(9) var<storage> material_ids: array<u32>; // TODO: Store material_id in instance_custom_index instead?
@group(0) @binding(10) var<storage> light_sources: array<LightSource>;
@group(0) @binding(11) var<storage> directional_lights: array<DirectionalLight>;
@group(0) @binding(12) var<storage> previous_frame_light_id_translations: array<u32>;
@group(0) @binding(13) var brdf_dfg_lut: texture_2d<f32>;
@group(0) @binding(14) var brdf_dfg_lut_sampler: sampler;

const RAY_T_MIN = 0.001f;
const RAY_T_MAX = 100000.0f;

const RAY_NO_CULL = 0xFFu;

fn trace_ray(ray_origin: vec3<f32>, ray_direction: vec3<f32>, ray_t_min: f32, ray_t_max: f32, ray_flag: u32) -> RayIntersection {
    let ray = RayDesc(ray_flag, RAY_NO_CULL, ray_t_min, ray_t_max, ray_origin, ray_direction);
    var rq: ray_query;
    rayQueryInitialize(&rq, tlas, ray);
    while (rayQueryProceed(&rq)) {
        let candidate = rayQueryGetCandidateIntersection(&rq);
        if candidate.kind == RAY_QUERY_INTERSECTION_TRIANGLE && candidate_passes_alpha(candidate) {
            rayQueryConfirmIntersection(&rq);
        }
    }
    return rayQueryGetCommittedIntersection(&rq);
}

fn sample_texture(id: u32, uv: vec2<f32>) -> vec4<f32> {
    return textureSampleLevel(textures[id], samplers[id], uv, 0.0); // TODO: Mipmap
}

struct SurfaceInput {
    uv0: vec2<f32>,
    world_position: vec3<f32>,
    world_normal: vec3<f32>,
}

struct SurfaceParams {
    lanes: array<vec4<f32>, 16>,
}

struct SurfaceSamples {
    values: array<vec4<f32>, 8>,
}

struct Surface {
    base_color: vec4<f32>,
    normal_ts: vec3<f32>,
    emissive: vec3<f32>,
    metallic: f32,
    perceptual_roughness: f32,
    reflectance: f32,
    occlusion: f32,
    model: u32,
}

fn surface_default() -> Surface {
    var surface: Surface;
    surface.base_color = vec4<f32>(1.0);
    surface.normal_ts = vec3<f32>(0.0, 0.0, 1.0);
    surface.emissive = vec3<f32>(0.0);
    surface.metallic = 0.0;
    surface.perceptual_roughness = 0.5;
    surface.reflectance = 0.5;
    surface.occlusion = 1.0;
    surface.model = 0u;
    return surface;
}

fn surface_texture_id(material: Material, index: u32) -> u32 {
    if index < 4u {
        return material.surface_texture_ids_0[index];
    }
    return material.surface_texture_ids_1[index - 4u];
}

fn surface_samples(material: Material, uv: vec2<f32>) -> SurfaceSamples {
    var samples: SurfaceSamples;
    for (var index = 0u; index < 8u; index += 1u) {
        let transform = material.surface_uniform.uv_transforms[index];
        let transformed_uv = uv * transform.xy + transform.zw;
        let texture_id = surface_texture_id(material, index);
        samples.values[index] = vec4<f32>(1.0);
        if texture_id != TEXTURE_MAP_NONE {
            samples.values[index] = sample_texture(texture_id, transformed_uv);
        }
    }
    return samples;
}

// BLACKSITE_SURFACE_EVALUATOR_LIBRARY

fn evaluate_surface(
    shader_id: u32,
    input: SurfaceInput,
    params: SurfaceParams,
    samples: SurfaceSamples,
) -> Surface {
// BLACKSITE_SURFACE_EVALUATOR_DISPATCH
    return surface_default();
}

struct ResolvedMaterial {
    base_color: vec3<f32>,
    emissive: vec3<f32>,
    reflectance: f32,
    perceptual_roughness: f32,
    roughness: f32,
    metallic: f32,
}

struct ResolvedSurface {
    material: ResolvedMaterial,
    normal_ts: vec3<f32>,
    alpha: f32,
}

struct ResolvedRayHitFull {
    world_position: vec3<f32>,
    previous_frame_world_position: vec3<f32>,
    world_normal: vec3<f32>,
    geometric_world_normal: vec3<f32>,
    world_tangent: vec4<f32>,
    uv: vec2<f32>,
    triangle_area: f32,
    triangle_count: u32,
    material: ResolvedMaterial,
}

fn resolve_standard_surface(material: Material, uv: vec2<f32>) -> ResolvedSurface {
    var m: ResolvedMaterial;

    var base_color = material.base_color;
    if material.texture_ids.y != TEXTURE_MAP_NONE {
        base_color *= sample_texture(material.texture_ids.y, uv);
    }
    m.base_color = base_color.rgb;

    m.emissive = material.emissive_roughness.rgb;
    if material.texture_ids.z != TEXTURE_MAP_NONE {
        m.emissive *= sample_texture(material.texture_ids.z, uv).rgb;
    }

    m.reflectance = material.pbr.y;

    m.perceptual_roughness = material.emissive_roughness.w;
    m.metallic = material.pbr.x;
    if material.texture_ids.w != TEXTURE_MAP_NONE {
        let metallic_roughness = sample_texture(material.texture_ids.w, uv).rgb;
        m.perceptual_roughness *= metallic_roughness.g;
        m.metallic *= metallic_roughness.b;
    }

    m.roughness = m.perceptual_roughness * m.perceptual_roughness;

    return ResolvedSurface(m, vec3<f32>(0.0, 0.0, 1.0), base_color.a);
}

fn resolve_surface(
    material: Material,
    uv: vec2<f32>,
    world_position: vec3<f32>,
    world_normal: vec3<f32>,
) -> ResolvedSurface {
    if (material.material_info.x & MATERIAL_FLAG_SURFACE) == 0u {
        return resolve_standard_surface(material, uv);
    }

    let input = SurfaceInput(uv, world_position, world_normal);
    let params = SurfaceParams(material.surface_uniform.params);
    let surface = evaluate_surface(
        material.surface_uniform.shader_id,
        input,
        params,
        surface_samples(material, uv),
    );
    var m: ResolvedMaterial;
    m.base_color = surface.base_color.rgb * clamp(surface.occlusion, 0.0, 1.0);
    m.emissive = surface.emissive;
    m.reflectance = clamp(surface.reflectance, 0.0, 1.0);
    m.perceptual_roughness = clamp(surface.perceptual_roughness, 0.001, 1.0);
    m.roughness = m.perceptual_roughness * m.perceptual_roughness;
    m.metallic = clamp(surface.metallic, 0.0, 1.0);
    if surface.model == BLACKSITE_SURFACE_UNLIT {
        m.emissive += surface.base_color.rgb;
        m.base_color = vec3<f32>(0.0);
        m.metallic = 0.0;
        m.perceptual_roughness = 1.0;
        m.roughness = 1.0;
    }
    return ResolvedSurface(m, surface.normal_ts, surface.base_color.a);
}

fn resolve_ray_hit_full(ray_hit: RayIntersection) -> ResolvedRayHitFull {
    let barycentrics = vec3(1.0 - ray_hit.barycentrics.x - ray_hit.barycentrics.y, ray_hit.barycentrics);
    return resolve_triangle_data_full(ray_hit.instance_index, ray_hit.primitive_index, barycentrics);
}

fn load_vertices(instance_geometry_ids: InstanceGeometryIds, triangle_id: u32) -> array<Vertex, 3> {
    let index_buffer = &index_buffers[instance_geometry_ids.index_buffer_id].indices;
    let vertex_buffer = &vertex_buffers[instance_geometry_ids.vertex_buffer_id].vertices;

    let indices_i = (triangle_id * 3u) + vec3(0u, 1u, 2u) + instance_geometry_ids.index_buffer_offset;
    let indices = vec3((*index_buffer)[indices_i.x], (*index_buffer)[indices_i.y], (*index_buffer)[indices_i.z]) + instance_geometry_ids.vertex_buffer_offset;

    return array<Vertex, 3>(
        unpack_vertex((*vertex_buffer)[indices.x]),
        unpack_vertex((*vertex_buffer)[indices.y]),
        unpack_vertex((*vertex_buffer)[indices.z])
    );
}

fn transform_positions(transform: mat4x4<f32>, vertices: array<Vertex, 3>) -> array<vec3<f32>, 3> {
    return array<vec3<f32>, 3>(
        (transform * vec4(vertices[0].position, 1.0)).xyz,
        (transform * vec4(vertices[1].position, 1.0)).xyz,
        (transform * vec4(vertices[2].position, 1.0)).xyz
    );
}

fn candidate_passes_alpha(candidate: RayIntersection) -> bool {
    let material = materials[material_ids[candidate.instance_index]];
    if (material.material_info.x & MATERIAL_FLAG_ALPHA_CUTOUT) == 0u {
        return true;
    }

    let barycentrics = vec3(
        1.0 - candidate.barycentrics.x - candidate.barycentrics.y,
        candidate.barycentrics,
    );
    let instance_geometry_ids = geometry_ids[candidate.instance_index];
    let vertices = load_vertices(instance_geometry_ids, candidate.primitive_index);
    let transform = transforms[candidate.instance_index];
    let world_vertices = transform_positions(transform, vertices);
    let world_position = mat3x3(world_vertices[0], world_vertices[1], world_vertices[2]) * barycentrics;
    let uv = mat3x2(vertices[0].uv, vertices[1].uv, vertices[2].uv) * barycentrics;
    let local_normal = mat3x3(vertices[0].normal, vertices[1].normal, vertices[2].normal) * barycentrics;
    let world_normal = normalize(mat3x3(transform[0].xyz, transform[1].xyz, transform[2].xyz) * local_normal);
    let resolved = resolve_surface(material, uv, world_position, world_normal);
    let cutoff = select(
        material.pbr.z,
        material.surface_uniform.alpha_cutoff,
        (material.material_info.x & MATERIAL_FLAG_SURFACE) != 0u,
    );
    return resolved.alpha >= cutoff;
}

fn resolve_triangle_data_full(instance_id: u32, triangle_id: u32, barycentrics: vec3<f32>) -> ResolvedRayHitFull {
    let material_id = material_ids[instance_id];
    let material = materials[material_id];

    let transform = transforms[instance_id];
    let previous_frame_transform = previous_frame_transforms[instance_id];

    let instance_geometry_ids = geometry_ids[instance_id];
    let vertices = load_vertices(instance_geometry_ids, triangle_id);

    let world_vertices = transform_positions(transform, vertices);
    let world_position = mat3x3(world_vertices[0], world_vertices[1], world_vertices[2]) * barycentrics;

    let previous_frame_world_vertices = transform_positions(previous_frame_transform, vertices);
    let previous_frame_world_position = mat3x3(previous_frame_world_vertices[0], previous_frame_world_vertices[1], previous_frame_world_vertices[2]) * barycentrics;

    let uv = mat3x2(vertices[0].uv, vertices[1].uv, vertices[2].uv) * barycentrics;

    let local_tangent = mat3x3(vertices[0].tangent.xyz, vertices[1].tangent.xyz, vertices[2].tangent.xyz) * barycentrics;
    let world_tangent = vec4(
        normalize(mat3x3(transform[0].xyz, transform[1].xyz, transform[2].xyz) * local_tangent),
        vertices[0].tangent.w,
    );

    let local_normal = mat3x3(vertices[0].normal, vertices[1].normal, vertices[2].normal) * barycentrics; // TODO: Use barycentric lerp, ray_hit.object_to_world, cross product geo normal
    var world_normal = normalize(mat3x3(transform[0].xyz, transform[1].xyz, transform[2].xyz) * local_normal);
    let geometric_world_normal = world_normal;
    let resolved_surface = resolve_surface(material, uv, world_position, world_normal);
    if (material.material_info.x & MATERIAL_FLAG_SURFACE) != 0u {
        let TBN = calculate_tbn_mikktspace(world_normal, world_tangent);
        let Nt = normalize(resolved_surface.normal_ts);
        world_normal = normalize(Nt.x * TBN[0] + Nt.y * TBN[1] + Nt.z * TBN[2]);
    } else if material.texture_ids.x != TEXTURE_MAP_NONE {
        let TBN = calculate_tbn_mikktspace(world_normal, world_tangent);
        let T = TBN[0];
        let B = TBN[1];
        let N = TBN[2];
        let Nt = sample_texture(material.texture_ids.x, uv).xyz;
        world_normal = normalize(Nt.x * T + Nt.y * B + Nt.z * N);
    }

    let triangle_edge0 = world_vertices[0] - world_vertices[1];
    let triangle_edge1 = world_vertices[0] - world_vertices[2];
    let triangle_area = length(cross(triangle_edge0, triangle_edge1)) / 2.0;

    return ResolvedRayHitFull(
        world_position,
        previous_frame_world_position,
        world_normal,
        geometric_world_normal,
        world_tangent,
        uv,
        triangle_area,
        instance_geometry_ids.triangle_count,
        resolved_surface.material,
    );
}
