#include "simulation.h"
#include "spacecraft.h"
#include "maneuver.h"
#include "orbital_mechanics.h"
#include <cassert>
#include <cstdlib>
#include <cstring>
#include <cstdio>
#include <cmath>

// Create a new simulation
SimulationState* create_simulation(int max_bodies, int max_craft, int max_maneuvers, double time_step) {
    SimulationState* sim = (SimulationState*)malloc(sizeof(SimulationState));
    sim->bodies = (CelestialBody*)malloc(sizeof(CelestialBody) * max_bodies);

    if (max_craft > 0) {
        sim->spacecraft = (Spacecraft*)malloc(sizeof(Spacecraft) * max_craft);
    } else {
        sim->spacecraft = NULL;
    }

    if (max_maneuvers > 0) {
        sim->maneuvers = (Maneuver*)malloc(sizeof(Maneuver) * max_maneuvers);
    } else {
        sim->maneuvers = NULL;
    }

    sim->body_count = 0;
    sim->max_bodies = max_bodies;
    sim->craft_count = 0;
    sim->max_craft = max_craft;
    sim->maneuver_count = 0;
    sim->max_maneuvers = max_maneuvers;
    sim->time = 0.0;
    sim->dt = time_step;
    sim->config_name[0] = '\0';
    return sim;
}

// Destroy simulation and free memory
void destroy_simulation(SimulationState* sim) {
    if (sim) {
        if (sim->bodies) {
            free(sim->bodies);
        }
        if (sim->max_craft > 0 && sim->spacecraft) {
            free(sim->spacecraft);
        }
        if (sim->max_maneuvers > 0 && sim->maneuvers) {
            free(sim->maneuvers);
        }
        free(sim);
    }
}

// Add a spacecraft to the simulation
int add_spacecraft(SimulationState* sim, Spacecraft* craft) {
    if (sim->craft_count >= sim->max_craft) {
        printf("Error: Cannot add spacecraft - simulation full (%d/%d)\n",
               sim->craft_count, sim->max_craft);
        return -1;
    }

    int new_idx = sim->craft_count;
    sim->spacecraft[new_idx] = *craft;
    sim->craft_count++;

    return new_idx;
}

// Add a body to the simulation at runtime
int add_body_to_simulation(SimulationState* sim, CelestialBody* body) {
    if (sim->body_count >= sim->max_bodies) {
        printf("Error: Cannot add body - simulation full (%d/%d)\n",
               sim->body_count, sim->max_bodies);
        return -1;
    }

    int new_idx = sim->body_count;
    sim->bodies[new_idx] = *body;
    sim->body_count++;

    if (body->parent_index >= 0 && body->parent_index < sim->body_count) {
        CelestialBody* parent = &sim->bodies[body->parent_index];
        sim->bodies[new_idx].local_position = vec3_sub(body->global_position, parent->global_position);
        sim->bodies[new_idx].local_velocity = vec3_sub(body->global_velocity, parent->global_velocity);
    } else {
        sim->bodies[new_idx].local_position = body->global_position;
        sim->bodies[new_idx].local_velocity = body->global_velocity;
    }

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

    sim->bodies[new_idx].global_position = body->global_position;
    sim->bodies[new_idx].global_velocity = body->global_velocity;

    return new_idx;
}

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 parent_idx = body->parent_index;

    // If parent is not root (not Sun): only check if still within parent's SOI
    if (parent_idx != 0) {
        if (parent_idx < 0 || parent_idx >= sim->body_count) {
            return -1;
        }

        CelestialBody* parent = &sim->bodies[parent_idx];
        double distance = vec3_distance(body->global_position, parent->global_position);

        // Stay with parent if within SOI, otherwise go to Sun
        if (distance < parent->soi_radius) {
            return parent_idx;
        } else {
            return 0;
        }
    }

    // Parent is root (Sun): check all bodies for SOI containment
    int new_parent = 0;
    double min_distance = INFINITY;

    for (int i = 0; i < sim->body_count; i++) {
        if (i == body_index) continue;

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

        // If within SOI and closer than current, switch to this body
        if (distance < potential->soi_radius && distance < min_distance) {
            min_distance = distance;
            new_parent = i;
        }
    }

    return new_parent;
}
// 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) {
    update_bodies_physics(sim);
    compute_global_coordinates(sim);
    update_spacecraft_physics(sim);
    execute_pending_maneuvers(sim);
    compute_spacecraft_globals(sim);
    sim->time += sim->dt;
}

// Calculate SOI radius for a single body
// r_soi = a * (m/M)^(2/5) where a = semi-major axis, m = body mass, M = parent mass
// Returns SOI radius in meters
double calculate_soi_radius(CelestialBody* body, CelestialBody* parent) {
    assert(body != nullptr && parent != nullptr);
    double mass_ratio = body->mass / parent->mass;
    return body->orbit.semi_major_axis * pow(mass_ratio, 0.4);  // 2/5 = 0.4
}

