use bevy::{
    core_pipeline::{bloom::Bloom, tonemapping::Tonemapping},
    input::mouse::{MouseScrollUnit, MouseWheel},
    math::DVec3,
    prelude::*,
    window::{WindowMode, WindowResolution},
};
use bevy_inspector_egui::{bevy_egui::EguiPlugin, quick::WorldInspectorPlugin};
use iyes_perf_ui::{PerfUiPlugin, prelude::*};
use serde::Deserialize;

mod jpl_horizon;

// Data structures for deserializing initial_state.ron
#[derive(Deserialize)]
struct InitialState {
    bodies: Vec<CelestialBodyData>,
}

#[derive(Deserialize)]
struct CelestialBodyData {
    name: String,
    mass: f64,      // kg
    radius: f64,    // km
    position: (f64, f64, f64), // AU
    velocity: (f64, f64, f64), // AU/day
}

// Scaling factor to convert AU to game units
// Neptune is approximately 30.1 astronomical units (AU) from the Sun
// Point light effective range is about 10 game units, so 10 game units = ~30 AU
const AU_TO_GAME_UNITS: f64 = 0.3;

// Unit wrapper types - all distances in AU
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct DistanceAu(pub f64);

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct PositionAu(pub DVec3);

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct VelocityAuPerDay(pub DVec3);

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct MassKg(pub f64);

#[derive(Clone, Debug, PartialEq)]
pub struct ObjectName(pub String);

// Component wrappers
#[derive(Component)]
struct Position(PositionAu);

#[derive(Component)]
struct Velocity(VelocityAuPerDay);

#[derive(Component)]
struct Mass(MassKg);

#[derive(Component)]
struct Radius(DistanceAu);

#[derive(Component)]
struct Name(ObjectName);

#[derive(Bundle)]
struct ObjectBundle {
    name: Name,
    position: Position,
    mass: Mass,
    radius: Radius,
    velocity: Velocity,
}

#[derive(Component)]
struct Star;
#[derive(Component)]
struct Earth;

// Component for UI name labels
#[derive(Component)]
struct ObjectLabel {
    target_entity: Entity,
}

// Component to mark objects that can be focused on with proper zoom levels
#[derive(Component)]
struct Trackable {}

// Resource to track which entity the camera should follow
#[derive(Resource)]
struct CameraFollow {
    target: Option<Entity>,
    distance: f32, // Current zoom distance from target
}

pub struct SolarRenderingPlugin;

impl Plugin for SolarRenderingPlugin {
    fn build(&self, app: &mut App) {
        app.insert_resource(CameraFollow {
            target: None,
            distance: 10.0,
        })
        .add_systems(Startup, (setup_rendering, setup_ui))
        .add_systems(
            FixedPostUpdate,
            (
                sync_radius_to_mesh,
                sync_position_to_transform,
                sync_name_labels,
                camera_follow_system.after(sync_position_to_transform),
            ),
        )
        .add_systems(PostUpdate, update_label_positions)
        .add_systems(
            Update,
            (
                handle_label_clicks,
                handle_scroll_zoom,
                handle_speed_controls,
            ),
        );
    }
}

fn sync_radius_to_mesh(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
    query: Query<(Entity, &Radius, Option<&Star>, Option<&Mesh3d>), Changed<Radius>>,
) {
    for (entity, radius, star, existing_mesh) in query.iter() {
        // Convert AU to game units for rendering
        let render_radius = radius.0.0 * AU_TO_GAME_UNITS;

        // Create or update sphere mesh
        let sphere_mesh = meshes.add(Sphere::new(render_radius as f32));
        let material = materials.add(if star.is_none() {
            StandardMaterial {
                base_color: Color::WHITE,
                ..default()
            }
        } else {
            StandardMaterial {
                base_color: Color::WHITE, // Sun-like color for stars
                emissive: LinearRgba::rgb(200.0, 200.0, 200.0),
                ..default()
            }
        });

        if existing_mesh.is_none() {
            // Add mesh and material components if they don't exist
            commands
                .entity(entity)
                .insert((Mesh3d(sphere_mesh), MeshMaterial3d(material)));
        } else {
            // Update existing mesh
            commands.entity(entity).insert(Mesh3d(sphere_mesh));
        }
    }
}

