//! A simple 3D scene with light shining over a cube sitting on a plane.

mod kneeboardplugin;
mod otdipcplugin;
mod vrcontrollerplugin;
mod vrplugin;

use bevy_cef::prelude::WebviewExtendStandardMaterial;
use vrplugin::VrPlugin;

use bevy::{
    asset::RenderAssetUsages,
    color::palettes::css,
    prelude::*,
    render::render_resource::{Extent3d, TextureDimension, TextureFormat},
};

use crate::otdipcplugin::OtdIpcPlugin;
use crate::{
    kneeboardplugin::KneeboardPlugin,
    otdipcplugin::{PenButtons, PenPosition, PenPressure},
    vrcontrollerplugin::VrControllersPlugin,
};

fn main() {
    App::new()
        .add_plugins((VrPlugin, MeshPickingPlugin))
        .add_plugins(VrControllersPlugin)
        .add_plugins(OtdIpcPlugin)
        .add_plugins(KneeboardPlugin)
        .insert_resource(ClearColor(Color::NONE))
        .insert_resource(LastPenPos(None))
        .insert_resource(PenColor(css::BLACK.into()))
        .insert_resource(PenSize(5.0))
        .add_systems(Startup, setup)
        .add_systems(Update, (color_changer, cursor, draw, plot).chain())
        .run();
}

#[derive(Resource)]
struct MyProcGenImage(Handle<Image>);

#[derive(Resource)]
struct PenSize(f32);

#[derive(Resource)]
struct PenColor(Color);

const IMAGE_WIDTH: i32 = (210.0 * 3.5) as i32;
const IMAGE_HEIGHT: i32 = (279.0 * 3.5) as i32;

fn setup(mut commands: Commands, mut images: ResMut<Assets<Image>>) {
    let image = Image::new_fill(
        Extent3d {
            width: IMAGE_WIDTH as u32,
            height: IMAGE_HEIGHT as u32,
            depth_or_array_layers: 1,
        },
        TextureDimension::D2,
        // Initialize it with a beige color
        &(Srgba::new(1.0, 1.0, 1.0, 0.0).to_u8_array()),
        // Use the same encoding as the color we set
        TextureFormat::Rgba8UnormSrgb,
        RenderAssetUsages::MAIN_WORLD | RenderAssetUsages::RENDER_WORLD,
    );

    let handle = images.add(image);
    commands.insert_resource(MyProcGenImage(handle));
    commands.insert_resource(CursorBuffer(
        [0; IMAGE_WIDTH as usize * IMAGE_HEIGHT as usize * 4],
    ));
    commands.insert_resource(DrawingBuffer(
        [0; IMAGE_WIDTH as usize * IMAGE_HEIGHT as usize * 4],
    ));
}

#[derive(Resource)]
struct CursorBuffer([u8; IMAGE_WIDTH as usize * IMAGE_HEIGHT as usize * 4]);

impl Drawing for CursorBuffer {
    fn set_color_at(&mut self, x: i32, y: i32, color: Color) {
        if !(0..=IMAGE_WIDTH).contains(&x) || !(0..=IMAGE_HEIGHT).contains(&y) {
            return;
        }

        let index = ((x + (IMAGE_WIDTH * y)) * 4) as usize;
        let srgba = Srgba::from(color);

        if index >= self.0.len() {
            return;
        }

        self.0[index] = (srgba.red * u8::MAX as f32) as u8;
        self.0[index + 1] = (srgba.green * u8::MAX as f32) as u8;
        self.0[index + 2] = (srgba.blue * u8::MAX as f32) as u8;
        self.0[index + 3] = (srgba.alpha * u8::MAX as f32) as u8;
    }

    fn clear(&mut self) {
        self.0.fill(0);
    }
}

#[derive(Resource)]
struct DrawingBuffer([u8; IMAGE_WIDTH as usize * IMAGE_HEIGHT as usize * 4]);

#[derive(Resource)]
pub struct LastPenPos(pub Option<(i32, i32)>);

