#include "renderer.h"
#include "raymath.h"
#include <cmath>
#include <cstdio>

// raygui implementation (header-only library)
#define RAYGUI_IMPLEMENTATION
#include "raygui.h"

// Forward declarations
static Vector3 sim_to_render(Vec3 pos, double scale);

// Initialize raylib window
void init_renderer(int width, int height, const char* title) {
    InitWindow(width, height, title);
    SetTargetFPS(60);
}

// Close raylib
void close_renderer() {
    CloseWindow();
}

// Setup the 3D camera
void setup_camera(RenderState* render_state) {
    render_state->camera.position = (Vector3){ 0.0f, 50.0f, 100.0f };
    render_state->camera.target = (Vector3){ 0.0f, 0.0f, 0.0f };
    render_state->camera.up = (Vector3){ 0.0f, 1.0f, 0.0f };
    render_state->camera.fovy = 45.0f;
    render_state->camera.projection = CAMERA_PERSPECTIVE;

    // Set scaling factors (same scale for distances and sizes)
    render_state->distance_scale = 1e-9;  // Meters to scaled units (1 unit = 1 billion meters)
    render_state->size_scale = 1e-9;      // Same scale for body sizes (minimum size still applies)
}

// Update camera with keyboard/mouse controls
void update_camera(RenderState* render_state, SimulationState* sim) {
    float angle_speed = 0.02f;

    // Handle camera follow state transitions
    if (render_state->camera_follow_body && !render_state->was_following_body) {
        // Just started following - store current offset
        if (render_state->selected_body_index >= 0 && render_state->selected_body_index < sim->body_count) {
            CelestialBody* body = &sim->bodies[render_state->selected_body_index];
            Vector3 body_pos = sim_to_render(body->position, render_state->distance_scale);
            render_state->camera.target = body_pos;
            render_state->camera_offset = Vector3Subtract(render_state->camera.position, body_pos);
        }
    }

    // Update target position when following
    if (render_state->camera_follow_body && render_state->selected_body_index >= 0 &&
        render_state->selected_body_index < sim->body_count) {
        CelestialBody* body = &sim->bodies[render_state->selected_body_index];
        Vector3 body_pos = sim_to_render(body->position, render_state->distance_scale);
        render_state->camera.target = body_pos;
        render_state->camera.position = Vector3Add(body_pos, render_state->camera_offset);
    }

    // Camera rotation using camera's up vector
    Vector3 to_camera = Vector3Subtract(render_state->camera.position, render_state->camera.target);
    float camera_distance = Vector3Length(to_camera);

    // Rotate around target using camera's up vector
    if (IsKeyDown(KEY_LEFT)) {
        Vector3 forward = Vector3Normalize(to_camera);

        // Rotate forward vector around camera's up axis (horizontal orbit)
        float cos_a = cosf(angle_speed);
        float sin_a = sinf(angle_speed);

        Vector3 new_forward = Vector3Add(
            Vector3Scale(forward, cos_a),
            Vector3Scale(Vector3CrossProduct(render_state->camera.up, forward), sin_a)
        );

        render_state->camera.position = Vector3Add(
            render_state->camera.target,
            Vector3Scale(new_forward, camera_distance)
        );

        if (render_state->camera_follow_body) {
            render_state->camera_offset = Vector3Subtract(
                render_state->camera.position,
                render_state->camera.target
            );
        }
    }
    if (IsKeyDown(KEY_RIGHT)) {
        Vector3 forward = Vector3Normalize(to_camera);

        // Rotate forward vector around camera's up axis (horizontal orbit)
        float cos_a = cosf(-angle_speed);
        float sin_a = sinf(-angle_speed);

        Vector3 new_forward = Vector3Add(
            Vector3Scale(forward, cos_a),
            Vector3Scale(Vector3CrossProduct(render_state->camera.up, forward), sin_a)
        );

        render_state->camera.position = Vector3Add(
            render_state->camera.target,
            Vector3Scale(new_forward, camera_distance)
        );

        if (render_state->camera_follow_body) {
            render_state->camera_offset = Vector3Subtract(
                render_state->camera.position,
                render_state->camera.target
            );
        }
    }

    // Zoom in/out with up/down keys
    if (IsKeyDown(KEY_UP) && camera_distance > 10.0f) {
        Vector3 direction = Vector3Normalize(Vector3Subtract(render_state->camera.target, render_state->camera.position));
        render_state->camera.position = Vector3Add(render_state->camera.position, Vector3Scale(direction, 2.0f));

        if (render_state->camera_follow_body) {
            render_state->camera_offset = Vector3Subtract(
                render_state->camera.position,
                render_state->camera.target
            );
        }
    }
    if (IsKeyDown(KEY_DOWN)) {
        Vector3 direction = Vector3Normalize(Vector3Subtract(render_state->camera.position, render_state->camera.target));
        render_state->camera.position = Vector3Add(render_state->camera.position, Vector3Scale(direction, 2.0f));

        if (render_state->camera_follow_body) {
            render_state->camera_offset = Vector3Subtract(
                render_state->camera.position,
                render_state->camera.target
            );
        }
    }

    // Store previous follow state
    render_state->was_following_body = render_state->camera_follow_body;
}

