#include "test_utilities.h"
#include <cstdlib>
#include <cmath>
#include <cstdio>

static const double MIN_DISTANCE_CLAMP = 1.0;

double calculate_kinetic_energy(CelestialBody* body) {
    Vec3 v = body->global_velocity;
    return 0.5 * body->mass * vec3_dot(v, v);
}

double calculate_potential_energy_pair(CelestialBody* body1, CelestialBody* body2) {
    double distance = vec3_distance(body1->global_position, body2->global_position);
    if (distance < MIN_DISTANCE_CLAMP) distance = MIN_DISTANCE_CLAMP;
    return -G * body1->mass * body2->mass / distance;
}

double calculate_system_total_energy(SimulationState* sim) {
    double kinetic = 0.0;
    double potential = 0.0;

    for (int i = 0; i < sim->body_count; i++) {
        kinetic += calculate_kinetic_energy(&sim->bodies[i]);

        for (int j = i + 1; j < sim->body_count; j++) {
            potential += calculate_potential_energy_pair(&sim->bodies[i], &sim->bodies[j]);
        }
    }

    return kinetic + potential;
}

OrbitalMetrics calculate_orbital_metrics(CelestialBody* body, CelestialBody* parent) {
    OrbitalMetrics metrics;

    Vec3 relative_pos = vec3_sub(body->global_position, parent->global_position);
    metrics.orbital_radius = vec3_magnitude(relative_pos);
    metrics.velocity_magnitude = vec3_magnitude(body->global_velocity);

    metrics.kinetic_energy = calculate_kinetic_energy(body);
    metrics.potential_energy = calculate_potential_energy_pair(body, parent);
    metrics.total_energy = metrics.kinetic_energy + metrics.potential_energy;

    return metrics;
}

OrbitTracker* create_orbit_tracker_with_min_time(int body_index, double min_time_seconds) {
    OrbitTracker* tracker = (OrbitTracker*)malloc(sizeof(OrbitTracker));
    tracker->body_index = body_index;
    tracker->initial_angle = 0.0;
    tracker->previous_angle = 0.0;
    tracker->accumulated_rotation = 0.0;
    tracker->wrap_count = 0;
    tracker->orbit_completed = false;
    tracker->time_at_completion = 0.0;
    tracker->min_time_seconds = min_time_seconds;
    tracker->inclination = 0.0;
    tracker->longitude_of_ascending_node = 0.0;
    tracker->argument_of_periapsis = 0.0;
    tracker->has_orbital_elements = false;
    return tracker;
}

OrbitTracker* create_orbit_tracker(int body_index) {
    return create_orbit_tracker_with_min_time(body_index, 86400.0);
}

OrbitTracker* create_orbit_tracker_3d(int body_index, double min_time_seconds,
                                       double inclination, double lon_ascending_node,
                                       double argument_of_periapsis) {
    OrbitTracker* tracker = create_orbit_tracker_with_min_time(body_index, min_time_seconds);
    tracker->inclination = inclination;
    tracker->longitude_of_ascending_node = lon_ascending_node;
    tracker->argument_of_periapsis = argument_of_periapsis;
    tracker->has_orbital_elements = true;
    return tracker;
}

static void reset_tracker_fields(OrbitTracker* tracker) {
    tracker->body_index = 0;
    tracker->initial_angle = 0.0;
    tracker->previous_angle = 0.0;
    tracker->accumulated_rotation = 0.0;
    tracker->wrap_count = 0;
    tracker->orbit_completed = false;
    tracker->time_at_completion = 0.0;
}

void reset_orbit_tracker(OrbitTracker* tracker) {
    reset_tracker_fields(tracker);
}