fn sync_position_to_transform(
    mut commands: Commands,
    mut query: Query<(Entity, &Position, Option<&mut Transform>), Changed<Position>>,
) {
    for (entity, position, transform) in query.iter_mut() {
        // Convert AU to game units for rendering
        let scaled_position = position.0.0 * AU_TO_GAME_UNITS;
        match transform {
            Some(mut t) => {
                // Update existing transform
                t.translation = scaled_position.as_vec3();
            }
            None => {
                // Insert a new Transform if it doesn't exist
                commands
                    .entity(entity)
                    .insert(Transform::from_translation(scaled_position.as_vec3()));
            }
        };
    }
}

fn setup_rendering(mut commands: Commands) {
    // Spawn camera with pan/orbit/zoom controls
    // Place it at a good distance to view the solar system
    commands.spawn((
        Camera3d::default(),
        Camera {
            hdr: true,
            clear_color: ClearColorConfig::Custom(Color::BLACK),
            ..default()
        },
        Projection::Perspective(PerspectiveProjection {
            near: 1e-9, // Very close near plane for extreme zooming
            ..default()
        }),
        Tonemapping::TonyMcMapface,
        Transform::from_translation(Vec3::new(0., 0., 10.0)),
        Bloom::NATURAL,
    ));

    commands.spawn(PointLight {
        color: Color::WHITE,
        shadows_enabled: true,
        ..default()
    });
    commands.spawn(PerfUiAllEntries::default());
}

fn setup_ui(mut commands: Commands) {
    // UI root node for labels
    commands.spawn(Node {
        width: Val::Percent(100.0),
        height: Val::Percent(100.0),
        position_type: PositionType::Absolute,
        ..default()
    });
}

fn sync_name_labels(
    mut commands: Commands,
    objects_with_names: Query<(Entity, &Name, &Position), Changed<Name>>,
    existing_labels: Query<&ObjectLabel>,
) {
    for (entity, name, _position) in objects_with_names.iter() {
        // Check if label already exists for this entity
        let has_label = existing_labels
            .iter()
            .any(|label| label.target_entity == entity);

        if !has_label {
            // Create new label with scaled font size
            commands.spawn((
                Text::new(name.0.0.clone()),
                TextColor(Color::WHITE),
                TextFont {
                    font_size: 16.0,
                    ..default()
                },
                Node {
                    position_type: PositionType::Absolute,
                    left: Val::Px(0.0), // Will be updated by update_label_positions
                    top: Val::Px(0.0),  // Will be updated by update_label_positions
                    ..default()
                },
                BackgroundColor(Color::srgba(0.0, 0.0, 0.0, 0.7)),
                Interaction::default(), // Make label clickable
                ObjectLabel {
                    target_entity: entity,
                },
            ));
        }
    }
}

fn update_label_positions(
    mut label_query: Query<(&mut Node, &ObjectLabel, &mut Text), With<ObjectLabel>>,
    objects_query: Query<(&GlobalTransform, &Name, &Radius)>,
    camera_query: Query<(&Camera, &GlobalTransform)>,
) {
    let Ok((camera, camera_transform)) = camera_query.single() else {
        return;
    };

    for (mut node, label, mut text) in label_query.iter_mut() {
        if let Ok((global_transform, name, _radius)) = objects_query.get(label.target_entity) {
            let world_pos = global_transform.translation();
            if let Ok(screen_pos) = camera.world_to_viewport(camera_transform, world_pos) {
                // Apply scale factor to label offset to maintain relative positioning
                let scaled_offset = 10.0;

                // Position the label on screen, offset slightly to avoid overlapping the object
                node.left = Val::Px(screen_pos.x + scaled_offset);
                node.top = Val::Px(screen_pos.y - scaled_offset);

                // Update text in case name changed
                text.0 = name.0.0.clone();
            } else {
                // Object is off-screen, hide label by moving it off-screen
                node.left = Val::Px(-1000.0);
                node.top = Val::Px(-1000.0);
            }
        } else {
            // Target entity no longer exists, remove label
            // Note: In a more complex system, you might want to handle this in a separate cleanup system
        }
    }
}

