use rapier2d::prelude::*;

use crate::protocol::*;

const LINEAR_DAMPING: f32 = 0.142;
const ANGULAR_DAMPING: f32 = 0.18;
const MAX_SIM_TIME: f32 = 30.0;
const REST_SPEED: f32 = 0.04;
const REST_ANGULAR_SPEED: f32 = 0.05;
// Rotation rate of the velocity vector, in rad/s.
// Positive curl_sign = right curl -> curves toward +x when moving up-sheet.
const CURL_RATE: f32 = 0.010;

pub struct PhysicsWorld {
    gravity: Vector,
    integration_parameters: IntegrationParameters,
    pipeline: PhysicsPipeline,
    islands: IslandManager,
    broad_phase: DefaultBroadPhase,
    narrow_phase: NarrowPhase,
    bodies: RigidBodySet,
    colliders: ColliderSet,
    impulse_joints: ImpulseJointSet,
    multibody_joints: MultibodyJointSet,
    ccd_solver: CCDSolver,
    next_stone_id: u32,
    stone_handles: Vec<(u32, RigidBodyHandle, Team, i8)>,
}

impl Default for PhysicsWorld {
    fn default() -> Self {
        Self::new()
    }
}

impl PhysicsWorld {
    pub fn new() -> Self {
        let mut integration_parameters = IntegrationParameters::default();
        integration_parameters.dt = PHYSICS_DT;
        integration_parameters.num_solver_iterations = 8;

        let mut world = Self {
            gravity: Vector::new(0.0, 0.0),
            integration_parameters,
            pipeline: PhysicsPipeline::new(),
            islands: IslandManager::new(),
            broad_phase: DefaultBroadPhase::new(),
            narrow_phase: NarrowPhase::new(),
            bodies: RigidBodySet::new(),
            colliders: ColliderSet::new(),
            impulse_joints: ImpulseJointSet::new(),
            multibody_joints: MultibodyJointSet::new(),
            ccd_solver: CCDSolver::new(),
            next_stone_id: 1,
            stone_handles: Vec::new(),
        };
        world.build_sheet();
        world
    }

    pub fn reset(&mut self) {
        self.bodies = RigidBodySet::new();
        self.colliders = ColliderSet::new();
        self.islands = IslandManager::new();
        self.broad_phase = DefaultBroadPhase::new();
        self.narrow_phase = NarrowPhase::new();
        self.impulse_joints = ImpulseJointSet::new();
        self.multibody_joints = MultibodyJointSet::new();
        self.ccd_solver = CCDSolver::new();
        self.stone_handles.clear();
        self.build_sheet();
    }

    pub fn reset_stone_ids(&mut self) {
        self.next_stone_id = 1;
    }

    fn build_sheet(&mut self) {
        let half = SHEET_WIDTH / 2.0 + 0.1;
        let left = ColliderBuilder::cuboid(0.1, SHEET_LENGTH / 2.0 + 1.0)
            .translation(Vector::new(-half, SHEET_LENGTH / 2.0))
            .friction(0.0)
            .restitution(0.0)
            .build();
        self.colliders.insert(left);

        let right = ColliderBuilder::cuboid(0.1, SHEET_LENGTH / 2.0 + 1.0)
            .translation(Vector::new(half, SHEET_LENGTH / 2.0))
            .friction(0.0)
            .restitution(0.0)
            .build();
        self.colliders.insert(right);

        let back = ColliderBuilder::cuboid(SHEET_WIDTH / 2.0 + 1.0, 0.1)
            .translation(Vector::new(0.0, SHEET_LENGTH + 0.1))
            .friction(0.0)
            .restitution(0.1)
            .build();
        self.colliders.insert(back);
    }

    pub fn throw(
        &mut self,
        team: Team,
        broom_x: f32,
        broom_y: f32,
        weight: u8,
        curl: i8,
        friction: f32,
    ) -> Result<Vec<StoneTrajectory>, String> {
        let weight = weight.clamp(1, 10) as f32;
        let t = (weight - 1.0) / 9.0;
        let speed = MIN_SPEED + t * (MAX_SPEED - MIN_SPEED);
        let dx = broom_x;
        let dy = broom_y - HACK_Y;
        let len = (dx * dx + dy * dy).sqrt().max(0.01);
        let vx = dx / len * speed;
        let vy = dy / len * speed;

        let curl_sign = if curl < 0 { -1 } else { 1 };
        let damping_mult = friction.clamp(0.5, 2.0);

        self.spawn_stone(team, 0.0, HACK_Y, vx, vy, curl_sign, damping_mult)
    }