// Transform from simulation coordinates (XY plane) to render coordinates (XZ plane)
// Rotation matrix: 90 degrees around X-axis maps Y -> Z
Vector3 sim_to_render(Vec3 pos, double scale) {
    return (Vector3){
        (float)(pos.x * scale),
        (float)(-pos.z * scale),
        (float)(pos.y * scale)
    };
}



// Scale a radius for rendering (with minimum visible size)
float scale_radius(double radius, double scale) {
    float scaled = (float)(radius * scale);
    float min_radius = 0.5f;  // Minimum visible radius
    return (scaled > min_radius) ? scaled : min_radius;
}

// Render a single celestial body
void render_body(CelestialBody* body, RenderState* render_state) {
    Vector3 position = sim_to_render(body->position, render_state->distance_scale);
    float radius = scale_radius(body->radius, render_state->size_scale);

    Color color = {
        (unsigned char)(body->color[0] * 255),
        (unsigned char)(body->color[1] * 255),
        (unsigned char)(body->color[2] * 255),
        255
    };

    DrawSphereWires(position, radius, 16, 16, color);
}

struct OrbitalBasis {
    Vec3 periapsis_dir;
    Vec3 normal;
    Vec3 q_vec;
};

static OrbitalBasis calculate_orbital_basis(Vec3 r_vec, Vec3 velocity, Vec3 e_vec) {
    OrbitalBasis basis;
    basis.periapsis_dir = vec3_normalize(e_vec);

    Vec3 h_vec = {
        r_vec.y * velocity.z - r_vec.z * velocity.y,
        r_vec.z * velocity.x - r_vec.x * velocity.z,
        r_vec.x * velocity.y - r_vec.y * velocity.x
    };
    basis.normal = vec3_normalize(h_vec);

    basis.q_vec = {
        basis.normal.y * basis.periapsis_dir.z - basis.normal.z * basis.periapsis_dir.y,
        basis.normal.z * basis.periapsis_dir.x - basis.normal.x * basis.periapsis_dir.z,
        basis.normal.x * basis.periapsis_dir.y - basis.normal.y * basis.periapsis_dir.x
    };

    return basis;
}

static Vec3 orbital_to_cartesian(double x, double y, OrbitalBasis basis, Vec3 parent_pos) {
    return {
        basis.periapsis_dir.x * x + basis.q_vec.x * y + parent_pos.x,
        basis.periapsis_dir.y * x + basis.q_vec.y * y + parent_pos.y,
        basis.periapsis_dir.z * x + basis.q_vec.z * y + parent_pos.z
    };
}

static void draw_orbit_segment(double x1, double y1, double x2, double y2,
                               OrbitalBasis basis, Vec3 parent_pos,
                               RenderState* render_state, Color color) {
    Vec3 p1_sim = orbital_to_cartesian(x1, y1, basis, parent_pos);
    Vec3 p2_sim = orbital_to_cartesian(x2, y2, basis, parent_pos);

    Vector3 p1 = sim_to_render(p1_sim, render_state->distance_scale);
    Vector3 p2 = sim_to_render(p2_sim, render_state->distance_scale);

    DrawLine3D(p1, p2, color);
}

static void render_elliptical_orbit(double a, double e, OrbitalBasis basis,
                                   Vec3 parent_pos, RenderState* render_state, Color color) {
    double b = a * sqrt(1.0 - e * e);
    double c = a * e;

    int segments = 100;
    for (int i = 0; i < segments; i++) {
        float theta1 = (float)i / segments * 2.0f * PI;
        float theta2 = (float)(i + 1) / segments * 2.0f * PI;

        double x1 = a * cos(theta1) - c;
        double y1 = b * sin(theta1);
        double x2 = a * cos(theta2) - c;
        double y2 = b * sin(theta2);

        draw_orbit_segment(x1, y1, x2, y2, basis, parent_pos, render_state, color);
    }
}

