use cyborg::camera::{Camera, Flycam};
use cyborg::pass::{mesh, RenderPassBox};
use cyborg::shader;
use cyborg::storage::mesh::*;
use cyborg::viewport::*;
use cyborg::Renderer;
use legion::*;
use rand::prelude::*;
use std::sync::Arc;
use winit::dpi::PhysicalSize;
use winit::{
    event::*,
    event_loop::{ControlFlow, EventLoop},
    window::WindowBuilder,
};

struct DeltaTime {
    pub dt: f32,
    last_update: std::time::Instant,
}

impl Default for DeltaTime {
    fn default() -> Self {
        DeltaTime {
            dt: 0.0,
            last_update: std::time::Instant::now(),
        }
    }
}

struct SpawnDistributions {
    theta: rand::distributions::Uniform<f32>,
    radius: rand::distributions::Uniform<f32>,
    up: rand::distributions::Uniform<f32>,
}

impl Default for SpawnDistributions {
    fn default() -> Self {
        Self {
            theta: rand::distributions::Uniform::new(0.0, std::f32::consts::TAU),
            radius: rand::distributions::Uniform::new(3.0, 4.0),
            up: rand::distributions::Uniform::new(8.0, 12.0),
        }
    }
}

struct Particle {
    position: glam::Vec3A,
    mass: f32,
}

struct Velocity {
    linear: glam::Vec3A,
}

#[system]
fn update_dt(#[resource] dt: &mut DeltaTime) {
    dt.dt = dt.last_update.elapsed().as_secs_f32();
    dt.last_update = std::time::Instant::now();
}

#[system]
fn update_app(#[resource] app: &mut Application) {
    app.update();
}

#[system(for_each)]
fn apply_gravity(#[resource] dt: &DeltaTime, velocity: &mut Velocity) {
    const GRAVITY: f32 = 6.0;
    velocity.linear.y -= GRAVITY * dt.dt;
}

#[system(for_each)]
fn move_particles(#[resource] dt: &DeltaTime, particle: &mut Particle, velocity: &Velocity) {
    particle.position += velocity.linear * (dt.dt / particle.mass);
    // TODO angular velocity
}

#[system(for_each)]
fn respawn_particles(
    #[resource] distributions: &SpawnDistributions,
    particle: &mut Particle,
    velocity: &mut Velocity,
) {
    if particle.position.y < -10.0 {
        particle.position = glam::Vec3A::ZERO;
        let rng = &mut rand::thread_rng();
        velocity.linear.y = distributions.up.sample(rng);

        let theta = distributions.theta.sample(rng);
        let radius = distributions.radius.sample(rng);
        velocity.linear.x = theta.sin() * radius;
        velocity.linear.z = theta.cos() * radius;
    }
}

#[system(for_each)]
fn update_transforms(transform: &mut cyborg::scene::Transform, particle: &Particle) {
    transform.transform = glam::Mat4::from_translation(particle.position.into());
}

#[system(for_each)]
fn update_debug(draw_list: &mut cyborg::scene::DebugDrawList, velocity: &Velocity) {
    draw_list.clear();
    draw_list.indices.extend_from_slice(&[0, 1]);

    let color = [1.0, 1.0, 1.0];

    draw_list.vertices.push(cyborg::scene::DebugVertex {
        position: [0.0, 0.0, 0.0],
        color,
    });

    draw_list.vertices.push(cyborg::scene::DebugVertex {
        position: velocity.linear.to_array(),
        color,
    });
}

fn build_update_schedule() -> Schedule {
    let mut builder = Schedule::builder();
    builder.add_system(update_dt_system());
    builder.add_system(update_app_system());
    builder.add_system(apply_gravity_system());
    builder.add_system(move_particles_system());
    builder.add_system(respawn_particles_system());
    builder.add_system(update_transforms_system());
    builder.add_system(update_debug_system());
    builder.build()
}