    fn spawn_stone(
        &mut self,
        team: Team,
        x: f32,
        y: f32,
        vx: f32,
        vy: f32,
        curl_sign: i8,
        damping_mult: f32,
    ) -> Result<Vec<StoneTrajectory>, String> {
        let id = self.next_stone_id;
        self.next_stone_id += 1;

        let body = RigidBodyBuilder::dynamic()
            .translation(Vector::new(x, y))
            .linvel(Vector::new(vx, vy))
            .angvel(0.0)
            .linear_damping(LINEAR_DAMPING * damping_mult)
            .angular_damping(ANGULAR_DAMPING)
            .can_sleep(false)
            .build();

        let handle = self.bodies.insert(body);

        let collider = ColliderBuilder::ball(STONE_RADIUS)
            .friction(STONE_FRICTION)
            .friction_combine_rule(CoefficientCombineRule::Average)
            .restitution(STONE_RESTITUTION)
            .restitution_combine_rule(CoefficientCombineRule::Average)
            .density(STONE_MASS / (std::f32::consts::PI * STONE_RADIUS * STONE_RADIUS))
            .build();

        self.colliders.insert_with_parent(collider, handle, &mut self.bodies);
        self.stone_handles.push((id, handle, team, curl_sign));

        self.simulate_until_rest(id)
    }

    fn simulate_until_rest(&mut self, thrown_id: u32) -> Result<Vec<StoneTrajectory>, String> {
        // All paths share the thrown stone's release instant as t=0. This keeps the
        // frontend's existing trajectory helpers (which expect the thrown stone to
        // start at x=0, y=HACK_Y with t=0) working unchanged while also giving every
        // other stone a consistent timeline.
        let sample_step = 1.0 / SAMPLE_RATE_HZ as f32;
        let mut sample_accum: f32 = 0.0;
        let mut time: f32 = 0.0;

        // Pre-allocate a path buffer for every stone currently in the world.
        let mut paths: Vec<(u32, RigidBodyHandle, Vec<(f32, f32, f32)>)> = self
            .stone_handles
            .iter()
            .map(|(id, handle, _, _)| (*id, *handle, Vec::new()))
            .collect();

        // Record the initial sample at t=0 for every stone.
        for (id, handle, path) in &mut paths {
            if let Some(body) = self.bodies.get(*handle) {
                let pos = body.translation();
                path.push((pos.x, pos.y, time));
            } else {
                // Body missing for an tracked stone; this should not happen.
                return Err(format!("stone {} has no rigid body", id));
            }
        }

        loop {
            self.step();
            self.apply_curl();
            time += PHYSICS_DT;
            sample_accum += PHYSICS_DT;

            if sample_accum >= sample_step {
                sample_accum -= sample_step;
                for (_, handle, path) in &mut paths {
                    if let Some(body) = self.bodies.get(*handle) {
                        let pos = body.translation();
                        path.push((pos.x, pos.y, time));
                    }
                }
            }

            if self.all_stones_at_rest() || time > MAX_SIM_TIME {
                break;
            }
        }

        self.prune_out_of_play();

        // The thrown stone is released at (0.0, HACK_Y). Shift every path in time so
        // that t=0 corresponds to that release instant. Because we already started
        // sampling at the release instant, the first sample time is 0.0 and no shift
        // is required; this comment documents the invariant.
        let thrown_first_t = paths
            .iter()
            .find(|(id, _, _)| *id == thrown_id)
            .and_then(|(_, _, path)| path.first().map(|(_, _, t)| *t))
            .unwrap_or(0.0);

        Ok(paths
            .into_iter()
            .map(|(id, _, mut path)| {
                if thrown_first_t != 0.0 {
                    for (_, _, t) in &mut path {
                        *t -= thrown_first_t;
                    }
                }
                StoneTrajectory { stone_id: id, path }
            })
            .collect())
    }

