#include "simulation.h"
#include <cstdlib>
#include <cstring>
#include <cstdio>
#include <cmath>

// Create a new simulation
SimulationState* create_simulation(int max_bodies, double time_step) {
    SimulationState* sim = (SimulationState*)malloc(sizeof(SimulationState));
    sim->bodies = (CelestialBody*)malloc(sizeof(CelestialBody) * max_bodies);
    sim->body_count = 0;
    sim->max_bodies = max_bodies;
    sim->time = 0.0;
    sim->dt = time_step;
    return sim;
}

// Destroy simulation and free memory
void destroy_simulation(SimulationState* sim) {
    if (sim) {
        if (sim->bodies) {
            free(sim->bodies);
        }
        free(sim);
    }
}

// Add a celestial body to the simulation
void add_body(SimulationState* sim, const char* name, double mass, double radius,
              Vec3 pos, Vec3 vel, int parent_index, float r, float g, float b,
              double eccentricity, double semi_major_axis) {
    if (sim->body_count >= sim->max_bodies) {
        return;  // No more space
    }

    CelestialBody* body = &sim->bodies[sim->body_count];
    strncpy(body->name, name, 63);
    body->name[63] = '\0';
    body->mass = mass;
    body->radius = radius;
    body->position = pos;
    body->velocity = vel;
    body->soi_radius = 0.0;  // Will be calculated later
    body->parent_index = parent_index;
    body->color[0] = r;
    body->color[1] = g;
    body->color[2] = b;
    body->eccentricity = eccentricity;
    body->semi_major_axis = semi_major_axis;

    sim->body_count++;
}

// Find which body is gravitationally dominant for the given body
int find_dominant_body(SimulationState* sim, int body_index) {
    if (body_index < 0 || body_index >= sim->body_count) {
        return -1;
    }

    CelestialBody* body = &sim->bodies[body_index];
    int dominant = body->parent_index;

    // Check all other bodies to see if we're within their SOI
    for (int i = 0; i < sim->body_count; i++) {
        if (i == body_index) continue;

        CelestialBody* potential_parent = &sim->bodies[i];
        double distance = vec3_distance(body->position, potential_parent->position);

        // If we're within this body's SOI and it's not our current parent
        if (distance < potential_parent->soi_radius && i != dominant) {
            // Check if this body is more dominant (closer or more massive)
            if (dominant == -1) {
                dominant = i;
            } else {
                CelestialBody* current_parent = &sim->bodies[dominant];
                double dist_to_current = vec3_distance(body->position, current_parent->position);

                // Switch if this potential parent is significantly closer
                if (distance < dist_to_current * 0.5) {
                    dominant = i;
                }
            }
        }
    }

    return dominant;
}

// Update sphere of influence radius using Hill sphere approximation
// r_soi = a * (m/M)^(2/5) where a = semi-major axis, m = body mass, M = parent mass
void update_soi(CelestialBody* body, CelestialBody* parent, double semi_major_axis) {
    if (parent == NULL || parent->mass <= 0.0) {
        // Root body (like Sun) has infinite SOI, use a large value
        body->soi_radius = 1e15;  // 1000 AU in meters
        return;
    }

    double mass_ratio = body->mass / parent->mass;
    body->soi_radius = semi_major_axis * pow(mass_ratio, 0.4);  // 2/5 = 0.4
}

void update_simulation(SimulationState* sim) {
    for (int i = 0; i < sim->body_count; i++) {
        CelestialBody* body = &sim->bodies[i];

        if (body->parent_index == -1) {
            AccelerationContext ctx;
            ctx.sim = sim;
            ctx.current_body = body;
            ctx.body_index = i;

            rk4_step(body, &ctx, sim->dt);
        }
    }

    for (int i = 0; i < sim->body_count; i++) {
        CelestialBody* body = &sim->bodies[i];

        if (body->parent_index == -1) {
            continue;
        }

        int new_parent = find_dominant_body(sim, i);
        if (new_parent != body->parent_index && new_parent != -1) {
            body->parent_index = new_parent;
        }

        if (body->parent_index >= 0 && body->parent_index < sim->body_count) {
            AccelerationContext ctx;
            ctx.sim = sim;
            ctx.current_body = body;
            ctx.body_index = i;

            rk4_step(body, &ctx, sim->dt);
        }
    }

    sim->time += sim->dt;
}