fn handle_label_clicks(
    interaction_query: Query<
        (&Interaction, &ObjectLabel),
        (Changed<Interaction>, With<ObjectLabel>),
    >,
    trackable_objects: Query<(&Transform, &Trackable, &Radius)>,
    mut camera_query: Query<&mut Transform, (With<Camera>, Without<Trackable>)>,
    mut camera_follow: ResMut<CameraFollow>,
) {
    for (interaction, label) in interaction_query.iter() {
        if *interaction == Interaction::Pressed {
            if let Ok((target_transform, _trackable, radius)) =
                trackable_objects.get(label.target_entity)
            {
                if let Ok(mut camera_transform) = camera_query.single_mut() {
                    let target_world_pos = target_transform.translation;

                    // Calculate appropriate distance based on entity radius (with some padding)
                    let entity_radius = (radius.0.0 * AU_TO_GAME_UNITS) as f32;
                    let desired_distance = entity_radius * 16.0; // 8x radius

                    // Position camera at desired distance from target
                    let camera_direction =
                        (camera_transform.translation - target_world_pos).normalize();
                    let new_camera_pos = target_world_pos + camera_direction * desired_distance;

                    camera_transform.translation = new_camera_pos;
                    camera_transform.look_at(target_world_pos, Vec3::Y);

                    // Set the follow target and initial distance
                    camera_follow.target = Some(label.target_entity);
                    camera_follow.distance = desired_distance;
                }
            }
        }
    }
}

fn camera_follow_system(
    camera_follow: Res<CameraFollow>,
    objects_query: Query<&Transform, With<Trackable>>,
    mut camera_query: Query<&mut Transform, (With<Camera>, Without<Trackable>)>,
) {
    if let Some(target_entity) = camera_follow.target {
        if let Ok(target_transform) = objects_query.get(target_entity) {
            if let Ok(mut camera_transform) = camera_query.single_mut() {
                // Fixed camera direction relative to tracked object (stationary relative position)
                let fixed_direction = Vec3::new(0.0, 0.0, 1.0);

                // Position camera at the specified distance from target in fixed direction
                let new_camera_pos =
                    target_transform.translation + fixed_direction * camera_follow.distance;
                camera_transform.translation = new_camera_pos;

                // Keep camera looking at the target entity
                camera_transform.look_at(target_transform.translation, Vec3::Y);
            }
        }
    }
}

fn handle_scroll_zoom(
    mut scroll_events: EventReader<MouseWheel>,
    mut camera_follow: ResMut<CameraFollow>,
) {
    for event in scroll_events.read() {
        let zoom_delta = match event.unit {
            MouseScrollUnit::Line => event.y * 0.1,
            MouseScrollUnit::Pixel => event.y * 0.01,
        };

        // Update zoom distance with limits
        camera_follow.distance = (camera_follow.distance * (1.0 - zoom_delta)).clamp(1e-9, 100.0);
    }
}

