// This example is modified from https://github.com/bevyengine/bevy/blob/latest/examples/games/breakout.rs
// Non-trivial points which differ for this example will be marked as NEW.
use bevy::{
    prelude::*,
    sprite::{
        collide_aabb::{collide, Collision},
        MaterialMesh2dBundle,
    },
    DefaultPlugins,
};
use bevy_wasm_scripting::*;
use wasmer::{imports, Function, FunctionEnv, FunctionEnvMut};

// Defines the amount of time that should elapse between each physics step.
const TIME_STEP: f32 = 1.0 / 60.0;

// These constants are defined in `Transform` units.
// Using the default 2D camera they correspond 1:1 with screen pixels.
const PADDLE_SIZE: Vec3 = Vec3::new(120.0, 20.0, 0.0);
const GAP_BETWEEN_PADDLE_AND_FLOOR: f32 = 60.0;
const PADDLE_SPEED: f32 = 500.0;
// How close can the paddle get to the wall
const PADDLE_PADDING: f32 = 10.0;

// We set the z-value of the ball to 1 so it renders on top in the case of overlapping sprites.
const BALL_STARTING_POSITION: Vec3 = Vec3::new(0.0, -50.0, 1.0);
const BALL_SIZE: Vec3 = Vec3::new(30.0, 30.0, 0.0);
const BALL_SPEED: f32 = 400.0;
const INITIAL_BALL_DIRECTION: Vec2 = Vec2::new(0.5, -0.5);

const WALL_THICKNESS: f32 = 10.0;
// x coordinates
const LEFT_WALL: f32 = -450.;
const RIGHT_WALL: f32 = 450.;
// y coordinates
const BOTTOM_WALL: f32 = -300.;
const TOP_WALL: f32 = 300.;

const BRICK_SIZE: Vec2 = Vec2::new(100., 30.);
// These values are exact
const GAP_BETWEEN_PADDLE_AND_BRICKS: f32 = 270.0;
const GAP_BETWEEN_BRICKS: f32 = 5.0;
// These values are lower bounds, as the number of bricks is computed
const GAP_BETWEEN_BRICKS_AND_CEILING: f32 = 20.0;
const GAP_BETWEEN_BRICKS_AND_SIDES: f32 = 20.0;

const SCOREBOARD_FONT_SIZE: f32 = 40.0;
const SCOREBOARD_TEXT_PADDING: Val = Val::Px(5.0);

const BACKGROUND_COLOR: Color = Color::rgb(0.9, 0.9, 0.9);
const PADDLE_COLOR: Color = Color::rgb(0.3, 0.3, 0.7);
const BALL_COLOR: Color = Color::rgb(1.0, 0.5, 0.5);
const BRICK_COLOR: Color = Color::rgb(0.5, 0.5, 1.0);
const WALL_COLOR: Color = Color::rgb(0.8, 0.8, 0.8);
const TEXT_COLOR: Color = Color::rgb(0.5, 0.5, 1.0);
const SCORE_COLOR: Color = Color::rgb(1.0, 0.5, 0.5);

fn main() {
    App::new()
        .add_plugins(DefaultPlugins)
        .add_plugin(WasmPlugin)
        .insert_resource(Scoreboard { score: 0 })
        .insert_resource(ClearColor(BACKGROUND_COLOR))
        .insert_resource(FixedTime::new_from_secs(TIME_STEP))
        .add_startup_system(setup)
        .add_event::<CollisionEvent>()
        .add_system(check_for_collisions.in_schedule(CoreSchedule::FixedUpdate))
        .add_system(
            move_paddle
                .in_schedule(CoreSchedule::FixedUpdate)
                .before(check_for_collisions),
        )
        .add_system(
            apply_velocity
                .in_schedule(CoreSchedule::FixedUpdate)
                .before(check_for_collisions),
        )
        .add_system(
            play_collision_sound
                .in_schedule(CoreSchedule::FixedUpdate)
                .after(check_for_collisions),
        )
        .add_system(update_scoreboard)
        .add_system(bevy::window::close_on_esc)
        // NEW - Update balls, using their script.
        .add_system(ball_on_update_script)
        .add_wasm_script_component::<BallScript>()
        .run();
}