static void compute_perpendicular_orbital_velocity(CelestialBody* body, Vec3 center,
                                                   double orbiting_mass) {
    Vec3 r = vec3_sub(body->position, center);
    double distance = vec3_magnitude(r);

    if (distance < 1.0) {
        body->velocity = {0.0, 0.0, 0.0};
        return;
    }

    double speed = sqrt(G * orbiting_mass / distance);
    Vec3 z_axis = {0.0, 0.0, 1.0};
    Vec3 vel_dir = vec3_cross(r, z_axis);

    if (vec3_magnitude(vel_dir) < 0.01) {
        Vec3 x_axis = {1.0, 0.0, 0.0};
        vel_dir = vec3_cross(r, x_axis);
    }

    vel_dir = vec3_normalize(vel_dir);
    body->velocity = vec3_scale(vel_dir, speed);
}

static void compute_orbital_velocity_from_vis_viva(CelestialBody* body,
                                                     CelestialBody* parent) {
    Vec3 r = vec3_sub(body->position, parent->position);
    double distance = vec3_magnitude(r);

    if (distance < 1.0) {
        body->velocity = {0.0, 0.0, 0.0};
        return;
    }

    double e = body->eccentricity;
    double a = body->semi_major_axis;
    double speed = sqrt(G * parent->mass * (2.0 / distance - 1.0 / a));

    if (e > 0.01) {
        printf("  %s: eccentric orbit (e=%.2f, a=%.3e m), speed at r=%.3e m: %.3f km/s\n",
               body->name, e, a, distance, speed / 1000.0);
    }

    Vec3 z_axis = {0.0, 0.0, 1.0};
    Vec3 vel_dir = vec3_cross(r, z_axis);

    if (vec3_magnitude(vel_dir) < 0.01) {
        Vec3 x_axis = {1.0, 0.0, 0.0};
        vel_dir = vec3_cross(r, x_axis);
    }

    vel_dir = vec3_normalize(vel_dir);
    body->velocity = vec3_scale(vel_dir, speed);
    body->velocity = vec3_add(body->velocity, parent->velocity);
}

static Vec3 compute_system_barycenter(SimulationState* sim, int* root_indices,
                                      int* root_count) {
    Vec3 barycenter = {0.0, 0.0, 0.0};
    *root_count = 0;

    for (int i = 0; i < sim->body_count; i++) {
        if (sim->bodies[i].parent_index == -1) {
            if (*root_count < 32) {
                root_indices[(*root_count)++] = i;
                Vec3 weighted_pos = vec3_scale(sim->bodies[i].position, sim->bodies[i].mass);
                barycenter = vec3_add(barycenter, weighted_pos);
            }
        }
    }

    double total_mass = 0.0;
    for (int i = 0; i < *root_count; i++) {
        total_mass += sim->bodies[root_indices[i]].mass;
    }

    if (total_mass > 0.0) {
        barycenter = vec3_scale(barycenter, 1.0 / total_mass);
    }

    return barycenter;
}

static double compute_total_root_mass(SimulationState* sim, int* root_indices,
                                      int root_count) {
    double total_mass = 0.0;
    for (int i = 0; i < root_count; i++) {
        total_mass += sim->bodies[root_indices[i]].mass;
    }
    return total_mass;
}

static void set_root_bodies_velocity(SimulationState* sim, int* root_indices,
                                     int root_count, Vec3 barycenter, double total_mass) {
    for (int i = 0; i < root_count; i++) {
        CelestialBody* body = &sim->bodies[root_indices[i]];
        double other_mass = total_mass - body->mass;
        compute_perpendicular_orbital_velocity(body, barycenter, other_mass);

        double distance = vec3_magnitude(vec3_sub(body->position, barycenter));
        double speed = vec3_magnitude(body->velocity);
        printf("  %s: distance from barycenter = %.3e m, orbital speed = %.3e m/s\n",
               body->name, distance, speed);
    }
}