impl Drawing for DrawingBuffer {
    fn set_color_at(&mut self, x: i32, y: i32, color: Color) {
        if !(0..=IMAGE_WIDTH).contains(&x) || !(0..=IMAGE_HEIGHT).contains(&y) {
            return;
        }

        let index = ((x + (IMAGE_WIDTH * y)) * 4) as usize;
        let srgba = Srgba::from(color);

        if index >= self.0.len() {
            return;
        }

        self.0[index] = (srgba.red * u8::MAX as f32) as u8;
        self.0[index + 1] = (srgba.green * u8::MAX as f32) as u8;
        self.0[index + 2] = (srgba.blue * u8::MAX as f32) as u8;
        self.0[index + 3] = (srgba.alpha * u8::MAX as f32) as u8;
    }

    fn clear(&mut self) {
        self.0.fill(0);
    }
}

fn color_changer(mut pen_color: ResMut<PenColor>, mut pen_buttons: MessageReader<PenButtons>) {
    let Some(buttons) = pen_buttons.read().next() else {
        return;
    };

    pen_color.0 = css::BLACK.into();

    if (buttons.state & 2) == 2 {
        let mut t = LinearRgba::from(pen_color.0);
        t.alpha = 0.0;
        pen_color.0 = t.into();
    }

    if (buttons.state & 4) == 4 {
        let mut t = LinearRgba::from(css::DARK_BLUE);
        t.alpha = 0.5;
        pen_color.0 = t.into();
    }
}

fn cursor(
    pen_color: Res<PenColor>,
    mut buffer: ResMut<CursorBuffer>,
    mut pen_size: ResMut<PenSize>,
    mut pen_buttons: MessageReader<PenButtons>,
    mut pen_position: MessageReader<PenPosition>,
    mut pen_pressure: MessageReader<PenPressure>,
) {
    let mut size: f32 = 1.0;
    let mut offset = 20.0;

    let c = pen_color.0;

    if let Some(buttons) = pen_buttons.read().next()
        && (buttons.state & 2 == 2)
    {
        offset *= 4.;
    };

    for penpres in pen_pressure.read() {
        // this needs log scaling

        let _in = penpres.pressure;
        let _out = penpres.pressure.powi(4);

        size *= 1. + (_out * offset);
    }

    pen_size.0 = size.clamp(1.0, 200.0);
    let s = pen_size.0 as i32;
    let b = css::BLACK.into();

    let cs = 10;
    let cd = 3;

    for penpos in pen_position.read() {
        let x = (penpos.position.x * (IMAGE_WIDTH as f32)) as i32;
        let y = (penpos.position.y * (IMAGE_HEIGHT as f32)) as i32;

        buffer.clear();
        draw_line(&mut buffer, x - (cs + cd) - s, y, x - (cd) - s, y, 4.0, b);
        draw_line(&mut buffer, x - (cs + cd) - s, y, x - (cd) - s, y, 1.0, c);
        draw_line(&mut buffer, x + (cd) + s, y, x + (cs + cd) + s, y, 4.0, b);
        draw_line(&mut buffer, x + (cd) + s, y, x + (cs + cd) + s, y, 1.0, c);

        draw_line(&mut buffer, x, y - (cs + cd) - s, x, y - (cd) - s, 4.0, b);
        draw_line(&mut buffer, x, y - (cs + cd) - s, x, y - (cd) - s, 1.0, c);
        draw_line(&mut buffer, x, y + (cs + cd) + s, x, y + (cd) + s, 4.0, b);
        draw_line(&mut buffer, x, y + (cs + cd) + s, x, y + (cd) + s, 1.0, c);

        draw_filled_circle(&mut buffer, x, y, (pen_size.0 / 2.) as i32, c)
    }
}

fn draw(
    pen_size: Res<PenSize>,
    pen_color: Res<PenColor>,
    mut lastloc: ResMut<LastPenPos>,
    mut buffer: ResMut<DrawingBuffer>,
    mut pen_buttons: MessageReader<PenButtons>,
    mut pen_position: MessageReader<PenPosition>,
) {
    let Some(buttons) = pen_buttons.read().next() else {
        return;
    };

    for penpos in pen_position.read() {
        let x = (penpos.position.x * (IMAGE_WIDTH as f32)) as i32;
        let y = (penpos.position.y * (IMAGE_HEIGHT as f32)) as i32;

        let Some(ll) = lastloc.0 else {
            lastloc.0 = Some((x, y));
            return;
        };

        if buttons.state & 1 == 1 {
            draw_line(&mut buffer, ll.0, ll.1, x, y, pen_size.0, pen_color.0);
        }

        lastloc.0 = Some((x, y));
    }
}

