// src/Stage.cpp

#include "Stage.h"
#include "Motor.h"
#include "Airframe.h"
#include "FinSet.h"
#include "RecoverySystem.h"
#include "Environment.h"
#include "FlightState.h"

#include <numeric> // for std::accumulate
#include <algorithm> // for std::all_of

// Constructor
Stage::Stage(const std::string& name)
    : name_(name),
      totalMass_(0.0),
      totalPropellantMass_(0.0),
      separationTriggered_(false),
      recoveryDeployed_(false)
{}

// Public Methods

void Stage::addMotor(std::unique_ptr<Motor> motor) {
    motors_.emplace_back(std::move(motor));
}

void Stage::setAirframe(std::unique_ptr<Airframe> airframe) {
    airframe_ = std::move(airframe);
}

void Stage::setFinSet(std::unique_ptr<FinSet> finSet) {
    finSet_ = std::move(finSet);
}

void Stage::setRecoverySystem(std::unique_ptr<RecoverySystem> recovery) {
    recoverySystem_ = std::move(recovery);
}

double Stage::getTotalMass() const {
    return totalMass_;
}

double Stage::getTotalPropellantMass() const {
    return totalPropellantMass_;
}

void Stage::updateMassProperties() {
    totalMass_ = 0.0;
    totalPropellantMass_ = 0.0;

    // Add airframe dry mass
    if (airframe_) {
        totalMass_ += airframe_->getDryMass();
    }

    // Add motor masses
    for (const auto& motor : motors_) {
        totalMass_ += motor->getRemainingPropellantMass();
        totalPropellantMass_ += motor->getRemainingPropellantMass();
    }

    // Add fin mass
    if (finSet_) {
        totalMass_ += finSet_->calculateMass();
    }
}

void Stage::prepareForFlight(const Environment& env) {
    // Ignite all motors immediately at start (for simplicity now)
    for (auto& motor : motors_) {
        motor->ignite();
    }
}

void Stage::applyFlightState(const FlightState& state) {
    // Placeholder for future advanced logic (flex, vibration)
    // Currently nothing needed at 3-DoF level
}

bool Stage::checkSeparationEvent() const {
    // Simple rule: if all motors burned out, stage is done
    bool allBurnedOut = std::all_of(motors_.begin(), motors_.end(),
                                    [](const std::unique_ptr<Motor>& motor) {
                                        return motor->isBurnedOut();
                                    });
    return allBurnedOut;
}

bool Stage::checkRecoveryEvent() const {
    if (!recoverySystem_ || recoverySystem_->isDeployed()) {
        return false;
    }
    // RecoverySystem handles checking altitude, velocity, etc.
    // Here we just provide example usage
    return false; // Actual deployment logic will be done in FlightSimulator::handleEvents()
}

const std::string& Stage::getName() const {
    return name_;
}

double Stage::getAirframeLength() const {
    if (airframe_) {
        return airframe_->getLength();
    }
    return 0.0;
}

double Stage::calculateNormalForceCoefficient() const {
    if (finSet_) {
        return finSet_->calculateNormalForceCoefficient();
    }
    return 0.0;
}

std::array<double, 3> Stage::getTotalThrust() const {
    std::array<double, 3> totalThrust = {0.0, 0.0, 0.0};

    for (const auto& motor : motors_) {
        // For now, assume thrust always points along +Z axis
        totalThrust[2] += motor->getCurrentThrust();
    }

    return totalThrust;
}