#include "Material/StandardLit.h"
#include "Algorithm/BSDF.h"
#include "Algorithm/MicrofacetGGX.h"
#include "Algorithm/GGXMultiScatter.h"
#include "Lighting/LightEvaluation.h"

static float oren_nayar_eval(vec3s l, vec3s v, vec3s n, float roughness, float n_dot_l, float n_dot_v)
{
    // Full Qualitative Oren-Nayar
    float sigma2 = roughness * roughness;
    float A = 1.0f - (sigma2 / (2.0f * (sigma2 + 0.33f)));
    float B = 0.45f * sigma2 / (sigma2 + 0.09f);

    // Cosine of azimuth difference (phi)
    // Project V and L onto the tangent plane
    vec3s v_plane = glms_vec3_normalize(glms_vec3_sub(v, glms_vec3_scale(n, n_dot_v)));
    vec3s l_plane = glms_vec3_normalize(glms_vec3_sub(l, glms_vec3_scale(n, n_dot_l)));
    float cos_phi_diff = fmaxf(0.0f, glms_vec3_dot(v_plane, l_plane));

    // Sine and Tangent terms
    // alpha = max(theta_l, theta_v), beta = min(theta_l, theta_v)
    // We use sin/tan relationships: sin(x) = sqrt(1-cos^2(x))
    float sin_theta_l = sqrtf(fmaxf(0.0f, 1.0f - n_dot_l * n_dot_l));
    float sin_theta_v = sqrtf(fmaxf(0.0f, 1.0f - n_dot_v * n_dot_v));

    float sin_alpha, tan_beta;
    if (n_dot_v > n_dot_l)
    {
        sin_alpha = sin_theta_l;
        tan_beta = sin_theta_v / fmaxf(n_dot_v, 0.0001f);
    }
    else
    {
        sin_alpha = sin_theta_v;
        tan_beta = sin_theta_l / fmaxf(n_dot_l, 0.0001f);
    }

    return (A + B * cos_phi_diff * sin_alpha * tan_beta) * INV_PI;
}

static void get_surface_data(const shading_context_t* context, const standard_lit_properties_t* properties, standard_lit_surface_data_t* data_out)
{
    // Use the ray cone width (footprint) instead of simple distance for mip selection
    float distance = glms_vec3_distance(context->camera_position, context->position);
    vec3s view = context->camera_direction;

    // Fetch geometry data for LOD calculation
    const triangle_t* triangle = &context->triangles->buffer[context->triangle_id];
    texture_sample_context_t sample_context = 
    {
        .view_direction = view,
        .normal = context->normal,
        .edge1 = glms_vec3_sub(triangle->vertices[1].position, triangle->vertices[0].position),
        .edge2 = glms_vec3_sub(triangle->vertices[2].position, triangle->vertices[0].position),
        .uv1 = glms_vec2_sub(triangle->vertices[1].uv, triangle->vertices[0].uv),
        .uv2 = glms_vec2_sub(triangle->vertices[2].uv, triangle->vertices[0].uv),
        .ray_width = context->cone_width,
        .distance = distance
    };

    data_out->albedo = properties->albedo;
    const texture_t* albedo_texture = get_texture(context->textures, properties->albedo_texture);
    if (albedo_texture != NULL && albedo_texture->data != NULL)
    {
        data_out->albedo = glms_vec3_mul(data_out->albedo, glms_vec3(texture_sample(albedo_texture, &sample_context, context->uv)));
    }

    data_out->normal = context->normal;
    const texture_t* normal_texture = get_texture(context->textures, properties->normal_texture);
    if (normal_texture != NULL && normal_texture->data != NULL)
    {
        vec3s normal_sample = glms_vec3(texture_sample(normal_texture, &sample_context, context->uv));
        normal_sample = normal_unpack(normal_sample);
        data_out->normal = normal_ts_to_ws(normal_sample, context->normal, context->tangent);
    }

    data_out->diffuse_roughness = properties->diffuse_roughness;

    data_out->roughness = properties->roughness;
    const texture_t* roughness_texture = get_texture(context->textures, properties->roughness_texture);
    if (roughness_texture != NULL && roughness_texture->data != NULL)
    {
        data_out->roughness = data_out->roughness * texture_sample(roughness_texture, &sample_context, context->uv).x;
    }

    data_out->roughness = fmaxf(data_out->roughness, 0.001f);

    data_out->metallic = properties->metallic;
    const texture_t* metallic_texture = get_texture(context->textures, properties->metallic_texture);
    if (metallic_texture != NULL && metallic_texture->data != NULL)
    {
        data_out->metallic = data_out->metallic * texture_sample(metallic_texture, &sample_context, context->uv).x;
    }
}