// Initialize orbital objects from orbital elements
// Converts orbital elements to local position/velocity and computes global coordinates
void initialize_orbital_objects(SimulationState* sim) {
    for (int i = 0; i < sim->body_count; i++) {
        CelestialBody* body = &sim->bodies[i];
        CelestialBody* parent = NULL;

        if (body->parent_index >= 0 && body->parent_index < sim->body_count) {
            parent = &sim->bodies[body->parent_index];

            Vec3 local_pos, local_vel;
            orbital_elements_to_cartesian(body->orbit, parent->mass, &local_pos, &local_vel);

            body->local_position = local_pos;
            body->local_velocity = local_vel;
            body->global_position = vec3_add(parent->global_position, local_pos);
            body->global_velocity = vec3_add(parent->global_velocity, local_vel);
            body->soi_radius = calculate_soi_radius(body, parent);
        } else {
            body->local_position = {0.0, 0.0, 0.0};
            body->local_velocity = {0.0, 0.0, 0.0};
            body->global_position = {0.0, 0.0, 0.0};
            body->global_velocity = {0.0, 0.0, 0.0};
            body->soi_radius = 1e15;
        }
    }

    for (int i = 0; i < sim->craft_count; i++) {
        Spacecraft* craft = &sim->spacecraft[i];

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

            Vec3 local_pos, local_vel;
            OrbitalElements elements = craft->orbit;

            orbital_elements_to_cartesian(elements, parent->mass, &local_pos, &local_vel);

            craft->local_position = local_pos;
            craft->local_velocity = local_vel;
            craft->global_position = vec3_add(parent->global_position, local_pos);
            craft->global_velocity = vec3_add(parent->global_velocity, local_vel);
        }
    }
}

// Simulation update helper functions
void update_bodies_physics(SimulationState* sim) {
    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) {
            if (body->parent_index >= 0 && body->parent_index < sim->body_count) {
                CelestialBody* old_parent = &sim->bodies[body->parent_index];
                body->global_position = vec3_add(body->local_position, old_parent->global_position);
                body->global_velocity = vec3_add(body->local_velocity, old_parent->global_velocity);
            } else {
                body->global_position = body->local_position;
                body->global_velocity = body->local_velocity;
            }

            body->parent_index = new_parent;

            if (body->parent_index >= 0 && body->parent_index < sim->body_count) {
                CelestialBody* new_parent_body = &sim->bodies[body->parent_index];
                body->local_position = vec3_sub(body->global_position, new_parent_body->global_position);
                body->local_velocity = vec3_sub(body->global_velocity, new_parent_body->global_velocity);
            } else {
                body->local_position = body->global_position;
                body->local_velocity = body->global_velocity;
            }
        }

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

            rk4_step(&body->local_position, &body->local_velocity,
                     sim->dt, body->mass, parent->mass);
        }
    }
}

void update_spacecraft_physics(SimulationState* sim) {
    for (int i = 0; i < sim->craft_count; i++) {
        Spacecraft* craft = &sim->spacecraft[i];

        if (craft->parent_index < 0 || craft->parent_index >= sim->body_count) {
            continue;
        }

        CelestialBody* parent = &sim->bodies[craft->parent_index];

        rk4_step(&craft->local_position, &craft->local_velocity,
                 sim->dt, craft->mass, parent->mass);
    }
}

void execute_pending_maneuvers(SimulationState* sim) {
    for (int i = 0; i < sim->maneuver_count; i++) {
        Maneuver* maneuver = &sim->maneuvers[i];

        if (maneuver->executed) {
            continue;
        }

        if (maneuver->craft_index < 0 || maneuver->craft_index >= sim->craft_count) {
            continue;
        }

        Spacecraft* craft = &sim->spacecraft[maneuver->craft_index];

        if (check_maneuver_trigger(maneuver, craft, sim)) {
            execute_maneuver(maneuver, craft, sim->time);
        }
    }
}

void compute_spacecraft_globals(SimulationState* sim) {
    for (int i = 0; i < sim->craft_count; i++) {
        Spacecraft* craft = &sim->spacecraft[i];

        if (craft->parent_index >= 0 && craft->parent_index < sim->body_count) {
            CelestialBody* parent = &sim->bodies[craft->parent_index];
            craft->global_position = vec3_add(parent->global_position, craft->local_position);
            craft->global_velocity = vec3_add(parent->global_velocity, craft->local_velocity);
        } else {
            craft->global_position = craft->local_position;
            craft->global_velocity = craft->local_velocity;
        }
    }
}

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

        if (body->parent_index == -1) {
            body->global_position = body->local_position;
            body->global_velocity = body->local_velocity;
        } else if (body->parent_index >= 0 && body->parent_index < sim->body_count) {
            CelestialBody* parent = &sim->bodies[body->parent_index];
            body->global_position = vec3_add(body->local_position, parent->global_position);
            body->global_velocity = vec3_add(body->local_velocity, parent->global_velocity);
        }
    }
}

OrbitalAnalysis 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;

    OrbitalAnalysis elem;
    elem.time_days = current_time / SECONDS_PER_DAY;

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

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

    if (optional_ref_body) {
        double dist_ref = vec3_distance(body->global_position, optional_ref_body->global_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->global_velocity.x + r_vec.y * body->global_velocity.y + r_vec.z * body->global_velocity.z;

        Vec3 e_vec;
        e_vec.x = (v_squared - G * M_sun / r) * r_vec.x - r_dot_v * body->global_velocity.x;
        e_vec.y = (v_squared - G * M_sun / r) * r_vec.y - r_dot_v * body->global_velocity.y;
        e_vec.z = (v_squared - G * M_sun / r) * r_vec.z - r_dot_v * body->global_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;
}