#ifndef STAGE_H
#define STAGE_H

#include <string>
#include <vector>
#include <memory>

class Motor;
class Airframe;
class FinSet;
class RecoverySystem;
class Environment;
class FlightState;

/**
 * @brief Represents a single stage of a (potentially) multi-stage rocket.
 *
 * A Stage contains structural components, motors, recovery systems,
 * and manages its own mass properties and separation events.
 */
class Stage {
public:
    /**
     * @brief Constructs a new Stage with a given name.
     * @param name The name of the stage.
     */
    Stage(const std::string& name);

    /**
     * @brief Default destructor.
     */
    ~Stage() = default;

    /**
     * @brief Adds a motor to the stage.
     * @param motor A unique_ptr to the Motor object.
     */
    void addMotor(std::unique_ptr<Motor> motor);

    /**
     * @brief Sets the airframe for the stage.
     * @param airframe A unique_ptr to the Airframe object.
     */
    void setAirframe(std::unique_ptr<Airframe> airframe);

    /**
     * @brief Sets the fin set for the stage.
     * @param finSet A unique_ptr to the FinSet object.
     */
    void setFinSet(std::unique_ptr<FinSet> finSet);

    /**
     * @brief Sets the recovery system for the stage.
     * @param recovery A unique_ptr to the RecoverySystem object.
     */
    void setRecoverySystem(std::unique_ptr<RecoverySystem> recovery);

    /**
     * @brief Returns the total mass of the stage (structure + motors + payload).
     * @return Total mass in kilograms.
     */
    double getTotalMass() const;

    /**
     * @brief Returns the current total propellant mass of all motors in the stage.
     * @return Propellant mass in kilograms.
     */
    double getTotalPropellantMass() const;

    /**
     * @brief Updates mass properties (e.g., after motor burn or separation).
     */
    void updateMassProperties();

    /**
     * @brief Prepares the stage for flight simulation (e.g., motor ignition sequencing).
     * @param env Launch environment.
     */
    void prepareForFlight(const Environment& env);

    /**
     * @brief Updates the stage based on current flight conditions.
     * @param state Current flight state.
     */
    void applyFlightState(const FlightState& state);

    /**
     * @brief Checks if the stage should separate (e.g., after burnout or trigger).
     * @return True if ready to separate.
     */
    bool checkSeparationEvent() const;

    /**
     * @brief Checks if the recovery system should deploy (e.g., apogee, velocity triggers).
     * @return True if recovery deployment condition met.
     */
    bool checkRecoveryEvent() const;

    /**
     * @brief Gets the length of the airframe.
     * @return Length in meters (0 if no airframe assigned).
     */
    double getAirframeLength() const;

    /**
     * @brief Calculates the normal force coefficient contribution from the stage's fins.
     * @return Normal force coefficient (dimensionless).
     */
    double calculateNormalForceCoefficient() const;

    /**
     * @brief Gets the name of the stage.
     * @return Stage name as a constant reference.
     */
    const std::string& getName() const;

    /**
     * @brief Calculates the total thrust vector produced by all motors in this stage.
     * @return Thrust vector (Newton) in body-fixed frame (Z-forward).
     */
    std::array<double, 3> getTotalThrust() const;

    // Only for testing access!
    const std::vector<std::unique_ptr<Motor>>& getMotorsForTesting() const { return motors_; }


private:
    std::string name_; ///< Name of the stage.

    // Core Components
    std::vector<std::unique_ptr<Motor>> motors_; ///< List of motors in this stage.
    std::unique_ptr<Airframe> airframe_; ///< Structural body and aerodynamic surfaces.
    std::unique_ptr<FinSet> finSet_; ///< Fins for aerodynamic stability.
    std::unique_ptr<RecoverySystem> recoverySystem_; ///< Recovery deployment system.

    // Cached Mass Properties
    double totalMass_; ///< Total mass [kg].
    double totalPropellantMass_; ///< Remaining propellant [kg].

    // Event Flags
    bool separationTriggered_; ///< Flag indicating stage separation event.
    bool recoveryDeployed_; ///< Flag indicating recovery system deployment.
};

#endif // STAGE_H