fn handle_speed_controls(
    keyboard_input: Res<ButtonInput<KeyCode>>,
    mut time: ResMut<Time<Virtual>>,
) {
    // Pause/Unpause with SPACE
    if keyboard_input.just_pressed(KeyCode::Space) {
        if time.is_paused() {
            time.unpause();
            println!("Simulation RESUMED (Speed: {}x)", time.relative_speed());
        } else {
            time.pause();
            println!("Simulation PAUSED");
        }
    }

    // Speed up with + or = key (double current speed, max 64x)
    if keyboard_input.just_pressed(KeyCode::Equal)
        || keyboard_input.just_pressed(KeyCode::NumpadAdd)
    {
        if !time.is_paused() {
            let current_speed = time.relative_speed();
            let new_speed = (current_speed * 2.0).min(64.0);
            time.set_relative_speed(new_speed);
            println!("Speed: {}x", new_speed);
        }
    }

    // Speed down with - key (halve current speed, min 1.0x)
    if keyboard_input.just_pressed(KeyCode::Minus)
        || keyboard_input.just_pressed(KeyCode::NumpadSubtract)
    {
        if !time.is_paused() {
            let current_speed = time.relative_speed();
            let new_speed = (current_speed * 0.5).max(1.0);
            time.set_relative_speed(new_speed);
            println!("Speed: {}x", new_speed);
        }
    }

    // Reset to normal speed with R
    if keyboard_input.just_pressed(KeyCode::KeyR) {
        time.set_relative_speed(1.0);
        time.unpause();
        println!("Speed reset to 1x");
    }
}

fn initialize_camera_follow(
    earth_query: Query<Entity, (With<Earth>, With<Trackable>)>,
    mut camera_follow: ResMut<CameraFollow>,
) {
    if let Ok(earth) = earth_query.single() {
        camera_follow.target = Some(earth);
        camera_follow.distance = 1e-3;
    }
}

fn setup_solar_system(mut commands: Commands) {
    // Load initial state from RON file
    let initial_state_content = match std::fs::read_to_string("assets/initial_state.ron") {
        Ok(content) => content,
        Err(err) => {
            error!("Failed to read initial_state.ron: {}", err);
            return;
        }
    };

    let initial_state: InitialState = match ron::from_str(&initial_state_content) {
        Ok(state) => state,
        Err(err) => {
            error!("Failed to parse initial_state.ron: {}", err);
            return;
        }
    };

    // Spawn all celestial bodies from the initial state
    for body_data in initial_state.bodies {
        // Convert radius from km to AU
        let radius_au = body_data.radius / 149597870.691; // km to AU conversion

        // Create base bundle
        let mut entity_commands = commands.spawn((
            ObjectBundle {
                name: Name(ObjectName(body_data.name.clone())),
                position: Position(PositionAu(DVec3::new(
                    body_data.position.0,
                    body_data.position.1, 
                    body_data.position.2,
                ))),
                velocity: Velocity(VelocityAuPerDay(DVec3::new(
                    body_data.velocity.0,
                    body_data.velocity.1,
                    body_data.velocity.2,
                ))),
                mass: Mass(MassKg(body_data.mass)),
                radius: Radius(DistanceAu(radius_au)),
            },
            Trackable {},
        ));

        // Add special components for Sun and Earth
        match body_data.name.as_str() {
            "Sun" => {
                entity_commands.insert(Star);
            }
            "Earth" => {
                entity_commands.insert(Earth);
            }
            _ => {}
        }

        info!("Spawned celestial body: {}", body_data.name);
    }
}

// High precision constants
const G_SI: f64 = 6.67430e-11; // m^3 / (kg s^2)
const AU_TO_M: f64 = 149_597_870_691.0; // m
const DAY_TO_S: f64 = 86_400.0; // s

// G in AU^3 / (kg day^2)
const G_AU: f64 = G_SI * (DAY_TO_S * DAY_TO_S) / (AU_TO_M * AU_TO_M * AU_TO_M);

const STEPS: usize = 100;
// If FixedUpdate is 64 Hz and you want 1 day per game second: DT = 1/(64*STEPS) day/step
const DT: f64 = 1.0 / (64.0 * STEPS as f64);

