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

double calculate_kinetic_energy(CelestialBody* body) {
    double v_squared = body->global_velocity.x * body->global_velocity.x +
                       body->global_velocity.y * body->global_velocity.y +
                       body->global_velocity.z * body->global_velocity.z;
    return 0.5 * body->mass * v_squared;
}

double calculate_potential_energy_pair(CelestialBody* body1, CelestialBody* body2) {
    double distance = vec3_distance(body1->global_position, body2->global_position);
    if (distance < 1.0) distance = 1.0;
    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.angular_position = atan2(relative_pos.y, relative_pos.x);

    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(int body_index) {
    OrbitTracker* tracker = (OrbitTracker*)malloc(sizeof(OrbitTracker));
    tracker->body_index = body_index;
    tracker->initial_angle = 0.0;
    tracker->previous_angle = 0.0;
    tracker->quadrant_transitions = 0;
    tracker->orbit_completed = false;
    tracker->time_at_completion = 0.0;
    tracker->min_time_days = 100.0;
    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_with_min_time(int body_index, double min_time_days) {
    OrbitTracker* tracker = (OrbitTracker*)malloc(sizeof(OrbitTracker));
    tracker->body_index = body_index;
    tracker->initial_angle = 0.0;
    tracker->previous_angle = 0.0;
    tracker->quadrant_transitions = 0;
    tracker->orbit_completed = false;
    tracker->time_at_completion = 0.0;
    tracker->min_time_days = min_time_days;
    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_3d(int body_index, double min_time_days,
                                       double inclination, double lon_ascending_node,
                                       double argument_of_periapsis) {
    OrbitTracker* tracker = create_orbit_tracker_with_min_time(body_index, min_time_days);
    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;
}

void reset_orbit_tracker(OrbitTracker* tracker) {
    tracker->initial_angle = 0.0;
    tracker->previous_angle = 0.0;
    tracker->quadrant_transitions = 0;
    tracker->orbit_completed = false;
    tracker->time_at_completion = 0.0;
}

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->quadrant_transitions == 0) {
        tracker->initial_angle = current_angle;
        tracker->previous_angle = current_angle;
        tracker->quadrant_transitions = 1;
        return;
    }

    double angle_diff = current_angle - tracker->previous_angle;

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

    double total_rotation = current_angle - tracker->initial_angle;
    if (total_rotation < -M_PI) total_rotation += 2.0 * M_PI;
    if (total_rotation > M_PI) total_rotation -= 2.0 * M_PI;

    const double SECONDS_PER_DAY = 86400.0;
    double min_time_seconds = tracker->min_time_days * SECONDS_PER_DAY;

    if (tracker->quadrant_transitions >= 2 &&
        current_time > min_time_seconds &&
        fabs(total_rotation) < 0.05) {
        tracker->orbit_completed = true;
        tracker->time_at_completion = current_time;
    }

    tracker->previous_angle = current_angle;
}

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

bool compare_double(double a, double b, double tolerance) {
    return fabs(a - b) <= tolerance;
}

bool compare_vec3(Vec3 a, Vec3 b, double tolerance) {
    return fabs(a.x - b.x) <= tolerance &&
           fabs(a.y - b.y) <= tolerance &&
           fabs(a.z - b.z) <= tolerance;
}

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

    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++) {
        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);
    for (int i = 0; i < sim->craft_count; i++) {
        Spacecraft* s = &sim->spacecraft[i];
        double r = sqrt(s->local_position.x*s->local_position.x +
                       s->local_position.y*s->local_position.y +
                       s->local_position.z*s->local_position.z);
        double v = sqrt(s->local_velocity.x*s->local_velocity.x +
                       s->local_velocity.y*s->local_velocity.y +
                       s->local_velocity.z*s->local_velocity.z);
        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);
        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);
    for (int i = 0; i < sim->maneuver_count; i++) {
        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;
}