    // Rotate each stone's velocity slightly based on its selected curl direction.
    // Right curl (curl_sign = +1) curves toward +x when moving up-sheet (positive y).
    fn apply_curl(&mut self) {
        for (_, handle, _, curl_sign) in &self.stone_handles {
            let body = match self.bodies.get_mut(*handle) {
                Some(b) => b,
                None => continue,
            };

            let v = body.linvel();
            let speed_sq = v.x * v.x + v.y * v.y;
            let speed = speed_sq.sqrt();
            if speed < 1e-4 {
                continue;
            }

            let angle = -(*curl_sign as f32) * CURL_RATE * PHYSICS_DT;
            let cos = angle.cos();
            let sin = angle.sin();
            let new_v = Vector::new(v.x * cos - v.y * sin, v.x * sin + v.y * cos);
            body.set_linvel(new_v, true);
        }
    }

    fn prune_out_of_play(&mut self) {
        let mut keep = Vec::new();
        for (id, handle, team, curl) in self.stone_handles.drain(..) {
            if let Some(body) = self.bodies.get(handle) {
                let pos = body.translation();
                let beyond_back = pos.y > BACK_LINE_Y;
                let short_of_hog = pos.y < HOG_LINE_Y;
                let outside = pos.x.abs() > SHEET_WIDTH / 2.0;
                if beyond_back || short_of_hog || outside {
                    self.bodies.remove(handle, &mut self.islands, &mut self.colliders, &mut self.impulse_joints, &mut self.multibody_joints, true);
                } else {
                    keep.push((id, handle, team, curl));
                }
            }
        }
        self.stone_handles = keep;
    }

    fn step(&mut self) {
        self.pipeline.step(
            self.gravity,
            &self.integration_parameters,
            &mut self.islands,
            &mut self.broad_phase,
            &mut self.narrow_phase,
            &mut self.bodies,
            &mut self.colliders,
            &mut self.impulse_joints,
            &mut self.multibody_joints,
            &mut self.ccd_solver,
            &(),
            &(),
        );
    }

    fn all_stones_at_rest(&self) -> bool {
        for (_, handle, _, _) in &self.stone_handles {
            if let Some(body) = self.bodies.get(*handle) {
                let v = body.linvel();
                let speed = (v.x * v.x + v.y * v.y).sqrt();
                if speed > REST_SPEED || body.angvel().abs() > REST_ANGULAR_SPEED {
                    return false;
                }
            }
        }
        true
    }

    pub fn current_stones(&self) -> Vec<StoneState> {
        let mut states = Vec::new();
        for (id, handle, team, _) in &self.stone_handles {
            if let Some(body) = self.bodies.get(*handle) {
                let pos = body.translation();
                states.push(StoneState {
                    id: *id,
                    team: *team,
                    x: pos.x,
                    y: pos.y,
                    rotation: body.rotation().angle(),
                    active: false,
                });
            }
        }
        states
    }