#[derive(Component)]
struct Paddle;

#[derive(Component)]
struct Ball;

// NEW - We're going to add scripts to balls!
#[derive(Component)]
struct BallScript(Handle<WasmScript>);

impl<'w> WasmScriptComponent for BallScript {
    type ImportQueriedComponents = (&'static Transform, &'static mut Velocity);

    type ImportResources = ();

    fn get_wasm_script_handle(&self) -> &Handle<WasmScript> {
        &self.0
    }

    fn get_imports_from_world(
        wasmer_store: &mut WasmerStore,
        world: &WorldPointer,
    ) -> wasmer::Imports {
        let env = FunctionEnv::new(&mut wasmer_store.0, world.clone());
        imports! {
            "env" => {
                "get_velocity_x" => Function::new_typed_with_env(&mut wasmer_store.0, &env, get_velocity_x),
                "get_velocity_y" => Function::new_typed_with_env(&mut wasmer_store.0, &env, get_velocity_y),
                "set_velocity" => Function::new_typed_with_env(&mut wasmer_store.0, &env, set_velocity),
                "spawn_new_ball" => Function::new_typed_with_env(&mut wasmer_store.0, &env, spawn_new_ball),
            }
        }
    }
}

// NEW - Defined for FFI and Component below.
// #[derive(Component, Deref, DerefMut)]
// struct Velocity(Vec2);

/** NEW
 * Define some structs and functions that the ball scripts can use!
 */
#[derive(Component, Clone)] // We could also have a Component struct and a FFI struct, instead.
#[repr(C)]
pub struct Velocity(pub f32, pub f32);

#[derive(Component)]
struct Collider;

// NEW - Add some more information, to support extra behaviors on collision.
pub enum CollisionEvent {
    BallWall(Entity, Entity),
    BallBrick(Entity, Entity),
}

#[derive(Component)]
struct Brick;

#[derive(Resource)]
struct CollisionSound(Handle<AudioSource>);

// This bundle is a collection of the components that define a "wall" in our game
#[derive(Bundle)]
struct WallBundle {
    // You can nest bundles inside of other bundles like this
    // Allowing you to compose their functionality
    sprite_bundle: SpriteBundle,
    collider: Collider,
}

/// Which side of the arena is this wall located on?
enum WallLocation {
    Left,
    Right,
    Bottom,
    Top,
}

impl WallLocation {
    fn position(&self) -> Vec2 {
        match self {
            WallLocation::Left => Vec2::new(LEFT_WALL, 0.),
            WallLocation::Right => Vec2::new(RIGHT_WALL, 0.),
            WallLocation::Bottom => Vec2::new(0., BOTTOM_WALL),
            WallLocation::Top => Vec2::new(0., TOP_WALL),
        }
    }

    fn size(&self) -> Vec2 {
        let arena_height = TOP_WALL - BOTTOM_WALL;
        let arena_width = RIGHT_WALL - LEFT_WALL;
        // Make sure we haven't messed up our constants
        assert!(arena_height > 0.0);
        assert!(arena_width > 0.0);

        match self {
            WallLocation::Left | WallLocation::Right => {
                Vec2::new(WALL_THICKNESS, arena_height + WALL_THICKNESS)
            }
            WallLocation::Bottom | WallLocation::Top => {
                Vec2::new(arena_width + WALL_THICKNESS, WALL_THICKNESS)
            }
        }
    }
}

impl WallBundle {
    // This "builder method" allows us to reuse logic across our wall entities,
    // making our code easier to read and less prone to bugs when we change the logic
    fn new(location: WallLocation) -> WallBundle {
        WallBundle {
            sprite_bundle: SpriteBundle {
                transform: Transform {
                    // We need to convert our Vec2 into a Vec3, by giving it a z-coordinate
                    // This is used to determine the order of our sprites
                    translation: location.position().extend(0.0),
                    // The z-scale of 2D objects must always be 1.0,
                    // or their ordering will be affected in surprising ways.
                    // See https://github.com/bevyengine/bevy/issues/4149
                    scale: location.size().extend(1.0),
                    ..default()
                },
                sprite: Sprite {
                    color: WALL_COLOR,
                    ..default()
                },
                ..default()
            },
            collider: Collider,
        }
    }
}