static void set_child_bodies_velocity(SimulationState* sim) {
    for (int i = 0; i < sim->body_count; i++) {
        CelestialBody* body = &sim->bodies[i];

        if (body->parent_index == -1) {
            continue;
        }

        if (body->parent_index >= 0 && body->parent_index < sim->body_count) {
            CelestialBody* parent = &sim->bodies[body->parent_index];
            compute_orbital_velocity_from_vis_viva(body, parent);
        }
    }
}

static void print_system_info_if_multiple_roots(int root_count, Vec3 barycenter) {
    if (root_count > 1) {
        printf("\nBinary/Multiple star system detected:\n");
        printf("  Number of root bodies: %d\n", root_count);
        printf("  Barycenter position: (%.3e, %.3e, %.3e) m\n",
               barycenter.x, barycenter.y, barycenter.z);
    }
}

void compute_initial_velocities(SimulationState* sim) {
    int root_indices[32];
    int root_count;
    Vec3 barycenter = compute_system_barycenter(sim, root_indices, &root_count);

    print_system_info_if_multiple_roots(root_count, barycenter);

    if (root_count > 1) {
        double total_mass = compute_total_root_mass(sim, root_indices, root_count);
        set_root_bodies_velocity(sim, root_indices, root_count, barycenter, total_mass);
    } else if (root_count == 1) {
        sim->bodies[root_indices[0]].velocity = {0.0, 0.0, 0.0};
    }

    set_child_bodies_velocity(sim);
}

void calculate_initial_velocities(SimulationState* sim) {
    compute_initial_velocities(sim);
}

// Calculate SOI radii for all bodies
void calculate_soi_radii(SimulationState* sim) {
    for (int i = 0; i < sim->body_count; i++) {
        CelestialBody* body = &sim->bodies[i];

        if (body->parent_index == -1) {
            body->soi_radius = 1e15;
        } else if (body->parent_index >= 0 && body->parent_index < sim->body_count) {
            CelestialBody* parent = &sim->bodies[body->parent_index];
            update_soi(body, parent, body->semi_major_axis);
        }
    }
}

OrbitalElements calculate_orbital_elements(CelestialBody* body, CelestialBody* primary,
                                           CelestialBody* optional_ref_body, double current_time) {
    const double AU = 1.496e11;
    const double SECONDS_PER_DAY = 86400.0;
    const double M_sun = primary->mass;

    OrbitalElements elem;
    elem.time_days = current_time / SECONDS_PER_DAY;

    Vec3 r_vec = vec3_sub(body->position, primary->position);
    double r = vec3_magnitude(r_vec);
    double v = vec3_magnitude(body->velocity);

    elem.distance_to_sun_au = r / AU;
    elem.velocity_magnitude = v;

    if (optional_ref_body) {
        double dist_ref = vec3_distance(body->position, optional_ref_body->position);
        elem.distance_to_ref_body_au = dist_ref / AU;
    } else {
        elem.distance_to_ref_body_au = -1.0;
    }

    elem.specific_energy = (v * v) / 2.0 - (G * M_sun) / r;

    if (elem.specific_energy < 0) {
        elem.semi_major_axis_au = -(G * M_sun) / (2.0 * elem.specific_energy) / AU;

        double v_squared = v * v;
        double r_dot_v = r_vec.x * body->velocity.x + r_vec.y * body->velocity.y + r_vec.z * body->velocity.z;

        Vec3 e_vec;
        e_vec.x = (v_squared - G * M_sun / r) * r_vec.x - r_dot_v * body->velocity.x;
        e_vec.y = (v_squared - G * M_sun / r) * r_vec.y - r_dot_v * body->velocity.y;
        e_vec.z = (v_squared - G * M_sun / r) * r_vec.z - r_dot_v * body->velocity.z;

        double e_mag = vec3_magnitude(e_vec) / (G * M_sun);
        elem.eccentricity = e_mag;
    } else {
        elem.semi_major_axis_au = 0.0;
        elem.eccentricity = 1.0;
    }

    return elem;
}