#include <stdint.h>
#include "rom/ets_sys.h"

/* GPIO register addresses */
volatile uint32_t *gpio_out    = (volatile uint32_t *)0x3FF44004;
volatile uint32_t *gpio_enable = (volatile uint32_t *)0x3FF44020;
volatile uint32_t *gpio_in     = (volatile uint32_t *)0x3FF4403C;

/* Pin masks */
uint32_t red    = (1 << 2);
uint32_t yellow = (1 << 4);
uint32_t green  = (1 << 5);
uint32_t button = (1 << 12);

uint32_t all_leds;

/*
  Delay while checking button every 50ms

  If button pressed:
  - Turn Yellow ON
  - Hold for 5 seconds
*/
void smart_delay(uint32_t total_us)
{
    uint32_t step = 50000; // 50 ms

    for (uint32_t t = 0; t < total_us; t += step)
    {
        // Check button during delay
        if ((*gpio_in & button) != 0)
        {
            // Yellow only
            *gpio_out = (*gpio_out & ~all_leds) | yellow;

            // Hold yellow 5 sec
            ets_delay_us(5000000);

            return;
        }

        ets_delay_us(step);
    }
}

void app_main(void)
{
    // All traffic LEDs together
    all_leds = red | yellow | green;

    // Set LED pins as outputs
    *gpio_enable |= all_leds;

    while (1)
    {
        // GREEN light
        *gpio_out = (*gpio_out & ~all_leds) | green;
        smart_delay(3000000);

        // YELLOW light
        *gpio_out = (*gpio_out & ~all_leds) | yellow;
        smart_delay(1000000);

        // RED light
        *gpio_out = (*gpio_out & ~all_leds) | red;
        smart_delay(3000000);
    }
}