static vec3s evaluate_bsdf_standard_lit(const shading_context_t* context, standard_lit_surface_data_t* surface_data, vec3s wi)
{
    vec3s n = surface_data->normal;
    vec3s v = glms_vec3_negate(context->wo);
    vec3s l = wi;

    float n_dot_l = fmaxf(glms_vec3_dot(n, l), 0.0f);
    float n_dot_v = fmaxf(glms_vec3_dot(n, v), 0.0f);
    if (n_dot_l <= 0.0f || n_dot_v <= 0.0f || glms_vec3_dot(context->normal, l) <= 0.0f)
    {
        return glms_vec3_zero();
    }
    
    vec3s h = glms_vec3_normalize(glms_vec3_add(v, l));
    float n_dot_h = fmaxf(glms_vec3_dot(n, h), 0.0f);
    float v_dot_h = fmaxf(glms_vec3_dot(v, h), 0.0f);

    // Fresnel
    vec3s f0 = glms_vec3_lerp(DIELECTRIC_REFLECTIVE, surface_data->albedo, surface_data->metallic);
    vec3s F = fresnel_schlick_vec3(f0, v_dot_h);

    // Specular (GGX)
    float D = ggx_distribution(n_dot_h, surface_data->roughness);
    float G = ggx_g_smith(n_dot_v, n_dot_l, surface_data->roughness);
    vec3s spec = glms_vec3_scale(glms_vec3_mul(F, (vec3s){D * G, D * G, D * G}), 1.0f / fmaxf(4.0f * n_dot_v * n_dot_l, 0.0001f));

    // Multi-scatter GGX (broad lobe)
    vec3s ms = ggx_multi_scatter_lambert(f0, n_dot_v, n_dot_l, surface_data->roughness);

    // Diffuse (Oren-Nayar)
    // Using (1 - F) here can make rough dielectrics look too dark because our specular is single-scatter GGX
    // (missing multi-scattering energy compensation). A stable approximation is (1 - F0_diel).
    float kd_scale = (1.0f - surface_data->metallic) * (1.0f - DIELECTRIC_REFLECTIVE_F0);
    vec3s kD = glms_vec3_scale(glms_vec3_one(), kd_scale);
    float on_val = oren_nayar_eval(l, v, n, surface_data->diffuse_roughness, n_dot_l, n_dot_v);
    vec3s diff = glms_vec3_scale(glms_vec3_mul(surface_data->albedo, kD), on_val);

    return glms_vec3_add(glms_vec3_add(diff, spec), ms);
}

