#include "pxl8_shader.h"
#include "pxl8_shader_builtins.h"

u8 pxl8_shader_lit(pxl8_shader_ctx* ctx) {
    u8 tex_idx = 0;
    if (ctx->bindings && ctx->bindings->texture) {
        tex_idx = pxl8_sample_indexed(ctx, ctx->v_uv);
        if (tex_idx == 0) return 0;
    } else {
        if (ctx->uniforms && ctx->uniforms->dither) {
            tex_idx = pxl8_gfx_dither(ctx->v_color, (u32)ctx->x, (u32)ctx->y);
        } else {
            f32 clamped = pxl8_clamp(ctx->v_color, 0.0f, 255.0f);
            tex_idx = (u8)(clamped);
        }
    }

    f32 light = ctx->v_light;

    const pxl8_shader_uniforms* u = ctx->uniforms;
    if (u) {
        f32 ambient = (f32)u->ambient / 255.0f;
        if (ambient > light) light = ambient;

        if (u->celestial_intensity > 0.0f) {
            f32 dx = u->celestial_dir.x;
            f32 dy = u->celestial_dir.y;
            f32 dz = u->celestial_dir.z;
            f32 len = pxl8_sqrt(dx * dx + dy * dy + dz * dz);
            if (len > 0.0001f) {
                dx /= len;
                dy /= len;
                dz /= len;
                f32 ndotl = -(ctx->v_normal.x * dx + ctx->v_normal.y * dy + ctx->v_normal.z * dz);
                if (ndotl > 0.0f) {
                    light += ndotl * u->celestial_intensity;
                }
            }
        }

        f32 dyn_strength = 0.0f;
        f32 dyn_r = 0.0f;
        f32 dyn_g = 0.0f;
        f32 dyn_b = 0.0f;

        for (u32 i = 0; i < u->lights_count; i++) {
            const pxl8_light* l = &u->lights[i];
            f32 lx = l->position.x - ctx->v_world.x;
            f32 ly = l->position.y - ctx->v_world.y;
            f32 lz = l->position.z - ctx->v_world.z;
            f32 dist_sq = lx * lx + ly * ly + lz * lz;
            if (dist_sq >= l->radius_sq) continue;

            f32 dist = pxl8_sqrt(dist_sq);
            f32 inv_dist = dist > 0.0001f ? (1.0f / dist) : 0.0f;
            f32 nx = lx * inv_dist;
            f32 ny = ly * inv_dist;
            f32 nz = lz * inv_dist;

            f32 ndotl = ctx->v_normal.x * nx + ctx->v_normal.y * ny + ctx->v_normal.z * nz;
            if (ndotl <= 0.0f) continue;

            f32 falloff = 1.0f - dist_sq * l->inv_radius_sq;
            if (falloff <= 0.0f) continue;

            f32 strength = ((f32)l->intensity / 255.0f) * falloff * ndotl;
            if (strength <= 0.0f) continue;

            dyn_strength += strength;
            dyn_r += strength * (f32)l->r;
            dyn_g += strength * (f32)l->g;
            dyn_b += strength * (f32)l->b;
        }

        if (dyn_strength > 0.0f) {
            f32 inv = 1.0f / dyn_strength;
            u8 r = (u8)pxl8_clamp(dyn_r * inv, 0.0f, 255.0f);
            u8 g = (u8)pxl8_clamp(dyn_g * inv, 0.0f, 255.0f);
            u8 b = (u8)pxl8_clamp(dyn_b * inv, 0.0f, 255.0f);
            u8 a = (u8)pxl8_clamp(dyn_strength * 255.0f, 0.0f, 255.0f);
            ctx->out_light_color = (u32)r | ((u32)g << 8) | ((u32)b << 16) | ((u32)a << 24);
            light += dyn_strength;
        }
    }

    if (light > 1.0f) light = 1.0f;
    if (light < 0.0f) light = 0.0f;

    f32 light_f = light * 255.0f;
    u8 light_u8 = (u8)light_f;
    if (ctx->uniforms && ctx->uniforms->dither) {
        light_u8 = pxl8_gfx_dither(light_f, (u32)ctx->x, (u32)ctx->y);
    }

    u8 shaded = pxl8_colormap_lookup(ctx, tex_idx, light_u8);

    if (ctx->uniforms && ctx->uniforms->emissive) {
        u32 rgb = 0x00FFFFFF;
        if (ctx->bindings && ctx->bindings->palette) {
            rgb = ctx->bindings->palette[tex_idx] & 0x00FFFFFF;
        }
        pxl8_set_light_tint(ctx, rgb, 1.0f);
    }

    return shaded;
}