// This resource tracks the game's score
#[derive(Resource)]
struct Scoreboard {
    score: usize,
}

/** NEW
 * Define some functions that the ball scripts can use!
 */
fn get_velocity_x(env: FunctionEnvMut<WorldPointer>, entity_id: EntityId) -> f32 {
    env.data()
        .read()
        .get::<Velocity>(entity_id.to_entity())
        .cloned()
        .map(|vel| vel.0)
        .unwrap_or(0.0)
}

fn get_velocity_y(env: FunctionEnvMut<WorldPointer>, entity_id: EntityId) -> f32 {
    env.data()
        .read()
        .get::<Velocity>(entity_id.to_entity())
        .cloned()
        .map(|vel| vel.1)
        .unwrap_or(0.0)
}

fn set_velocity(env: FunctionEnvMut<WorldPointer>, entity_id: EntityId, vx: f32, vy: f32) {
    if let Some(mut entity_velocity) = env
        .data()
        .write()
        .get_mut::<Velocity>(entity_id.to_entity())
    {
        *entity_velocity = Velocity(vx, vy);
    }
    println!("Setting {} {}", vx, vy);
}

fn spawn_new_ball(
    env: FunctionEnvMut<WorldPointer>,
    entity_id: EntityId,
    vx: f32,
    vy: f32,
    speed: f32,
) -> EntityId {
    let entity = entity_id.to_entity();
    let world_pointer = env.data();
    let mut commands = world_pointer.commands::<BallScript>();
    let mut meshes = world_pointer.write().resource_mut::<Assets<Mesh>>();
    let mut materials = world_pointer
        .write()
        .resource_mut::<Assets<ColorMaterial>>();
    if let Some(ball_transform) = world_pointer.read().get::<Transform>(entity) {
        f64::from_bits(
            spawn_ball(
                &mut commands,
                &mut meshes,
                &mut materials,
                ball_transform.translation,
                (vx, vy),
                speed,
            )
            .to_bits(),
        )
    } else {
        f64::from_bits(u64::MAX)
    }
}

// NEW - Both the imported function and the setup function will call this general function.
// This is a refactor of the original code into a function, with some added parameters.
fn spawn_ball(
    commands: &mut Commands,
    meshes: &mut Assets<Mesh>,
    materials: &mut Assets<ColorMaterial>,
    location: Vec3,
    direction: (f32, f32),
    speed: f32,
) -> Entity {
    commands
        .spawn((
            MaterialMesh2dBundle {
                mesh: meshes.add(shape::Circle::default().into()).into(),
                material: materials.add(ColorMaterial::from(BALL_COLOR)),
                transform: Transform::from_translation(location).with_scale(BALL_SIZE),
                ..default()
            },
            Ball,
            Velocity(direction.0 * speed, direction.1 * speed),
        ))
        .id()
}

