from pydantic import BaseModel

from .acs import AttitudeControlSystem
from .communications import CommunicationsSystem
from .data_generator import DataGeneration
from .power import PowerDraw
from .thermal import Heater


class SpacecraftBus(BaseModel):
    name: str = "Default Bus"
    power_draw: PowerDraw = PowerDraw()
    attitude_control: AttitudeControlSystem = AttitudeControlSystem()
    communications: CommunicationsSystem | None = None
    heater: Heater | None = None
    data_generation: DataGeneration = DataGeneration()

    def power(self, mode: int | None = None, in_eclipse: bool = False) -> float:
        """Get the power draw for the spacecraft bus in the given mode.

        Args:
            mode: Operational mode (None for nominal)
            in_eclipse: Whether spacecraft is in eclipse

        Returns:
            Total power draw in watts
        """
        base_power = self.power_draw.power(mode, in_eclipse=in_eclipse)
        heater_power = (
            self.heater.power(mode, in_eclipse=in_eclipse) if self.heater else 0.0
        )
        return base_power + heater_power

    def data_generated(self, duration_seconds: float) -> float:
        """Calculate data generated by the spacecraft bus over a given duration.

        Args:
            duration_seconds: Duration of data generation in seconds.

        Returns:
            Amount of data generated in Gigabits.
        """
        return self.data_generation.data_generated(duration_seconds)
