#define_import_path nightshade::area_lighting

#import nightshade::material_sampling::sample_srgb_index

struct AreaLight {
    position: vec4<f32>,
    right: vec4<f32>,
    up: vec4<f32>,
    color: vec4<f32>,
    shape: u32,
    range: f32,
    radius: f32,
    two_sided: u32,
    shadow_index: i32,
    emissive_layer: u32,
    area_light_pad_a: f32,
    area_light_pad_b: f32,
};

struct AreaUniforms {
    count: u32,
    area_uniforms_pad_a: u32,
    area_uniforms_pad_b: u32,
    area_uniforms_pad_c: u32,
};

@group(1) @binding(10)
var<storage, read> area_lights: array<AreaLight>;

@group(1) @binding(11)
var<uniform> area_uniforms: AreaUniforms;

const AREA_PI: f32 = 3.14159265359;
const AREA_SHAPE_RECTANGLE: u32 = 0u;
const AREA_SHAPE_DISK: u32 = 1u;
const AREA_SHAPE_SPHERE: u32 = 2u;
const AREA_SHAPE_TUBE: u32 = 3u;
const AREA_EMISSIVE_LAYER_NONE: u32 = 0xFFFFFFFFu;

fn area_light_count() -> u32 {
    return area_uniforms.count;
}

fn area_light_at(index: u32) -> AreaLight {
    return area_lights[index];
}

fn area_normal(light: AreaLight) -> vec3<f32> {
    return vec3<f32>(light.position.w, light.right.w, light.up.w);
}

fn area_d_ggx(n_dot_h: f32, roughness: f32) -> f32 {
    let alpha = roughness * roughness;
    let alpha_sq = alpha * alpha;
    let denom = n_dot_h * n_dot_h * (alpha_sq - 1.0) + 1.0;
    return alpha_sq / max(AREA_PI * denom * denom, 1e-7);
}

fn area_v_smith(n_dot_v: f32, n_dot_l: f32, roughness: f32) -> f32 {
    let alpha = roughness * roughness;
    let k = alpha * 0.5;
    let gv = n_dot_v / max(n_dot_v * (1.0 - k) + k, 1e-5);
    let gl = n_dot_l / max(n_dot_l * (1.0 - k) + k, 1e-5);
    return gv * gl;
}

fn area_fresnel(cos_theta: f32, f0: vec3<f32>) -> vec3<f32> {
    let f = pow(clamp(1.0 - cos_theta, 0.0, 1.0), 5.0);
    return f0 + (vec3<f32>(1.0) - f0) * f;
}

fn area_range_attenuation(range: f32, distance: f32) -> f32 {
    if range <= 0.0 {
        return 1.0;
    }
    let window = clamp(1.0 - pow(distance / range, 4.0), 0.0, 1.0);
    return window * window;
}

fn area_integrate_edge(v1: vec3<f32>, v2: vec3<f32>) -> f32 {
    let x = dot(v1, v2);
    let y = abs(x);
    let a = 0.8543985 + (0.4965155 + 0.0145206 * y) * y;
    let b = 3.4175940 + (4.1616724 + y) * y;
    let v = a / b;
    var theta_sintheta = v;
    if x <= 0.0 {
        theta_sintheta = 0.5 * inverseSqrt(max(1.0 - x * x, 1e-7)) - v;
    }
    return cross(v1, v2).z * theta_sintheta;
}

fn area_polygon_form_factor(
    surface_normal: vec3<f32>,
    view: vec3<f32>,
    surface_pos: vec3<f32>,
    corner_a: vec3<f32>,
    corner_b: vec3<f32>,
    corner_c: vec3<f32>,
    corner_d: vec3<f32>,
    two_sided: bool,
) -> f32 {
    let tangent = normalize(view - surface_normal * dot(view, surface_normal));
    let bitangent = cross(surface_normal, tangent);
    let basis = transpose(mat3x3<f32>(tangent, bitangent, surface_normal));

    let l0 = normalize(basis * (corner_a - surface_pos));
    let l1 = normalize(basis * (corner_b - surface_pos));
    let l2 = normalize(basis * (corner_c - surface_pos));
    let l3 = normalize(basis * (corner_d - surface_pos));

    var sum = 0.0;
    sum = sum + area_integrate_edge(l0, l1);
    sum = sum + area_integrate_edge(l1, l2);
    sum = sum + area_integrate_edge(l2, l3);
    sum = sum + area_integrate_edge(l3, l0);

    if two_sided {
        sum = abs(sum);
    } else {
        sum = max(0.0, sum);
    }
    return sum / (2.0 * AREA_PI);
}

struct AreaQuad {
    center: vec3<f32>,
    axis_u: vec3<f32>,
    axis_v: vec3<f32>,
    normal: vec3<f32>,
    corner_a: vec3<f32>,
    corner_b: vec3<f32>,
    corner_c: vec3<f32>,
    corner_d: vec3<f32>,
    extent: f32,
};