fn n_body(mut query: Query<(&Mass, &mut Position, &mut Velocity)>) {
    // Pull data out once. This copy lets us integrate freely without borrow issues.
    // (If you want zero copies, use a Resource scratch buffer and reuse it across frames.)
    let mut bodies: Vec<(f64, DVec3, DVec3)> = query
        .iter()
        .map(|(mass, pos, vel)| (mass.0.0, pos.0.0, vel.0.0))
        .collect();

    let n = bodies.len();
    if n < 2 {
        // Nothing to do; write back if needed and return
        for (i, (_m, mut pos, mut vel)) in query.iter_mut().enumerate() {
            if let Some(b) = bodies.get(i) {
                pos.0.0 = b.1;
                vel.0.0 = b.2;
            }
        }
        return;
    }

    // Single reusable acceleration buffer
    let mut acc: Vec<DVec3> = vec![DVec3::ZERO; n];

    // Pairwise acceleration accumulation (O(n^2)/2) with Newton's 3rd law
    let compute_accelerations = |bodies: &[(f64, DVec3, DVec3)], acc: &mut [DVec3]| {
        acc.fill(DVec3::ZERO);

        for i in 0..n - 1 {
            let (mi, pi, _) = bodies[i];
            for j in (i + 1)..n {
                let (mj, pj, _) = bodies[j];

                let r = pj - pi; // AU
                let r2 = r.length_squared(); // AU^2
                // Tighter epsilon in AU^2; avoids division by ~0 without killing legit close passes
                if r2 > 1e-20 {
                    // a_i = G * m_j * r / r^3, a_j = -G * m_i * r / r^3
                    let inv_r = r2.sqrt().recip(); // 1 / r
                    let inv_r3 = inv_r * inv_r * inv_r; // 1 / r^3

                    let a_i = r * (G_AU * mj * inv_r3);
                    let a_j = r * (G_AU * mi * inv_r3);

                    acc[i] += a_i;
                    acc[j] -= a_j; // opposite direction
                }
            }
        }
    };

    // Leapfrog / velocity-Verlet
    for _ in 0..STEPS {
        // v += a(q) * dt/2
        compute_accelerations(&bodies, &mut acc);
        for i in 0..n {
            bodies[i].2 += acc[i] * (DT * 0.5);
        }

        // q += v * dt
        for i in 0..n {
            let vel = bodies[i].2; // copy DVec3 (cheap, it's just 3 f64s)
            bodies[i].1 += vel * DT;
        }

        // v += a(q_new) * dt/2
        compute_accelerations(&bodies, &mut acc);
        for i in 0..n {
            bodies[i].2 += acc[i] * (DT * 0.5);
        }
    }

    // Write back
    for (i, (_mass, mut pos, mut vel)) in query.iter_mut().enumerate() {
        let (_, p, v) = bodies[i];
        pos.0.0 = p;
        vel.0.0 = v;
    }
}

fn main() {
    App::new()
        .add_plugins(DefaultPlugins.set(WindowPlugin {
            primary_window: Some(Window {
                title: "Solar Sim".to_string(),
                mode: WindowMode::BorderlessFullscreen(MonitorSelection::Primary),
                resolution: WindowResolution::default().with_scale_factor_override(2.0),
                ..default()
            }),
            ..default()
        }))
        .add_plugins(EguiPlugin::default())
        .add_plugins(WorldInspectorPlugin::new())
        .add_plugins(SolarRenderingPlugin) // we want Bevy to measure these values for us:
        .add_plugins(bevy::diagnostic::FrameTimeDiagnosticsPlugin::default())
        .add_plugins(bevy::diagnostic::EntityCountDiagnosticsPlugin)
        .add_plugins(bevy::diagnostic::SystemInformationDiagnosticsPlugin)
        .add_plugins(bevy::render::diagnostic::RenderDiagnosticsPlugin)
        .add_plugins(PerfUiPlugin)
        .add_systems(
            Startup,
            (setup_solar_system, initialize_camera_follow).chain(),
        )
        .add_systems(FixedUpdate, n_body)
        .run();
}