// Add the game's entities to our world
fn setup(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<ColorMaterial>>,
    asset_server: Res<AssetServer>,
) {
    // Camera
    commands.spawn(Camera2dBundle::default());

    // Sound
    let ball_collision_sound = asset_server.load("sounds/breakout_collision.ogg");
    commands.insert_resource(CollisionSound(ball_collision_sound));

    // Paddle
    let paddle_y = BOTTOM_WALL + GAP_BETWEEN_PADDLE_AND_FLOOR;

    commands.spawn((
        SpriteBundle {
            transform: Transform {
                translation: Vec3::new(0.0, paddle_y, 0.0),
                scale: PADDLE_SIZE,
                ..default()
            },
            sprite: Sprite {
                color: PADDLE_COLOR,
                ..default()
            },
            ..default()
        },
        Paddle,
        Collider,
    ));

    // Scoreboard
    commands.spawn(
        TextBundle::from_sections([
            TextSection::new(
                "Score: ",
                TextStyle {
                    font: asset_server.load("fonts/FiraSans-Bold.ttf"),
                    font_size: SCOREBOARD_FONT_SIZE,
                    color: TEXT_COLOR,
                },
            ),
            TextSection::from_style(TextStyle {
                font: asset_server.load("fonts/FiraMono-Medium.ttf"),
                font_size: SCOREBOARD_FONT_SIZE,
                color: SCORE_COLOR,
            }),
        ])
        .with_style(Style {
            position_type: PositionType::Absolute,
            position: UiRect {
                top: SCOREBOARD_TEXT_PADDING,
                left: SCOREBOARD_TEXT_PADDING,
                ..default()
            },
            ..default()
        }),
    );

    // Ball
    let ball = spawn_ball(
        &mut commands,
        meshes.as_mut(),
        materials.as_mut(),
        BALL_STARTING_POSITION,
        INITIAL_BALL_DIRECTION.into(),
        BALL_SPEED,
    );
    commands.entity(ball).insert(BallScript(
        asset_server.load("breakout/scripts/gravity.wasm"),
    ));

    // Walls
    commands.spawn(WallBundle::new(WallLocation::Left));
    commands.spawn(WallBundle::new(WallLocation::Right));
    commands.spawn(WallBundle::new(WallLocation::Bottom));
    commands.spawn(WallBundle::new(WallLocation::Top));

    // Bricks
    // Negative scales result in flipped sprites / meshes,
    // which is definitely not what we want here
    assert!(BRICK_SIZE.x > 0.0);
    assert!(BRICK_SIZE.y > 0.0);

    let total_width_of_bricks = (RIGHT_WALL - LEFT_WALL) - 2. * GAP_BETWEEN_BRICKS_AND_SIDES;
    let bottom_edge_of_bricks = paddle_y + GAP_BETWEEN_PADDLE_AND_BRICKS;
    let total_height_of_bricks = TOP_WALL - bottom_edge_of_bricks - GAP_BETWEEN_BRICKS_AND_CEILING;

    assert!(total_width_of_bricks > 0.0);
    assert!(total_height_of_bricks > 0.0);

    // Given the space available, compute how many rows and columns of bricks we can fit
    let n_columns = (total_width_of_bricks / (BRICK_SIZE.x + GAP_BETWEEN_BRICKS)).floor() as usize;
    let n_rows = (total_height_of_bricks / (BRICK_SIZE.y + GAP_BETWEEN_BRICKS)).floor() as usize;
    let n_vertical_gaps = n_columns - 1;

    // Because we need to round the number of columns,
    // the space on the top and sides of the bricks only captures a lower bound, not an exact value
    let center_of_bricks = (LEFT_WALL + RIGHT_WALL) / 2.0;
    let left_edge_of_bricks = center_of_bricks
        // Space taken up by the bricks
        - (n_columns as f32 / 2.0 * BRICK_SIZE.x)
        // Space taken up by the gaps
        - n_vertical_gaps as f32 / 2.0 * GAP_BETWEEN_BRICKS;

    // In Bevy, the `translation` of an entity describes the center point,
    // not its bottom-left corner
    let offset_x = left_edge_of_bricks + BRICK_SIZE.x / 2.;
    let offset_y = bottom_edge_of_bricks + BRICK_SIZE.y / 2.;

    for row in 0..n_rows {
        for column in 0..n_columns {
            let brick_position = Vec2::new(
                offset_x + column as f32 * (BRICK_SIZE.x + GAP_BETWEEN_BRICKS),
                offset_y + row as f32 * (BRICK_SIZE.y + GAP_BETWEEN_BRICKS),
            );

            // brick
            commands.spawn((
                SpriteBundle {
                    sprite: Sprite {
                        color: BRICK_COLOR,
                        ..default()
                    },
                    transform: Transform {
                        translation: brick_position.extend(0.0),
                        scale: Vec3::new(BRICK_SIZE.x, BRICK_SIZE.y, 1.0),
                        ..default()
                    },
                    ..default()
                },
                Brick,
                Collider,
            ));
        }
    }
}