trait Drawing {
    fn set_color_at(&mut self, px: i32, py: i32, color: Color);
    fn clear(&mut self);
}

fn draw_line<T: Drawing + Resource>(
    buffer: &mut ResMut<T>,
    x1: i32,
    y1: i32,
    x2: i32,
    y2: i32,
    thickness: f32,
    color: Color,
) {
    let (mut x0, mut y0) = (x1, y1);
    let (x1, y1) = (x2, y2);

    let dx = (x1 - x0).abs();
    let dy = -(y1 - y0).abs();
    let sx = if x0 < x1 { 1 } else { -1 };
    let sy = if y0 < y1 { 1 } else { -1 };
    let mut err = dx + dy;

    let half_t = (thickness / 2.0).ceil() as i32;

    loop {
        draw_thick_point(buffer, x0, y0, half_t, color);

        if x0 == x1 && y0 == y1 {
            break;
        }

        let e2 = 2 * err;
        if e2 >= dy {
            err += dy;
            x0 += sx;
        }
        if e2 <= dx {
            err += dx;
            y0 += sy;
        }
    }
}

// Draw a small square brush (fast, branch-free inner loop)
fn draw_thick_point<T: Drawing + Resource>(
    buffer: &mut ResMut<T>,
    cx: i32,
    cy: i32,
    radius: i32,
    color: Color,
) {
    draw_filled_circle(buffer, cx, cy, radius, color);
}

fn draw_filled_circle<T: Drawing + Resource>(
    buffer: &mut ResMut<T>,
    cx: i32,
    cy: i32,
    radius: i32,
    color: Color,
) {
    let r2 = radius * radius;

    for dy in -radius..=radius {
        let y = cy + dy;
        if y < 0 {
            continue;
        }

        let dx_limit = ((r2 - dy * dy) as f32).sqrt() as i32;

        let start_x = cx - dx_limit;
        let end_x = cx + dx_limit;

        for x in start_x..=end_x {
            if x >= 0 {
                buffer.set_color_at(x, y, color);
            }
        }
    }
}

fn draw_circle<T: Drawing + Resource>(
    buffer: &mut ResMut<T>,
    cx: i32,
    cy: i32,
    radius: i32,
    thickness: f32,
    color: Color,
) {
    if radius <= 0 {
        return;
    }

    let half_t = (thickness / 2.0).max(0.5);
    let outer_r = radius as f32 + half_t;
    let inner_r = (radius as f32 - half_t).max(0.0);

    let outer_r2 = (outer_r * outer_r) as i32;
    let inner_r2 = (inner_r * inner_r) as i32;

    let max_r = outer_r.ceil() as i32;

    for dy in -max_r..=max_r {
        let y = cy + dy;
        if y < 0 {
            continue;
        }

        let dy2 = dy * dy;
        if dy2 > outer_r2 {
            continue;
        }

        let outer_dx = ((outer_r2 - dy2) as f32).sqrt() as i32;
        let inner_dx = if dy2 < inner_r2 {
            ((inner_r2 - dy2) as f32).sqrt() as i32
        } else {
            -1
        };

        let left_outer = cx - outer_dx;
        let right_outer = cx + outer_dx;

        if inner_dx >= 0 {
            let left_inner = cx - inner_dx;
            let right_inner = cx + inner_dx;

            // left segment
            for x in left_outer..left_inner {
                if x >= 0 {
                    buffer.set_color_at(x, y, color);
                }
            }

            // right segment
            for x in (right_inner + 1)..=right_outer {
                if x >= 0 {
                    buffer.set_color_at(x, y, color);
                }
            }
        } else {
            // fully filled span (very thin or small radius)
            for x in left_outer..=right_outer {
                if x >= 0 {
                    buffer.set_color_at(x, y, color);
                }
            }
        }
    }
}