static void render_hyperbolic_orbit(double p, double e, OrbitalBasis basis,
                                     Vec3 parent_pos, RenderState* render_state, Color color) {
    double max_true_anomaly = (e > 1.01) ? acos(-1.0 / e) * 0.95 : PI * 0.48;

    int segments = 60;
    for (int i = 0; i < segments; i++) {
        float theta1 = -max_true_anomaly + (float)i / segments * 2.0f * max_true_anomaly;
        float theta2 = -max_true_anomaly + (float)(i + 1) / segments * 2.0f * max_true_anomaly;

        double r1 = p / (1.0 + e * cos(theta1));
        double r2 = p / (1.0 + e * cos(theta2));

        double x1 = r1 * cos(theta1);
        double y1 = r1 * sin(theta1);
        double x2 = r2 * cos(theta2);
        double y2 = r2 * sin(theta2);

        draw_orbit_segment(x1, y1, x2, y2, basis, parent_pos, render_state, color);
    }
}

static void render_parabolic_orbit(double p, OrbitalBasis basis,
                                   Vec3 parent_pos, RenderState* render_state, Color color) {
    double max_true_anomaly = PI * 0.95;

    int segments = 80;
    for (int i = 0; i < segments; i++) {
        float theta1 = -max_true_anomaly + (float)i / segments * 2.0f * max_true_anomaly;
        float theta2 = -max_true_anomaly + (float)(i + 1) / segments * 2.0f * max_true_anomaly;

        double r1 = p / (1.0 + cos(theta1));
        double r2 = p / (1.0 + cos(theta2));

        double x1 = r1 * cos(theta1);
        double y1 = r1 * sin(theta1);
        double x2 = r2 * cos(theta2);
        double y2 = r2 * sin(theta2);

        draw_orbit_segment(x1, y1, x2, y2, basis, parent_pos, render_state, color);
    }
}

// Render orbit path for a body
void render_orbit(CelestialBody* body, CelestialBody* parent, RenderState* render_state) {
    if (body->parent_index == -1 || parent == NULL) {
        return;
    }

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

    if (r < 1.0) return;

    double mu = G * parent->mass;
    double specific_energy = (v * v) / 2.0 - mu / r;

    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 = {
        (v_squared - mu / r) * r_vec.x - r_dot_v * body->velocity.x,
        (v_squared - mu / r) * r_vec.y - r_dot_v * body->velocity.y,
        (v_squared - mu / r) * r_vec.z - r_dot_v * body->velocity.z
    };
    double e = vec3_magnitude(e_vec) / mu;

    Color orbit_color = {
        (unsigned char)(body->color[0] * 128),
        (unsigned char)(body->color[1] * 128),
        (unsigned char)(body->color[2] * 128),
        128
    };

    OrbitalBasis basis = calculate_orbital_basis(r_vec, body->velocity, e_vec);

    if (e < 0.98) {
        double a = -mu / (2.0 * specific_energy);
        if (a <= 0.0) return;

        render_elliptical_orbit(a, e, basis, parent->position, render_state, orbit_color);
    } else if (e > 1.02) {
        double a = mu / (2.0 * (-specific_energy));
        double p = a * (1.0 - e * e);

        if (p <= 0.0) return;

        render_hyperbolic_orbit(p, e, basis, parent->position, render_state, orbit_color);
    } else {
        Vec3 h_vec = {
            r_vec.y * body->velocity.z - r_vec.z * body->velocity.y,
            r_vec.z * body->velocity.x - r_vec.x * body->velocity.z,
            r_vec.x * body->velocity.y - r_vec.y * body->velocity.x
        };
        double h_squared = h_vec.x * h_vec.x + h_vec.y * h_vec.y + h_vec.z * h_vec.z;
        double p = h_squared / mu;

        if (p <= 0.0) return;

        render_parabolic_orbit(p, basis, parent->position, render_state, orbit_color);
    }
}