fn move_paddle(
    keyboard_input: Res<Input<KeyCode>>,
    mut query: Query<&mut Transform, With<Paddle>>,
) {
    let mut paddle_transform = query.single_mut();
    let mut direction = 0.0;

    if keyboard_input.pressed(KeyCode::Left) {
        direction -= 1.0;
    }

    if keyboard_input.pressed(KeyCode::Right) {
        direction += 1.0;
    }

    // Calculate the new horizontal paddle position based on player input
    let new_paddle_position = paddle_transform.translation.x + direction * PADDLE_SPEED * TIME_STEP;

    // Update the paddle position,
    // making sure it doesn't cause the paddle to leave the arena
    let left_bound = LEFT_WALL + WALL_THICKNESS / 2.0 + PADDLE_SIZE.x / 2.0 + PADDLE_PADDING;
    let right_bound = RIGHT_WALL - WALL_THICKNESS / 2.0 - PADDLE_SIZE.x / 2.0 - PADDLE_PADDING;

    paddle_transform.translation.x = new_paddle_position.clamp(left_bound, right_bound);
}

fn apply_velocity(mut query: Query<(&mut Transform, &Velocity)>) {
    for (mut transform, velocity) in &mut query {
        transform.translation.x += velocity.0 * TIME_STEP;
        transform.translation.y += velocity.1 * TIME_STEP;
    }
}

fn ball_on_update_script(
    mut script_env: WasmScriptComponentEnv<BallScript, Without<BallScript>>,
    ball_query: Query<(Entity, &BallScript)>,
    time: Res<Time>,
) {
    for (entity, script) in ball_query.iter() {
        if let Err(err) = script_env.call_if_instantiated_2::<f64, f32, ()>(
            &script.0,
            "on_update",
            f64::from_bits(entity.to_bits()),
            time.delta_seconds(),
        ) {
            bevy::log::error!("Failed to run script: {}", err);
        }
    }
}

fn update_scoreboard(scoreboard: Res<Scoreboard>, mut query: Query<&mut Text>) {
    let mut text = query.single_mut();
    text.sections[1].value = scoreboard.score.to_string();
}

fn check_for_collisions(
    mut commands: Commands,
    mut scoreboard: ResMut<Scoreboard>,
    mut ball_query: Query<(Entity, &mut Velocity, &Transform), With<Ball>>,
    collider_query: Query<(Entity, &Transform, Option<&Brick>), With<Collider>>,
    mut collision_events: EventWriter<CollisionEvent>,
) {
    let (ball_entity, mut ball_velocity, ball_transform) = ball_query.single_mut();
    let ball_size = ball_transform.scale.truncate();

    // check collision with walls
    for (collider_entity, transform, maybe_brick) in &collider_query {
        let collision = collide(
            ball_transform.translation,
            ball_size,
            transform.translation,
            transform.scale.truncate(),
        );
        if let Some(collision) = collision {
            // Sends a collision event so that other systems can react to the collision
            // NEW - Don't send default, send based on what we collide with.
            // collision_events.send_default();

            // Bricks should be despawned and increment the scoreboard on collision
            if maybe_brick.is_some() {
                scoreboard.score += 1;
                collision_events.send(CollisionEvent::BallBrick(ball_entity, collider_entity));
                commands.entity(collider_entity).despawn();
            } else {
                collision_events.send(CollisionEvent::BallWall(ball_entity, collider_entity));
            }

            // reflect the ball when it collides
            let mut reflect_x = false;
            let mut reflect_y = false;

            // only reflect if the ball's velocity is going in the opposite direction of the
            // collision
            match collision {
                Collision::Left => reflect_x = ball_velocity.0 > 0.0,
                Collision::Right => reflect_x = ball_velocity.0 < 0.0,
                Collision::Top => reflect_y = ball_velocity.1 < 0.0,
                Collision::Bottom => reflect_y = ball_velocity.1 > 0.0,
                Collision::Inside => { /* do nothing */ }
            }

            // reflect velocity on the x-axis if we hit something on the x-axis
            if reflect_x {
                ball_velocity.0 = -ball_velocity.0;
            }

            // reflect velocity on the y-axis if we hit something on the y-axis
            if reflect_y {
                ball_velocity.1 = -ball_velocity.1;
            }
        }
    }
}

fn play_collision_sound(
    mut collision_events: EventReader<CollisionEvent>,
    audio: Res<Audio>,
    sound: Res<CollisionSound>,
) {
    // Play a sound once per frame if a collision occurred.
    if !collision_events.is_empty() {
        // This prevents events staying active on the next frame.
        collision_events.clear();
        audio.play(sound.0.clone());
    }
}