fn make_terrain(attributes: &mesh::Attributes) -> MeshBuffer {
    const MOUNTAINS_RADIUS: f32 = 25.0;
    const MOUNTAINS_WIDTH: f32 = 10.0;
    const MOUNTAINS_HEIGHT: f32 = 5.0;
    const MOUNTAINS_FREQUENCY: f32 = 1.0;

    use noise::NoiseFn;
    let noise = noise::OpenSimplex::new();

    let sample = |xy: glam::Vec2| -> f32 {
        let mountains_dist = (xy.length() - MOUNTAINS_RADIUS).abs() / MOUNTAINS_WIDTH;
        let mountains_scale = 1.0 - mountains_dist;
        if mountains_scale > 0.0 {
            let noise_input = (xy * MOUNTAINS_FREQUENCY).as_dvec2();
            let sampled = noise.get(noise_input.to_array()) as f32;
            ((sampled + 0.5) / 0.5) * MOUNTAINS_HEIGHT * mountains_scale
        } else {
            0.0
        }
    };

    const GRID_FREQUENCY: f32 = 0.2;
    const GRID_RADIUS: isize = ((MOUNTAINS_RADIUS + MOUNTAINS_WIDTH) / GRID_FREQUENCY) as isize;
    const GRID_WIDTH: isize = GRID_RADIUS * 2 + 1;

    let tile_num = (GRID_RADIUS * GRID_RADIUS) as usize;
    let mut vertices = Vec::with_capacity(tile_num);
    let mut indices = Vec::with_capacity(tile_num * 6);

    for x in -GRID_RADIUS..=GRID_RADIUS {
        for z in -GRID_RADIUS..=GRID_RADIUS {
            let x = x as f32 * GRID_FREQUENCY;
            let z = z as f32 * GRID_FREQUENCY;
            let y = sample(glam::Vec2::new(x, z));

            vertices.push(mesh::Vertex {
                position: [x, y, z],
                tan_frame: 0,
            });
        }
    }

    let mut cursor: mesh::Index = 0;
    for _x in -GRID_RADIUS..GRID_RADIUS {
        for _z in -GRID_RADIUS..GRID_RADIUS {
            let next_line = cursor + GRID_WIDTH as mesh::Index;
            indices.extend_from_slice(&[
                cursor,
                next_line,
                cursor + 1,
                next_line,
                cursor + 1,
                next_line + 1,
            ]);
            cursor += 1;
        }
        cursor += 1;
    }

    let vertices = AttrBuffer {
        id: attributes.vertex,
        count: vertices.len(),
        data: bytemuck::cast_slice(&vertices).to_vec(),
    };

    let indices = AttrBuffer {
        id: attributes.index,
        count: indices.len(),
        data: bytemuck::cast_slice(&indices).to_vec(),
    };

    let mut mesh = MeshBuffer::default();
    mesh.attributes.push(vertices);
    mesh.attributes.push(indices);
    mesh
}

struct Application {
    window: winit::window::Window,
    viewport: WinitViewport,
    flycam: Flycam,
    is_grabbed: bool,
}

impl Application {
    pub fn update(&mut self) {
        self.flycam.update();
        self.window.request_redraw();
    }

    pub fn on_mouse_motion(&mut self, x: f64, y: f64) {
        if self.is_grabbed {
            self.flycam.process_mouse(x, y);
        }
    }

    pub fn on_mouse_event(&mut self, button: MouseButton, state: ElementState) {
        if button == MouseButton::Left && state == ElementState::Pressed && !self.is_grabbed {
            self.window.set_cursor_grab(true).unwrap();
            self.window.set_cursor_visible(false);
            self.is_grabbed = true;
        }
    }

    pub fn on_key_event(&mut self, key: VirtualKeyCode, state: ElementState) {
        if key == VirtualKeyCode::Escape && state == ElementState::Pressed {
            if self.is_grabbed {
                self.window.set_cursor_grab(false).unwrap();
                self.window.set_cursor_visible(true);
                self.is_grabbed = false;
            }
        } else {
            self.flycam.process_keyboard(key, state);
        }
    }

    pub fn on_resize(&mut self, new_size: PhysicalSize<u32>) {
        self.viewport.resize(new_size);
        self.flycam.resize(new_size.width, new_size.height);
    }

    pub fn on_lost_focus(&mut self) {
        self.flycam.defocus();
    }

    pub fn on_surface_lost(&mut self) {
        self.viewport.resize(self.viewport.size);
    }
}

impl cyborg::legion::RenderCallbacks for Application {
    fn get_viewports(&mut self) -> Vec<(&dyn Viewport, Camera)> {
        vec![(&self.viewport as &dyn Viewport, self.flycam.get_camera())]
    }

    fn present(&mut self) {
        self.viewport.present();
    }
}

