Onzabr

Добавьте псевдоэлемент after, который будет заполнять свободное пространство:

.container::after {
  content: "";
  flex: auto;
}

Если состояний всего 2, как в вашем случае, то лучше использовать один класс и проверять на его наличие (если класс active присутствует - лампа горит, если нет, соответственно, не горит). Для переключения лучше использовать classList.toggle
element.classList.toggle("active");