finite_light_bevy 0.1.2

use std::f32::consts::{PI, TAU};

use bevy::{
    asset::RenderAssetUsages,
    image::ImageSampler,
    prelude::*,
    render::render_resource::{Extent3d, TextureDimension, TextureFormat},
};

use crate::HideWorldLine;

/// A skybox that deforms under relativistic aberration.
///
/// Provide six cubemap face images as positive/negative pairs per axis. The
/// plugin converts them to an equirectangular panorama, creates a UV sphere
/// mesh, and feeds it through the standard relativistic compute pipeline.
#[derive(Component)]
pub struct RelativisticSkybox {
    /// X-axis face images: [positive, negative].
    pub x: [Handle<Image>; 2],
    /// Y-axis face images: [positive, negative].
    pub y: [Handle<Image>; 2],
    /// Z-axis face images: [positive, negative].
    pub z: [Handle<Image>; 2],

    /// Sphere radius. Must be within the camera's far clipping plane and
    /// larger than any scene geometry.
    pub radius: f32,
}

/// Convert cubemap faces and spawn the skybox mesh.
pub(crate) fn assemble_skybox(
    mut commands: Commands,
    mut images: ResMut<Assets<Image>>,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
    query: Query<(Entity, &RelativisticSkybox), Without<Mesh3d>>,
) {
    for (entity, skybox) in &query {
        // Wait until all six faces are loaded.
        let handles = [
            &skybox.x[0],
            &skybox.x[1],
            &skybox.y[0],
            &skybox.y[1],
            &skybox.z[0],
            &skybox.z[1],
        ];
        let face_images: Vec<&Image> = handles.iter().filter_map(|h| images.get(*h)).collect();

        if face_images.len() != 6 {
            continue;
        }

        let face_size = face_images[0].width();
        assert_eq!(
            face_size,
            face_images[0].height(),
            "Skybox faces must be square"
        );

        let faces: [&Image; 6] = face_images.try_into().unwrap();
        let equirect = cubemap_to_equirectangular(faces, face_size);

        commands.entity(entity).insert((
            Mesh3d(meshes.add(Sphere::new(skybox.radius).mesh().uv(128, 64))),
            MeshMaterial3d(materials.add(StandardMaterial {
                base_color_texture: Some(images.add(equirect)),
                unlit: true,
                double_sided: true,
                cull_mode: None,
                ..default()
            })),
            // Apply a transformation to the skybox to orient it correctly.
            Transform::from_rotation(
                Quat::from_rotation_y(core::f32::consts::PI)
                    * Quat::from_rotation_x(core::f32::consts::FRAC_PI_2),
            ),
            HideWorldLine,
        ));
        commands.entity(entity).remove::<RelativisticSkybox>();
    }
}

/// Re-project six cubemap faces into a single equirectangular panorama.
///
/// Output size is `face_size * 4` wide by `face_size * 2` tall. Each output
/// pixel is mapped to a 3D direction via spherical coordinates, then sampled
/// from the appropriate cubemap face.
fn cubemap_to_equirectangular(faces: [&Image; 6], face_size: u32) -> Image {
    let width = face_size * 4;
    let height = face_size * 2;
    let mut data = vec![0u8; (width * height * 4) as usize];
    let max_coord = (face_size - 1) as f32;

    for py in 0..height {
        for px in 0..width {
            let u = (px as f32 + 0.5) / width as f32;
            let v = (py as f32 + 0.5) / height as f32;

            // Spherical coordinates matching Bevy's UV sphere parametrization.
            let theta = v * PI;
            let phi = u * TAU;
            let (sin_theta, cos_theta) = theta.sin_cos();
            let (sin_phi, cos_phi) = phi.sin_cos();
            let dir = Vec3::new(sin_theta * cos_phi, -cos_theta, sin_theta * sin_phi);

            let (face_index, face_u, face_v) = cubemap_face_uv(dir);

            let face_data = rgba8_bytes(faces[face_index]);
            let pixel =
                bilinear_sample(face_data, face_size, face_u * max_coord, face_v * max_coord);
            let dst = ((py * width + px) * 4) as usize;
            data[dst..dst + 4].copy_from_slice(&pixel);
        }
    }

    let mut image = Image::new(
        Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        },
        TextureDimension::D2,
        data,
        TextureFormat::Rgba8UnormSrgb,
        RenderAssetUsages::RENDER_WORLD,
    );
    image.sampler = ImageSampler::linear();
    image
}

/// Map a direction to a cubemap face index (WebGPU order) and face-local UV.
fn cubemap_face_uv(dir: Vec3) -> (usize, f32, f32) {
    let abs = dir.abs();

    // Select face by dominant axis. `sc` and `tc` are the signed components
    // that map to U and V within the face; `ma` is the major-axis magnitude.
    let (face, sc, tc, ma) = if abs.x >= abs.y && abs.x >= abs.z {
        if dir.x > 0. {
            (0, -dir.z, -dir.y, abs.x)
        } else {
            (1, dir.z, -dir.y, abs.x)
        }
    } else if abs.y >= abs.z {
        if dir.y > 0. {
            (2, dir.x, dir.z, abs.y)
        } else {
            (3, dir.x, -dir.z, abs.y)
        }
    } else if dir.z > 0. {
        (4, dir.x, -dir.y, abs.z)
    } else {
        (5, -dir.x, -dir.y, abs.z)
    };

    (face, (sc / ma + 1.) / 2., (tc / ma + 1.) / 2.)
}

/// Bilinear sample of an RGBA8 face at fractional pixel coordinates.
fn bilinear_sample(data: &[u8], size: u32, fx: f32, fy: f32) -> [u8; 4] {
    let x0 = (fx.floor() as u32).min(size - 1);
    let y0 = (fy.floor() as u32).min(size - 1);
    let x1 = (x0 + 1).min(size - 1);
    let y1 = (y0 + 1).min(size - 1);
    let tx = fx - fx.floor();
    let ty = fy - fy.floor();

    let fetch = |x: u32, y: u32| -> [f32; 4] {
        let i = ((y * size + x) * 4) as usize;
        [
            data[i] as f32,
            data[i + 1] as f32,
            data[i + 2] as f32,
            data[i + 3] as f32,
        ]
    };

    let p00 = fetch(x0, y0);
    let p10 = fetch(x1, y0);
    let p01 = fetch(x0, y1);
    let p11 = fetch(x1, y1);

    let mut result = [0u8; 4];
    for c in 0..4 {
        let top = p00[c] * (1. - tx) + p10[c] * tx;
        let bottom = p01[c] * (1. - tx) + p11[c] * tx;
        result[c] = (top * (1. - ty) + bottom * ty).round() as u8;
    }
    result
}

/// Return the image's pixel data, asserting RGBA8 format.
fn rgba8_bytes(image: &Image) -> &[u8] {
    match image.texture_descriptor.format {
        TextureFormat::Rgba8UnormSrgb | TextureFormat::Rgba8Unorm => {}
        format => panic!("Unexpected skybox face format {format:?}"),
    }
    image.data.as_deref().unwrap_or(&[])
}