fn main() {
    let event_loop = EventLoop::new();
    let window = WindowBuilder::new().build(&event_loop).unwrap();
    let viewport = pollster::block_on(WinitViewport::from_window(&window));

    let mut world = World::new(Default::default());
    let mut resources = Resources::default();

    resources.insert::<DeltaTime>(Default::default());
    resources.insert::<SpawnDistributions>(Default::default());
    let mut update_schedule = build_update_schedule();

    let renderer = Renderer::new(viewport.device.clone(), viewport.queue.clone());
    resources.insert(renderer);

    resources.insert(viewport.get_info().to_owned());

    let flycam = cyborg::camera::Flycam::new(0.002, 10.0, 0.25);

    let shader_store = Arc::new(shader::ShaderStore::new(viewport.device.clone()));
    let shaders_dir = std::env::current_dir().unwrap();
    let shaders_dir = shaders_dir.join("shaders/");
    let shader_watcher = shader::ShaderWatcher::new(shader_store.to_owned(), shaders_dir).unwrap();

    let mesh_forward = shader_watcher.add_file("mesh_forward.wgsl").unwrap();
    let mesh_skinning = shader_watcher.add_file("mesh_skinning.wgsl").unwrap();

    let mesh_shaders = mesh::ShaderInfo {
        store: shader_store.clone(),
        forward: mesh_forward,
        skinning: mesh_skinning,
    };

    resources.insert(mesh_shaders);

    let application = Application {
        window,
        viewport,
        flycam,
        is_grabbed: false,
    };

    let mut render_schedule = Schedule::builder();
    cyborg::legion::build_renderer(application, &mut resources, &mut render_schedule);
    let mut render_schedule = render_schedule.build();

    let example_vertices = vec![
        mesh::Vertex {
            position: [-0.5, 0.0, 0.5],
            tan_frame: 0,
        },
        mesh::Vertex {
            position: [0.5, 0.0, 0.5],
            tan_frame: 0,
        },
        mesh::Vertex {
            position: [0.0, -0.5, -0.5],
            tan_frame: 0,
        },
        mesh::Vertex {
            position: [0.0, 0.5, -0.5],
            tan_frame: 0,
        },
    ];

    let mesh_pass = resources.get::<RenderPassBox<mesh::MeshPass>>().unwrap();
    let attributes = mesh_pass.get_attributes();

    let example_vertices = AttrBuffer {
        id: attributes.vertex,
        count: example_vertices.len(),
        data: bytemuck::cast_slice(&example_vertices).to_vec(),
    };

    let example_indices: Vec<mesh::Index> = vec![0, 1, 2, 1, 2, 3, 2, 3, 0, 0, 1, 3];
    let example_indices = AttrBuffer {
        id: attributes.index,
        count: example_indices.len(),
        data: bytemuck::cast_slice(&example_indices).to_vec(),
    };

    let mut example_mesh = MeshBuffer::default();
    example_mesh.attributes.push(example_vertices);
    example_mesh.attributes.push(example_indices);
    let example_mesh = mesh_pass.get_mesh_pool().load(example_mesh).unwrap();

    world.push((
        cyborg::scene::Mesh {
            mesh: mesh_pass
                .get_mesh_pool()
                .load(make_terrain(attributes))
                .unwrap(),
        },
        cyborg::scene::Transform {
            transform: glam::Mat4::from_translation(glam::Vec3::new(0.0, -10.0, 0.0)),
        },
    ));

    drop(mesh_pass);

    let r = 6;
    for x in -r..r {
        for y in -r..r {
            for z in -r..r {
                let translation = glam::Vec3::new(x as f32, y as f32, z as f32);
                let transform = glam::Mat4::from_translation(translation);
                world.push((
                    cyborg::scene::Mesh {
                        mesh: example_mesh.clone(),
                    },
                    cyborg::scene::Transform { transform },
                    cyborg::scene::DebugDrawList::default(),
                    Particle {
                        position: translation.into(),
                        mass: 1.0,
                    },
                    Velocity {
                        linear: glam::Vec3A::ZERO,
                    },
                ));
            }
        }
    }

    event_loop.run(move |event, _, control_flow| match event {
        Event::RedrawRequested(_) => {
            let mut application = resources.get_mut::<Application>().unwrap();
            match application.viewport.acquire() {
                Err(wgpu::SurfaceError::Lost) => application.on_surface_lost(),
                Err(wgpu::SurfaceError::OutOfMemory) => *control_flow = ControlFlow::Exit,
                Err(e) => eprintln!("error: {:?}", e),
                Ok(_) => {
                    drop(application);
                    render_schedule.execute(&mut world, &mut resources);
                }
            }
        }
        Event::MainEventsCleared => {
            update_schedule.execute(&mut world, &mut resources);
            shader_watcher.watch();
        }
        Event::DeviceEvent { ref event, .. } => match event {
            DeviceEvent::MouseMotion { delta } => {
                let mut application = resources.get_mut::<Application>().unwrap();
                application.on_mouse_motion(delta.0, delta.1);
            }
            _ => {}
        },
        Event::WindowEvent {
            ref event,
            window_id,
        } => {
            let mut application = resources.get_mut::<Application>().unwrap();
            if window_id == application.window.id() {
                match event {
                    WindowEvent::KeyboardInput {
                        input:
                            KeyboardInput {
                                virtual_keycode: Some(key),
                                state,
                                ..
                            },
                        ..
                    } => {
                        application.on_key_event(*key, *state);
                    }
                    WindowEvent::MouseInput { button, state, .. } => {
                        application.on_mouse_event(*button, *state);
                    }
                    WindowEvent::CloseRequested => *control_flow = ControlFlow::Exit,
                    WindowEvent::Resized(physical_size) => {
                        application.on_resize(*physical_size);
                    }
                    WindowEvent::ScaleFactorChanged { new_inner_size, .. } => {
                        application.on_resize(**new_inner_size);
                    }
                    WindowEvent::Focused(false) => {
                        application.on_lost_focus();
                    }
                    _ => {}
                }
            }
        }
        _ => {}
    });
}