use argh::FromArgs;
use bevy::{prelude::*, utils::HashMap, utils::HashSet};
use bevy_spatial::{kdtree::KDTree3, SpatialAccess};

pub type NNTree = KDTree3<Toid>;

// toid
const SPEED: f32 = 1.0;
const SPEED_DIFF_RANGE: f32 = 0.08; // +/- 8%

// how far from origin before really wanting to come back
const RADIUS: f32 = 800.0;

// how close to try to stay to your buddies
const BUDDY_RADIUS: f32 = 15.0;

// how much room to try to maintain between toids
const PERSONAL_SPACE: f32 = 1.5;

const MIN_ALTITUDE: f32 = 1.5;

#[derive(Debug, FromArgs, Resource)]
/// Toid Watching
pub struct Config {
    /// how many Toids to spawn
    #[argh(option, short = 't', default = "100")]
    pub toids: usize,
}

#[derive(Clone, Debug, Default, Deref, DerefMut, Resource)]
pub struct Positions(pub HashMap<Entity, Vec3>);

#[derive(Component, Debug, Clone, Default)]
pub struct Buddies(HashSet<Entity>);

#[derive(Component, Debug, Clone, Deref, DerefMut, Default)]
pub struct Velocity(Vec3);

#[derive(Component)]
pub struct Toid {
    pub speed: f32,
    pub buddies: usize,
}

pub fn turkey_time(
    commands: &mut Commands,
    scene: &Handle<Scene>,
    r: &mut impl rand::prelude::Rng,
) -> Entity {
    let speed_diff = r.gen_range(-SPEED_DIFF_RANGE..=SPEED_DIFF_RANGE);
    let speed = SPEED + (SPEED * speed_diff);
    let vel = unit_vec(r) * speed;
    let buddies = r.gen_range(6..=8);
    let x = r.gen_range(-5.0..=5.0);
    let z = r.gen_range(-5.0..=5.0);
    let y = r.gen_range(0.1..=2.5);
    let spatial_bundle = SpatialBundle {
        transform: Transform::from_xyz(x, MIN_ALTITUDE + y, z),
        ..Default::default()
    };
    commands
        .spawn(spatial_bundle)
        .insert((Velocity(vel), Buddies::default(), Toid { speed, buddies }))
        .with_children(|t| {
            t.spawn(SceneBundle {
                scene: scene.to_owned(),
                ..Default::default()
            })
            .insert(Transform::from_rotation(Quat::from_axis_angle(
                Vec3::Y,
                -std::f32::consts::FRAC_PI_2,
            )));
        })
        .id()
}

pub fn update_vel(
    mut toids: Query<(&Transform, &mut Velocity, &Toid, &Buddies)>,
    positions: Res<Positions>,
) {
    for (xform, mut vel, toid, buddies) in toids.iter_mut() {
        let speed = toid.speed;
        let mut dir = vel.normalize();
        let pos = xform.translation;

        // nudge up if too low
        if pos.y < MIN_ALTITUDE {
            let dh = MIN_ALTITUDE - pos.y;
            let pct = (dh / MIN_ALTITUDE).min(1.0).powi(2);
            let rot = Quat::from_rotation_arc(dir, Vec3::Y);
            let rot = Quat::IDENTITY.slerp(rot, pct);
            dir = rot.mul_vec3(dir).normalize();
        }

        // nudge toward origin if too far
        {
            let dist = pos.length();
            let toward_origin = -pos.normalize();
            let rot = Quat::from_rotation_arc(dir, toward_origin);
            let pct = (dist / RADIUS).min(1.0);
            let rot = Quat::IDENTITY.slerp(rot, pct);
            dir = rot.mul_vec3(dir).normalize();
        }

        // find buddies and orient; point more towards further-away buddies
        for buddy in buddies.0.iter() {
            let bp = positions.get(buddy).unwrap_or(&pos);
            let bdir = *bp - pos;
            let dist = bdir.length();
            let rot = Quat::from_rotation_arc(dir, bdir.normalize());
            let s = (dist / (BUDDY_RADIUS * 1.2)).min(1.0);
            let rot = Quat::IDENTITY.slerp(rot, s);
            dir = rot.mul_vec3(dir).normalize();
        }

        **vel = dir * speed;
    }
}

pub fn update_pos(
    mut toids: Query<(&mut Transform, &Velocity, Entity), With<Toid>>,
    mut positions: ResMut<Positions>,
    time: Res<Time>,
) {
    let dt = time.delta_seconds();
    for (mut xform, vel, entity) in toids.iter_mut() {
        let look_at = xform.translation + vel.0;
        xform.translation += vel.0 * dt;
        xform.look_at(look_at, Vec3::Y);
        let _ = positions.insert(entity, xform.translation);
    }
}

pub fn update_buddies(
    mut toids: Query<(&Transform, Entity, &Toid, &mut Buddies)>,
    index: Res<NNTree>,
    positions: Res<Positions>,
) {
    let d2 = (BUDDY_RADIUS * 1.5).powi(2);
    for (xform, entity, toid, mut buddies) in toids.iter_mut() {
        let pos = xform.translation;

        for buddy in buddies.0.clone().iter() {
            let bp = positions.get(buddy).unwrap_or(&pos);
            let bd2 = (*bp - pos).length_squared();
            if bd2 > d2 {
                buddies.0.remove(buddy);
            }
        }

        if buddies.0.len() < toid.buddies {
            let mut neighbors = index
                .within_distance(pos, BUDDY_RADIUS)
                .into_iter()
                .filter(|n| n.1.is_some() && n.1.unwrap() != entity);
            let diff = toid.buddies - buddies.0.len();
            for _ in 0..diff {
                if let Some(neighbor) = neighbors.next() {
                    buddies.0.insert(neighbor.1.unwrap());
                }
            }
        }
    }
}

//-************************************************************************
// util
//-************************************************************************

pub fn unit_vec(r: &mut impl rand::Rng) -> Vec3 {
    let mut x1: f32 = 0.0;
    let mut x2: f32 = 0.0;
    let mut ssum = std::f32::MAX;
    while ssum >= 1.0 {
        x1 = r.gen_range(-1.0..=1.0);
        x2 = r.gen_range(-1.0..=1.0);
        ssum = x1.powi(2) + x2.powi(2);
    }
    let sqrt = (1.0 - ssum).sqrt();
    let x = 2.0 * x1 * sqrt;
    let y = 2.0 * x2 * sqrt;
    let z = 1.0 - 2.0 * ssum;
    Vec3::new(x, y, z).normalize()
}