// Render the entire simulation
void render_simulation(SimulationState* sim, RenderState* render_state) {
    BeginDrawing();
    ClearBackground(BLACK);

    BeginMode3D(render_state->camera);

    // Draw a subtle reference grid
    for (int i = -50; i <= 50; i++) {
        if (i == 0) continue;
        DrawLine3D((Vector3){(float)i * 10.0f, 0, -500.0f},
                   (Vector3){(float)i * 10.0f, 0, 500.0f},
                   (Color){20, 20, 20, 255});
        DrawLine3D((Vector3){-500.0f, 0, (float)i * 10.0f},
                   (Vector3){500.0f, 0, (float)i * 10.0f},
                   (Color){20, 20, 20, 255});
    }
    DrawLine3D((Vector3){0, 0, -500.0f}, (Vector3){0, 0, 500.0f}, (Color){40, 40, 40, 255});
    DrawLine3D((Vector3){-500.0f, 0, 0}, (Vector3){500.0f, 0, 0}, (Color){40, 40, 40, 255});

    // Render orbit paths first
    for (int i = 0; i < sim->body_count; i++) {
        CelestialBody* body = &sim->bodies[i];
        if (body->parent_index >= 0 && body->parent_index < sim->body_count) {
            CelestialBody* parent = &sim->bodies[body->parent_index];
            render_orbit(body, parent, render_state);
        }
    }

    // Render all bodies
    for (int i = 0; i < sim->body_count; i++) {
        render_body(&sim->bodies[i], render_state);
    }

    EndMode3D();

    // Render UI panels (always shown)
    render_info(sim, "solar_system.txt");
    render_body_list_ui(sim, render_state);
    render_body_info_ui(sim, render_state);

    EndDrawing();
}

// Render simulation information overlay
void render_info(SimulationState* sim, const char* config_name) {
    // Panel dimensions
    int panel_width = 200;
    int panel_height = 280;
    int panel_x = 10;
    int panel_y = GetScreenHeight() - panel_height - 10;

    Rectangle panel_bounds = { (float)panel_x, (float)panel_y, (float)panel_width, (float)panel_height };

    // Draw window
    if (GuiWindowBox(panel_bounds, "Info")) {
        // Window box close button doesn't close this panel
    }

    // Prepare info text
    char info_text[1024];
    char temp_buffer[256];
    int offset = 0;

    snprintf(info_text, sizeof(info_text), "Orbital Mechanics Simulation\n\n");

    // Simulation time (in days)
    double days = sim->time / 86400.0;  // seconds to days
    snprintf(temp_buffer, sizeof(temp_buffer), "Time: %.2f days\n", days);
    strcat(info_text, temp_buffer);

    // Body count
    snprintf(temp_buffer, sizeof(temp_buffer), "Bodies: %d\n", sim->body_count);
    strcat(info_text, temp_buffer);

    // FPS
    snprintf(temp_buffer, sizeof(temp_buffer), "FPS: %d\n\n", GetFPS());
    strcat(info_text, temp_buffer);

    // Controls
    strcat(info_text, "Controls:\n");
    strcat(info_text, "  Arrows: Rotate/Zoom\n");
    strcat(info_text, "  Space: Pause/Resume\n");
    strcat(info_text, "  +/-: Speed control\n");
    strcat(info_text, "  ESC: Quit\n");

    // Draw info text (inside window, leave space for title bar)
    Rectangle text_bounds = {
        (float)panel_x + 8,
        (float)panel_y + 25,
        (float)panel_width - 16,
        (float)panel_height - 29
    };

    GuiLabel(text_bounds, info_text);
}

