javascript breaks menu behaviour

I have nav bar that shrinks depending on the screen width and the button appears that toggles the dropdown menu. When I triggered the dropdown menu using checkbox (pure CSS), it worked perfectly. But after I started to use js instead of checkbox, it started to break the nav bar behaviour. When I shrink the window then press the button to toggle the dropdown menu and then expand window back, the nav bar doesn’t revert back to it’s “default” state. How to make the nav bar revert back to its state when the screen width increases? Is there any ability to apply style ONLY to media query css style but not to the main style? Here’s my code:

const menuButton = document.querySelector('.menuButton')
const menuDisplay = document.querySelector('nav ul')

menuButton.addEventListener('click', () => {
  if (menuDisplay.style.display == 'none') {
    menuDisplay.style.display = 'initial'
    setTimeout(() => {
      menuDisplay.style.height = '100vh'
    }, 0)
  } else {
    setTimeout(() => {
      menuDisplay.style.display = 'none'
    }, 300)
    menuDisplay.style.height = '0'
  }
})
header {
  background-color: rgba(24, 24, 24, 0.95);
  position: fixed;
  width: 100%;
  top: 0;
  z-index: 10;
}

.menuButton {
  display: none;
  position: fixed;
  right: 0;
  color: white;
  font-size: 1.5em;
  line-height: 65px;
  margin: 10 30px;
  cursor: pointer;
}

.logo {
  color: white;
  font-size: 30px;
  font-weight: bold;
  font-family: 'Fredoka One', sans-serif;
  line-height: 65px;
}

nav {
  padding: 10px 5%;
  display: flex;
  max-width: 900px;
  margin: auto;
  justify-content: space-between;
  align-items: center;
  -moz-user-select: none;
  -webkit-user-select: none;
  user-select: none;
}

nav ul {
  display: flex;
  list-style: none;
  font-weight: bold;
}

nav ul li a {
  display: block;
  text-align: center;
  padding: 14px 16px;
  margin: 5px;
}

.orderButton {
  border-radius: 5px;
  background-color: #e8e8e8;
  color: black;
  border: 2px solid #e8e8e8;
  text-align: center;
  padding: 14px 16px;
  margin: 5px;
  cursor: pointer;
  transition: all.3s ease;
  -webkit-transition: all.3s ease;
  -o-transition: all.3s ease;
  -moz-transition: all.3s ease;
}

.orderButton:hover {
  background-color: rgba(0, 0, 0, 0);
  color: white;
}

nav ul li .active {
  color: #999999;
}

@media (max-width: 600px) {
  nav ul {
    width: 100%;
    text-align: center;
    transition: margin-top 0.3s ease-in-out, height 0.3s ease-in-out;
    -webkit-transition: margin-top 0.3s ease-in-out, height 0.3s ease-in-out;
    -moz-transition: margin-top 0.3s ease-in-out, height 0.3s ease-in-out;
    flex-direction: column;
    overflow: hidden;
    height: 0;
    display: none;
  }
  nav ul li {
    margin: 20px;
  }
  nav {
    flex-direction: column;
  }
  .logo {
    text-align: center;
  }
  .menuButton {
    display: block;
  }
}
<header>
  <div class="menuButton">&#9776;</div>
  <nav>
    <a href="index.html">
      <h1 class="logo">Cite</h1>
    </a>
    <ul>
      <li>
        <div class="orderButton">Заказать</div>
      </li>
      <div class="orderForm">
        <div class="form-wrap">
          <h2>Данные для связи</h2>
          <form method="post" action="order.php">

            <p>Имя</p>
            <input type="text" name="name" value="" placeholder="Имя" />

            <p>Телефон</p>
            <input type="tel" name="phone" value="" placeholder="Телефон" />

            <p>Почта</p>
            <input type="email" name="email" value="" placeholder="E-mail" />

            <p>Комментарий</p>
            <textarea name="other" placeholder="Комментарий к заявке"></textarea>
          </form>
        </div>
      </div>
      <li><a href="about.html">О нас</a></li>
      <li><a href="contact.html">Контакты</a></li>
    </ul>
  </nav>

  <script src="js/script.js" type="text/javascript"></script>

</header>