void update_orbit_tracker(OrbitTracker* tracker, CelestialBody* body, CelestialBody* parent, double current_time) {
    if (tracker->orbit_completed) return;

    Vec3 relative_pos = vec3_sub(body->global_position, parent->global_position);
    double current_angle;

    if (tracker->has_orbital_elements) {
        Mat3 rotation = mat3_rotation_orbital(tracker->argument_of_periapsis,
                                              tracker->inclination,
                                              tracker->longitude_of_ascending_node);
        // Transpose to get inverse rotation (back to orbital plane)
        Mat3 rotation_T = {rotation.m00, rotation.m10, rotation.m20,
                           rotation.m01, rotation.m11, rotation.m21,
                           rotation.m02, rotation.m12, rotation.m22};
        Vec3 pos_orbital = mat3_multiply_vec3(rotation_T, relative_pos);
        current_angle = atan2(pos_orbital.y, pos_orbital.x);
    } else {
        current_angle = atan2(relative_pos.y, relative_pos.x);
    }

    if (tracker->wrap_count == 0) {
        tracker->initial_angle = current_angle;
        tracker->previous_angle = current_angle;
        tracker->accumulated_rotation = 0.0;
        tracker->wrap_count = 1;
        return;
    }

    double angle_diff = current_angle - tracker->previous_angle;

    if (angle_diff > M_PI) {
        angle_diff -= 2.0 * M_PI;
        tracker->wrap_count++;
    }
    if (angle_diff < -M_PI) {
        angle_diff += 2.0 * M_PI;
        tracker->wrap_count++;
    }

    tracker->accumulated_rotation += angle_diff;

    if (current_time > tracker->min_time_seconds &&
        (tracker->accumulated_rotation >= 2.0 * M_PI ||
         tracker->accumulated_rotation <= -2.0 * M_PI)) {
        tracker->orbit_completed = true;
        tracker->time_at_completion = current_time;
    }

    tracker->previous_angle = current_angle;
}

void destroy_orbit_tracker(OrbitTracker* tracker) {
    free(tracker);
}

int dump_simulation_state(SimulationState* sim, const char* label,
                          char* buffer, int buffer_size) {
    int offset = 0;

    if (offset >= buffer_size) return offset;

    offset += snprintf(buffer + offset, buffer_size - offset,
                       "\n=== %s (t=%.0f s) ===\n", label, sim->time);

    offset += snprintf(buffer + offset, buffer_size - offset,
                       "Bodies (%d):\n", sim->body_count);
    for (int i = 0; i < sim->body_count; i++) {
        if (offset >= buffer_size) break;
        offset += snprintf(buffer + offset, buffer_size - offset,
                           "  [%d] %s: mass=%.2e kg\n",
                           i, sim->bodies[i].name, sim->bodies[i].mass);
    }

    offset += snprintf(buffer + offset, buffer_size - offset,
                       "Spacecraft (%d):\n", sim->craft_count);
    if (sim->spacecraft) {
        for (int i = 0; i < sim->craft_count; i++) {
            if (offset >= buffer_size) break;
            Spacecraft* s = &sim->spacecraft[i];
            double r = vec3_magnitude(s->local_position);
            double v = vec3_magnitude(s->local_velocity);
            offset += snprintf(buffer + offset, buffer_size - offset,
                               "  [%d] %s: r=%.1f v=%.1f nu=%.5f a=%.1f e=%.6f, omega=%.6f\n",
                               i, s->name, r, v,
                               s->orbit.true_anomaly,
                               s->orbit.semi_major_axis,
                               s->orbit.eccentricity,
                               s->orbit.argument_of_periapsis);
            if (offset >= buffer_size) break;
            offset += snprintf(buffer + offset, buffer_size - offset,
                               "       pos=(%.1f, %.1f, %.1f) vel=(%.1f, %.1f, %.1f)\n",
                               s->local_position.x, s->local_position.y, s->local_position.z,
                               s->local_velocity.x, s->local_velocity.y, s->local_velocity.z);
        }
    }

    offset += snprintf(buffer + offset, buffer_size - offset,
                       "Maneuvers (%d):\n", sim->maneuver_count);
    if (sim->maneuvers) {
        for (int i = 0; i < sim->maneuver_count; i++) {
            if (offset >= buffer_size) break;
            Maneuver* m = &sim->maneuvers[i];
            offset += snprintf(buffer + offset, buffer_size - offset,
                               "  [%d] %s: craft=%d dir=%d dv=%.4f trigger=%d val=%.2f exec=%d\n",
                               i, m->name, m->craft_index, m->direction, m->delta_v,
                               m->trigger_type, m->trigger_value, m->executed);
        }
    }

    return offset;
}