// Render body list UI panel
void render_body_list_ui(SimulationState* sim, RenderState* render_state) {
    // Panel dimensions
    int panel_width = 200;
    int panel_height = 400;
    int panel_x = 10;
    int panel_y = 10;

    Rectangle panel_bounds = { (float)panel_x, (float)panel_y, (float)panel_width, (float)panel_height };

    // Create body list string for raygui (NULL-terminated)
    int buffer_size = sim->body_count * (64 + 1) + 1;
    char* body_list_text = (char*)malloc(buffer_size);
    if (!body_list_text) return;

    body_list_text[0] = '\0';
    int offset = 0;

    for (int i = 0; i < sim->body_count; i++) {
        if (i > 0) {
            offset += snprintf(body_list_text + offset, buffer_size - offset, ";");
        }
        offset += snprintf(body_list_text + offset, buffer_size - offset, "%s", sim->bodies[i].name);
    }

    // Draw window (no close button - panels always shown)
    GuiWindowBox(panel_bounds, "Bodies");

    // Draw list view (inside window, leave space for title bar)
    Rectangle list_bounds = {
        (float)panel_x + 4,
        (float)panel_y + 25,
        (float)panel_width - 8,
        (float)panel_height - 29
    };

    int previous_active = render_state->body_list_active;
    GuiListView(list_bounds, body_list_text, &render_state->body_list_scroll, &render_state->body_list_active);

    // Check if a body was selected (active changed from -1 or previous value)
    if (render_state->body_list_active >= 0 && render_state->body_list_active < sim->body_count &&
        render_state->body_list_active != previous_active) {
        render_state->selected_body_index = render_state->body_list_active;
        // Enable camera follow when body is selected
        render_state->camera_follow_body = true;
        printf("Camera follow enabled for: %s\n", sim->bodies[render_state->selected_body_index].name);
    }

    free(body_list_text);
}

// Render body information UI panel
void render_body_info_ui(SimulationState* sim, RenderState* render_state) {
    if (render_state->selected_body_index < 0 || render_state->selected_body_index >= sim->body_count) {
        // No body selected - render empty panel
        int panel_width = 250;
        int panel_height = 300;
        int panel_x = GetScreenWidth() - panel_width - 10;
        int panel_y = 10;

        Rectangle panel_bounds = { (float)panel_x, (float)panel_y, (float)panel_width, (float)panel_height };
        GuiWindowBox(panel_bounds, "Info");
        return;
    }

    // Panel dimensions
    int panel_width = 250;
    int panel_height = 300;
    int panel_x = GetScreenWidth() - panel_width - 10;
    int panel_y = 10;

    Rectangle panel_bounds = { (float)panel_x, (float)panel_y, (float)panel_width, (float)panel_height };

    CelestialBody* body = &sim->bodies[render_state->selected_body_index];

    // Draw window (no close button - panels always shown)
    GuiWindowBox(panel_bounds, body->name);

    // Prepare info text
    char info_text[1024];
    char temp_buffer[256];

    snprintf(info_text, sizeof(info_text), "Name: %s", body->name);

    snprintf(temp_buffer, sizeof(temp_buffer), "Mass: %.2e kg", body->mass);
    strcat(info_text, "\n");
    strcat(info_text, temp_buffer);

    snprintf(temp_buffer, sizeof(temp_buffer), "Radius: %.2e m", body->radius);
    strcat(info_text, "\n");
    strcat(info_text, temp_buffer);

    // Position
    snprintf(temp_buffer, sizeof(temp_buffer), "Position: (%.2e, %.2e, %.2e)",
             body->position.x, body->position.y, body->position.z);
    strcat(info_text, "\n");
    strcat(info_text, temp_buffer);

    // Velocity
    double vel_mag = vec3_magnitude(body->velocity);
    snprintf(temp_buffer, sizeof(temp_buffer), "Velocity: %.2f m/s", vel_mag);
    strcat(info_text, "\n");
    strcat(info_text, temp_buffer);

    // Orbital elements
    snprintf(temp_buffer, sizeof(temp_buffer), "Eccentricity: %.3f", body->eccentricity);
    strcat(info_text, "\n");
    strcat(info_text, temp_buffer);

    snprintf(temp_buffer, sizeof(temp_buffer), "Semi-major axis: %.2e m", body->semi_major_axis);
    strcat(info_text, "\n");
    strcat(info_text, temp_buffer);

    // Parent body
    if (body->parent_index >= 0 && body->parent_index < sim->body_count) {
        snprintf(temp_buffer, sizeof(temp_buffer), "Parent: %s", sim->bodies[body->parent_index].name);
    } else {
        snprintf(temp_buffer, sizeof(temp_buffer), "Parent: None");
    }
    strcat(info_text, "\n");
    strcat(info_text, temp_buffer);

    // SOI
    snprintf(temp_buffer, sizeof(temp_buffer), "SOI radius: %.2e m", body->soi_radius);
    strcat(info_text, "\n");
    strcat(info_text, temp_buffer);

    // Draw info text (inside window, leave space for title bar)
    Rectangle text_bounds = {
        (float)panel_x + 8,
        (float)panel_y + 25,
        (float)panel_width - 16,
        (float)panel_height - 29
    };

    GuiLabel(text_bounds, info_text);
}