static float sample_bsdf_pdf(const standard_lit_surface_data_t* surface_data, vec3s V, vec3s L)
{
    float n_dot_l = fmaxf(glms_vec3_dot(surface_data->normal, L), 0.0f);
    if (n_dot_l <= 0.0f)
    {
        return 0.0f;
    }

    // Lobe probabilities (single-scatter spec vs cosine)
    // We allocate some cosine probability for multi-scatter spec, especially for rough metals.
    vec3s f0 = glms_vec3_lerp(DIELECTRIC_F0, surface_data->albedo, surface_data->metallic);
    float n_dot_v = fmaxf(glms_vec3_dot(surface_data->normal, V), 0.0001f);
    vec3s F_est = fresnel_schlick_vec3(f0, n_dot_v);

    float spec_strength = luminance(F_est);
    float diff_strength = (1.0f - surface_data->metallic) * (1.0f - spec_strength);

    float Eo = ggx_ms_E(n_dot_v, surface_data->roughness);
    float w_ss = spec_strength * Eo;
    float w_cos = spec_strength * (1.0f - Eo) + diff_strength;

    float sum_w = w_ss + w_cos;
    if (sum_w < FLT_EPSILON)
    {
        return 0.0f;
    }
    w_ss /= sum_w;
    w_cos /= sum_w;

    // Specular PDF (GGX VNDF reflection)
    // GGX VNDF reflection pdf:
    // p_h(h)  = D(h) * G1(v) * (N·H)/(N·V)
    // p_w(wi) = p_h(h) / (4 * (V·H))
    vec3s H = glms_vec3_normalize(glms_vec3_add(L, V));
    float v_dot_h = glms_vec3_dot(V, H);
    if (v_dot_h <= 1e-6f)
    {
        return 0.0f;
    }
    float n_dot_h = fmaxf(glms_vec3_dot(surface_data->normal, H), 0.0f);
    float D = ggx_distribution(n_dot_h, surface_data->roughness);
    float G1v = ggx_g1(n_dot_v, surface_data->roughness);

    float pdf_h = (D * G1v * n_dot_h) / fmaxf(n_dot_v, 1e-6f);
    float pdf_spec = pdf_h / (4.0f * fmaxf(v_dot_h, 1e-6f));

    // Cosine PDF (used for diffuse + multi-scatter)
    float pdf_cos = n_dot_l * INV_PI;

    return w_ss * pdf_spec + w_cos * pdf_cos;
}