fn plot(
    cursor: Res<CursorBuffer>,
    buffer: Res<DrawingBuffer>,
    mut images: ResMut<Assets<Image>>,
    mut materials: ResMut<Assets<WebviewExtendStandardMaterial>>,
    webviews: Query<&MeshMaterial3d<WebviewExtendStandardMaterial>>,
    my_handle: Res<MyProcGenImage>,
) {
    let image = images.get_mut(&my_handle.0).expect("Image not found");
    image.clear(&[0, 0, 0, 0]);

    for (i, c) in buffer.0.chunks(4).enumerate() {
        let x = i as i32 % IMAGE_WIDTH;
        let y = i as i32 / IMAGE_WIDTH;

        let red: u8 = c[0];
        let green = c[1];
        let blue = c[2];
        let alpha = c[3];

        if alpha > 0 {
            let color = Color::srgba_u8(red, green, blue, alpha);
            image.set_color_at(x as u32, y as u32, color).unwrap();
        }
    }

    for (i, c) in cursor.0.chunks(4).enumerate() {
        let x = i as i32 % IMAGE_WIDTH;
        let y = i as i32 / IMAGE_WIDTH;

        let red: u8 = c[0];
        let green = c[1];
        let blue = c[2];
        let alpha = c[3];

        if alpha > 0 {
            let color = Color::srgba_u8(red, green, blue, alpha);
            image.set_color_at(x as u32, y as u32, color).unwrap();
        }
    }

    // Poke redraw
    for handle in webviews {
        if let Some(material) = materials.get_mut(handle.id()) {
            material.extension.overlay = Some(my_handle.0.clone());
        }
    }
}

// #[allow(clippy::type_complexity)]
// fn head_pointer(
//     browsers: NonSend<Browsers>,
//     webviews: Query<(Entity, &WebviewSize), With<WebviewSource>>,
//     headset: Query<&GlobalTransform, With<XrCamera>>,
//     kneeboard: Query<(&GlobalTransform, Option<&LookedAt>), With<Kneeboard>>,
//     aabb: MeshAabb,
// ) {
//     let Ok((webview, size)) = webviews.single() else {
//         return;
//     };

//     let tex_size = size.0;

//     if let Some(gt) = headset.into_iter().next() {
//         let Ok((plane_tf, looked_at)) = kneeboard.single() else {
//             return;
//         };

//         if looked_at.is_some() {
//             // this is inverted for some reason wtf
//             return;
//         }

//         let (min, max) = aabb.calculate_local(webview);
//         let plane_size = Vec2::new(max.x - min.x, max.y - min.y);

//         let ray = Ray3d::new(gt.translation(), gt.forward());

//         let n = plane_tf.forward().as_vec3();
//         let Some(t) = ray.intersect_plane(
//             plane_tf.translation(),
//             InfinitePlane3d::new(plane_tf.forward()),
//         ) else {
//             return;
//         };
//         let hit_world = ray.origin + ray.direction * t;
//         let local_hit = plane_tf.affine().inverse().transform_point(hit_world);
//         let local_normal = plane_tf.affine().inverse().transform_vector3(n).normalize();
//         let abs_normal = local_normal.abs();
//         let (u_coord, v_coord) = if abs_normal.z > abs_normal.x && abs_normal.z > abs_normal.y {
//             (local_hit.x, local_hit.y)
//         } else if abs_normal.y > abs_normal.x {
//             (local_hit.x, local_hit.z)
//         } else {
//             (local_hit.y, local_hit.z)
//         };

//         let w = plane_size.x;
//         let h = plane_size.y;
//         let u = (u_coord + w * 0.5) / w;
//         let v = (v_coord + h * 0.5) / h;
//         if !(0.0..=1.0).contains(&u) || !(0.0..=1.0).contains(&v) {
//             // outside plane bounds
//             return;
//         }
//         let px = (1.0 - u) * tex_size.x;
//         let py = (1.0 - v) * tex_size.y;

//         let pos = Vec2::new(px, py);

//         browsers.send_mouse_move(&webview, &[], pos, false);
//     }
// }