// tests/test_stage.cpp

#include "vendor/catch_amalgamated.hpp"
#include "Stage.h"
#include "Motor.h"
#include "Airframe.h"
#include "FinSet.h"
#include "RecoverySystem.h"
#include "Environment.h"

// Bring Approx into scope
using Catch::Approx;

TEST_CASE("Stage basic behavior", "[stage]") {
    Stage stage("TestStage");

    SECTION("Stage initially has zero mass and no components") {
        REQUIRE(stage.getTotalMass() == Approx(0.0));
        REQUIRE(stage.getTotalPropellantMass() == Approx(0.0));
        REQUIRE(stage.getAirframeLength() == Approx(0.0));
        REQUIRE(stage.calculateNormalForceCoefficient() == Approx(0.0));
    }

    SECTION("Adding a motor updates propellant mass") {
        auto motor = std::make_unique<Motor>("TestMotor", 0.05, 50.0); // 50 grams
        stage.addMotor(std::move(motor));

        stage.updateMassProperties();

        REQUIRE(stage.getTotalPropellantMass() > 0.0);
        REQUIRE(stage.getTotalMass() > 0.0);
    }

    SECTION("Adding an airframe updates dry mass and length") {
        auto airframe = std::make_unique<Airframe>("BodyTube",
            1.2, // length meters
            0.1, // diameter meters
            1.5, // dry mass kg
            0.75 // drag coefficient
        );
        double expectedLength = airframe->getLength();
        double expectedMass = airframe->getDryMass();

        stage.setAirframe(std::move(airframe));
        stage.updateMassProperties();

        REQUIRE(stage.getAirframeLength() == Approx(expectedLength));
        REQUIRE(stage.getTotalMass() >= Approx(expectedMass));
    }

    SECTION("Adding fins provides stability contribution") {
        auto fins = std::make_unique<FinSet>("BasicFins",
            3,  // 3 fins
            0.15, // root chord
            0.10, // tip chord
            0.10, // span
            0.05, // sweep length
            0.003, // thickness
            500.0  // material density
        );

        stage.setFinSet(std::move(fins));

        REQUIRE(stage.calculateNormalForceCoefficient() > 0.0);
    }

    SECTION("Stage thrust calculation sums motor thrusts") {
        auto motor1 = std::make_unique<Motor>("Motor1", 0.03, 30.0);
        auto motor2 = std::make_unique<Motor>("Motor2", 0.02, 20.0);

        motor1->addThrustDataPoint(0.0, 0.0);
        motor1->addThrustDataPoint(0.5, 50.0);
        motor1->addThrustDataPoint(1.0, 0.0);

        motor2->addThrustDataPoint(0.0, 0.0);
        motor2->addThrustDataPoint(0.5, 30.0);
        motor2->addThrustDataPoint(1.0, 0.0);

        stage.addMotor(std::move(motor1));
        stage.addMotor(std::move(motor2));

        stage.updateMassProperties();
        stage.prepareForFlight(*(new Environment(101325.0, 288.15, 0.0))); // Dummy environment

        // Simulate a half second into flight
        for (auto& motor : stage.getMotorsForTesting()) {
            motor->update(0.5);
        }

        auto thrustVector = stage.getTotalThrust();
        REQUIRE(thrustVector[2] > 0.0);
    }
}