fn area_build_quad(light: AreaLight, surface_pos: vec3<f32>) -> AreaQuad {
    var quad: AreaQuad;
    let center = light.position.xyz;
    quad.center = center;

    var axis_u = light.right.xyz;
    var axis_v = light.up.xyz;
    var normal = area_normal(light);

    if light.shape == AREA_SHAPE_SPHERE {
        let face = normalize(surface_pos - center);
        let tangent = normalize(cross(select(vec3<f32>(0.0, 1.0, 0.0), vec3<f32>(1.0, 0.0, 0.0), abs(face.y) > 0.99), face));
        let bitangent = cross(tangent, face);
        axis_u = tangent * light.radius;
        axis_v = bitangent * light.radius;
        normal = face;
    }

    if light.shape == AREA_SHAPE_TUBE {
        let axis = normalize(light.right.xyz);
        let to_surface = surface_pos - center;
        var perp = to_surface - axis * dot(to_surface, axis);
        if dot(perp, perp) < 1e-5 {
            perp = area_normal(light);
        }
        perp = normalize(perp);
        normal = perp;
        axis_u = light.right.xyz;
        axis_v = normalize(cross(axis, perp)) * light.radius;
    }

    quad.axis_u = axis_u;
    quad.axis_v = axis_v;
    quad.normal = normal;
    quad.corner_a = center - axis_u - axis_v;
    quad.corner_b = center + axis_u - axis_v;
    quad.corner_c = center + axis_u + axis_v;
    quad.corner_d = center - axis_u + axis_v;
    quad.extent = max(length(axis_u), length(axis_v));
    return quad;
}

fn area_representative_point(
    quad: AreaQuad,
    surface_pos: vec3<f32>,
    surface_normal: vec3<f32>,
    view: vec3<f32>,
) -> vec3<f32> {
    let reflection = reflect(-view, surface_normal);
    let denom = dot(reflection, quad.normal);
    var hit = quad.center;
    if abs(denom) > 1e-4 {
        let t = dot(quad.center - surface_pos, quad.normal) / denom;
        if t > 0.0 {
            hit = surface_pos + reflection * t;
        }
    }
    let len_u = length(quad.axis_u);
    let len_v = length(quad.axis_v);
    let dir_u = quad.axis_u / max(len_u, 1e-5);
    let dir_v = quad.axis_v / max(len_v, 1e-5);
    let offset = hit - quad.center;
    let du = clamp(dot(offset, dir_u), -len_u, len_u);
    let dv = clamp(dot(offset, dir_v), -len_v, len_v);
    return quad.center + dir_u * du + dir_v * dv;
}

fn evaluate_one_area_light(
    light: AreaLight,
    surface_pos: vec3<f32>,
    surface_normal: vec3<f32>,
    view: vec3<f32>,
    albedo: vec3<f32>,
    roughness: f32,
    metallic: f32,
    f0: vec3<f32>,
) -> vec3<f32> {
    let quad = area_build_quad(light, surface_pos);
    let two_sided = light.two_sided != 0u;

    let to_center = quad.center - surface_pos;
    let distance = length(to_center);
    if light.range > 0.0 && distance > light.range + quad.extent {
        return vec3<f32>(0.0);
    }

    let facing = dot(normalize(surface_pos - quad.center), quad.normal);
    if !two_sided && facing <= 0.0 {
        return vec3<f32>(0.0);
    }

    let attenuation = area_range_attenuation(light.range, distance);
    var radiance = light.color.rgb * attenuation;

    if light.emissive_layer != AREA_EMISSIVE_LAYER_NONE {
        let local = surface_pos - quad.center;
        let len_u = max(length(quad.axis_u), 1e-5);
        let len_v = max(length(quad.axis_v), 1e-5);
        let uv = vec2<f32>(
            clamp(dot(local, quad.axis_u / len_u) / len_u, -1.0, 1.0) * 0.5 + 0.5,
            clamp(dot(local, quad.axis_v / len_v) / len_v, -1.0, 1.0) * 0.5 + 0.5,
        );
        radiance = radiance * sample_srgb_index(light.emissive_layer & 0xFFFFu, uv).rgb;
    }

    let diffuse_factor = area_polygon_form_factor(
        surface_normal,
        view,
        surface_pos,
        quad.corner_a,
        quad.corner_b,
        quad.corner_c,
        quad.corner_d,
        two_sided,
    );

    let k_diffuse = (vec3<f32>(1.0) - f0) * (1.0 - metallic);
    var result = k_diffuse * albedo * diffuse_factor * radiance;

    let representative = area_representative_point(quad, surface_pos, surface_normal, view);
    let light_dir = normalize(representative - surface_pos);
    let n_dot_l = dot(surface_normal, light_dir);
    if n_dot_l > 0.0 {
        let halfway = normalize(view + light_dir);
        let n_dot_v = max(dot(surface_normal, view), 1e-4);
        let n_dot_h = max(dot(surface_normal, halfway), 0.0);
        let v_dot_h = max(dot(view, halfway), 0.0);

        let rep_distance = max(length(representative - surface_pos), 1e-3);
        let alpha = roughness * roughness;
        let alpha_prime = clamp(alpha + quad.extent / (2.0 * rep_distance), 0.0, 1.0);
        let sphere_norm = (alpha * alpha) / max(alpha_prime * alpha_prime, 1e-5);

        let distribution = area_d_ggx(n_dot_h, roughness) * sphere_norm;
        let visibility = area_v_smith(n_dot_v, n_dot_l, roughness);
        let fresnel = area_fresnel(v_dot_h, f0);
        let specular = distribution * visibility * fresnel * 0.25 / max(n_dot_v * n_dot_l, 1e-4);
        result = result + specular * radiance * n_dot_l;
    }

    return result;
}