    pub fn stone_states_for_scoring(&self) -> Vec<(u32, Team, f32, f32)> {
        let mut out = Vec::new();
        for (id, handle, team, _) in &self.stone_handles {
            if let Some(body) = self.bodies.get(*handle) {
                let pos = body.translation();
                out.push((*id, *team, pos.x, pos.y));
            }
        }
        out
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn final_y(world: &PhysicsWorld, id: u32) -> f32 {
        world.stone_handles.iter()
            .find(|(sid, _, _, _)| *sid == id)
            .map(|(_, h, _, _)| {
                let b = &world.bodies[*h];
                b.translation().y
            })
            .unwrap_or(f32::NAN)
    }

    fn final_x(world: &PhysicsWorld, id: u32) -> f32 {
        world.stone_handles.iter()
            .find(|(sid, _, _, _)| *sid == id)
            .map(|(_, h, _, _)| {
                let b = &world.bodies[*h];
                b.translation().x
            })
            .unwrap_or(f32::NAN)
    }

    #[test]
    fn weight_7_lands_on_tee_line() {
        let mut world = PhysicsWorld::new();
        world.throw(Team::Red, 0.0, HOUSE_CENTER.1, 7, 1, 1.0).unwrap();
        let id = world.next_stone_id - 1;
        let y = final_y(&world, id);
        println!("weight 7 final y={}", y);
        assert!(
            (y - HOUSE_CENTER.1).abs() <= 0.5,
            "weight-7 draw shot should finish on the tee line, got y={}",
            y
        );
    }

    #[test]
    fn curl_direction_mirrors_x_offset() {
        let mut right = PhysicsWorld::new();
        right.throw(Team::Red, 0.0, HOUSE_CENTER.1, 7, 1, 1.0).unwrap();
        let right_id = right.next_stone_id - 1;
        let right_x = final_x(&right, right_id);

        let mut left = PhysicsWorld::new();
        left.throw(Team::Red, 0.0, HOUSE_CENTER.1, 7, -1, 1.0).unwrap();
        let left_id = left.next_stone_id - 1;
        let left_x = final_x(&left, left_id);

        println!("right curl final x={} left curl final x={}", right_x, left_x);
        assert!(
            right_x > left_x + 0.05,
            "right curl should finish to the right of left curl: right={} left={}",
            right_x,
            left_x
        );
    }

    #[test]
    fn stones_persist_after_multiple_throws() {
        let mut world = PhysicsWorld::new();
        world.throw(Team::Red, 0.0, HOUSE_CENTER.1, 7, 1, 1.0).unwrap();
        world.throw(Team::Red, 0.0, HOUSE_CENTER.1, 7, -1, 1.0).unwrap();

        let stones = world.current_stones();
        assert_eq!(stones.len(), 2, "both stones should remain in the physics world");
        assert_eq!(stones[0].id, 1);
        assert_eq!(stones[1].id, 2);
    }

    #[test]
    fn out_of_play_stone_is_pruned() {
        // A very light, high-friction throw should stop short of the hog line and be removed.
        let mut world = PhysicsWorld::new();
        world.throw(Team::Red, 0.0, HOUSE_CENTER.1, 1, 0, 2.0).unwrap();
        let stones = world.current_stones();
        assert!(stones.is_empty(), "stones short of the hog line should be pruned");
    }

    #[test]
    fn collision_records_trajectories_for_both_stones() {
        // Place a stationary stone on the center line and throw a second stone
        // straight at it so they collide. Both stones must have sampled paths.
        let mut world = PhysicsWorld::new();

        // First stone: place it far enough up-sheet to stay in play after impact.
        world
            .throw(Team::Red, 0.0, HOUSE_CENTER.1, 7, 0, 1.0)
            .unwrap();
        let first_id = world.next_stone_id - 1;

        // Second stone: aimed directly at the first stone's final position.
        let target_y = final_y(&world, first_id);
        let target_x = final_x(&world, first_id);
        world
            .throw(Team::Yellow, target_x, target_y, 10, 0, 1.0)
            .unwrap();
        let second_id = world.next_stone_id - 1;

        // Re-run the collision throw and capture trajectories.
        let mut world = PhysicsWorld::new();
        world
            .throw(Team::Red, 0.0, HOUSE_CENTER.1, 7, 0, 1.0)
            .unwrap();
        let first_id = world.next_stone_id - 1;
        let target_y = final_y(&world, first_id);
        let target_x = final_x(&world, first_id);
        let trajectories = world
            .throw(Team::Yellow, target_x, target_y, 10, 0, 1.0)
            .unwrap();

        let by_id: std::collections::HashMap<u32, Vec<(f32, f32, f32)>> = trajectories
            .into_iter()
            .map(|st| (st.stone_id, st.path))
            .collect();

        assert!(
            by_id.contains_key(&first_id),
            "trajectories should contain the first stone (id={})",
            first_id
        );
        assert!(
            by_id.contains_key(&second_id),
            "trajectories should contain the thrown stone (id={})",
            second_id
        );

        let first_path = by_id.get(&first_id).unwrap();
        let second_path = by_id.get(&second_id).unwrap();
        assert!(
            first_path.len() > 1,
            "first stone path should have multiple samples, got {}",
            first_path.len()
        );
        assert!(
            second_path.len() > 1,
            "thrown stone path should have multiple samples, got {}",
            second_path.len()
        );

        // Both paths should share the same t=0 reference (the thrown stone's release).
        assert_eq!(first_path[0].2, 0.0, "first stone path should start at t=0");
        assert_eq!(second_path[0].2, 0.0, "thrown stone path should start at t=0");
    }
}