path_output standard_lit_render_loop(const standard_lit_properties_t* properties, const shading_context_t* context)
{
    standard_lit_surface_data_t surface_data; // Assuming you reuse your struct
    get_surface_data(context, properties, &surface_data);

    // Keep shading normal in the same hemisphere as the geometric normal to avoid invalid transport.
    if (glms_vec3_dot(surface_data.normal, context->normal) < 0.0f)
    {
        surface_data.normal = glms_vec3_negate(surface_data.normal);
    }

    path_output output = {0};
    vec3s V = glms_vec3_negate(context->wo);
    // Keep shading normal in the same hemisphere as the geometric normal.
    if (glms_vec3_dot(surface_data.normal, context->normal) < 0.0f)
    {
        surface_data.normal = glms_vec3_negate(surface_data.normal);
    }

    // If shading normal faces away from the view, fall back to geometric normal (prevents VNDF/pdf blow-ups).
    if (glms_vec3_dot(surface_data.normal, V) <= 0.0f)
    {
        surface_data.normal = context->normal;
    }

    float n_dot_v = fmaxf(glms_vec3_dot(surface_data.normal, V), 0.0001f);

    // Ensure LUT is ready (thread-safe, one-time).
    ggx_ms_init_lut_once();

    light_shading_context_t light_context = {
        .position = context->position,
        .normal = surface_data.normal,
        .geometric_normal = context->normal,
        .tangent = context->tangent,
        .uv = context->uv,
        .wo = context->wo,

        .bounce_depth = context->bounce_depth,

        .scene = context->scene,

        .bvh_tree = context->bvh_tree,
        .textures = context->textures,
    };

    uint32_t scramble = hash_position(context->position);

    // Running the light loop.
    // TODO: Implementing other light types.
    for (uint32_t i = 0; i < context->lights->directional_light_count; i++)
    {
        path_output light_output = evaluate_bsdf_directional(context->lights->directional_lights[i], &light_context, context->throughput, context->sample_index);
        if (light_output.state == PS_SUCCESS)
        {
            vec3s bsdf_dir_light = evaluate_bsdf_standard_lit(context, &surface_data, light_output.wi);
            float pdf_bsdf = sample_bsdf_pdf(&surface_data, V, light_output.wi);
            vec3s light_contribute = weight_nee_light(bsdf_dir_light, light_output.direct_lighting, pdf_bsdf, light_output.pdf);
            output.direct_lighting = glms_vec3_add(output.direct_lighting, light_contribute);
        }
    }

    // Sky
    path_output sky_output = evaluate_bsdf_sky(context->lights, &light_context, context->throughput, context->sample_index);
    if (sky_output.state == PS_SUCCESS)
    {
        vec3s bsdf_sky_light = evaluate_bsdf_standard_lit(context, &surface_data, sky_output.wi);
        float pdf_bsdf = sample_bsdf_pdf(&surface_data, V, sky_output.wi);
        vec3s sky_light = weight_nee_light(bsdf_sky_light, sky_output.direct_lighting, pdf_bsdf, sky_output.pdf);
        output.direct_lighting = glms_vec3_add(output.direct_lighting, sky_light);
    }

    // ----------------------------------------------------
    // Indirect Lighting (Sampling Next Ray)
    // ----------------------------------------------------

    // 1. Choose lobe: VNDF single-scatter spec vs cosine (diffuse + multi-scatter spec)
    vec3s f0 = glms_vec3_lerp(DIELECTRIC_F0, surface_data.albedo, surface_data.metallic);
    vec3s F_est = fresnel_schlick_vec3(f0, n_dot_v);

    float spec_strength = luminance(F_est);
    float diff_strength = (1.0f - surface_data.metallic) * (1.0f - spec_strength);

    float Eo = ggx_ms_E(n_dot_v, surface_data.roughness);
    float w_ss = spec_strength * Eo;
    float w_cos = spec_strength * (1.0f - Eo) + diff_strength;

    float sum_w = w_ss + w_cos;
    if (sum_w < FLT_EPSILON)
    {
        output.state = PS_TERMINATE;
        return output;
    }

    w_ss /= sum_w;
    w_cos /= sum_w;

    vec3s f_eval = glms_vec3_zero();
    float n_dot_l = 0.0f;
    float r_lobe = sobol_sample_scrambled(context->sample_index, sobol_get_dimension(context->bounce_depth, PRNG_BSDF), scramble);
    bool is_specular = (r_lobe < w_ss);

    if (is_specular)
    {
        // Sample GGX
        uint32_t d1 = sobol_get_dimension(context->bounce_depth, PRNG_BSDF_U);
        uint32_t d2 = sobol_get_dimension(context->bounce_depth, PRNG_BSDF_V);

        float u1 = sobol_sample_scrambled(context->sample_index, d1, scramble);
        float u2 = sobol_sample_scrambled(context->sample_index, d2, scramble);

        vec3s H = ggx_sample_vndf(surface_data.normal, V, surface_data.roughness, u1, u2);
        output.wi = glms_vec3_reflect(context->wo, H); // reflect(-V, H) -> V is wo inverted

        if (glms_vec3_dot(output.wi, surface_data.normal) <= 0.0f || glms_vec3_dot(output.wi, context->normal) <= 0.0f)
        {
            output.state = PS_TERMINATE;
            return output;
        }

        // Recalculate dots
        n_dot_l = fmaxf(glms_vec3_dot(surface_data.normal, output.wi), 0.0001f);
        vec3s H_new = glms_vec3_normalize(glms_vec3_add(output.wi, V));
        float n_dot_h = fmaxf(glms_vec3_dot(surface_data.normal, H_new), 0.0001f);
        float v_dot_h = fmaxf(glms_vec3_dot(V, H_new), 0.0001f);

        // Evaluate BSDF
        float D = ggx_distribution(n_dot_h, surface_data.roughness);
        float G1v = ggx_g1(n_dot_v, surface_data.roughness);
        float G = ggx_g_smith(n_dot_v, n_dot_l, surface_data.roughness);
        vec3s f0 = glms_vec3_lerp(DIELECTRIC_F0, surface_data.albedo, surface_data.metallic);
        vec3s F = fresnel_schlick_vec3(f0, v_dot_h);

        vec3s spec_f = glms_vec3_scale(glms_vec3_mul(F, (vec3s){D * G, D * G, D * G}),
                                       1.0f / fmaxf(4.0f * n_dot_v * n_dot_l, 1e-6f));

        f_eval = spec_f;

        // Propagate spread angle for ray cones
        // Heuristic: spread increases with roughness
        output.spread_angle = context->spread_angle + surface_data.roughness * QUARTER_PI;
    }
    else
    {
        // Sample Cosine hemisphere (Diffuse + Multi-scatter spec)
        // Note: We use cosine sampling for Oren-Nayar and the broad MS term.
        uint32_t d1 = sobol_get_dimension(context->bounce_depth, PRNG_BSDF_U);
        uint32_t d2 = sobol_get_dimension(context->bounce_depth, PRNG_BSDF_V);
        output.wi = random_cosine_direction(surface_data.normal, context->sample_index, d1, d2, scramble);

        n_dot_l = fmaxf(glms_vec3_dot(surface_data.normal, output.wi), 0.0001f);
        if (glms_vec3_dot(output.wi, context->normal) <= 0.0f)
        {
            output.state = PS_TERMINATE;
            return output;
        }

        float kd_scale = (1.0f - surface_data.metallic) * (1.0f - DIELECTRIC_REFLECTIVE_F0);
        vec3s kD = glms_vec3_scale(glms_vec3_one(), kd_scale);
        float on = oren_nayar_eval(output.wi, V, surface_data.normal, surface_data.diffuse_roughness, n_dot_l, n_dot_v);


        // Diffuse bounce significantly increases spread (effectively resets or becomes very wide)
        output.spread_angle = context->spread_angle + 0.5f;
        vec3s diff_f = glms_vec3_scale(glms_vec3_mul(surface_data.albedo, kD), on);

        // Multi-scatter GGX term (broad): sampled here with cosine.
        vec3s ms_f = ggx_multi_scatter_lambert(f0, n_dot_v, n_dot_l, surface_data.roughness);

        // Throughput multiplier: (f * NoL) / pdf
        vec3s f_sum = glms_vec3_add(diff_f, ms_f);
        f_eval = f_sum;
    }

    output.pdf = sample_bsdf_pdf(&surface_data, V, output.wi);
    if (output.pdf < 1e-12f)
    {
        output.state = PS_TERMINATE;
        return output;
    }

    // Throughput multiplier must be: (f * NoL) / pdf_total (mixture PDF)
    output.bsdf = glms_vec3_scale(f_eval, n_dot_l / output.pdf);
    output.state = PS_SUCCESS;
    return output;
}

void standard_lit_render_aov(const standard_lit_properties_t* properties, const shading_context_t* context, aov_output_t* aov_output)
{
    standard_lit_surface_data_t surface_data; // Assuming you reuse your struct
    get_surface_data(context, properties, &surface_data);
    // Keep shading normal in the same hemisphere as the geometric normal to avoid invalid transport.
    if (glms_vec3_dot(surface_data.normal, context->normal) < 0.0f)
    {
        surface_data.normal = glms_vec3_negate(surface_data.normal);
    }

    aov_output->albedo = glms_vec4(surface_data.albedo, 1.0f);
    vec3s n_ws = glms_vec3_normalize(surface_data.normal);
    aov_output->normal = (vec4s){n_ws.x * 0.5f + 0.5f, n_ws.y * 0.5f + 0.5f, n_ws.z * 0.5f + 0.5f, 1.0f};
}