#include <avr/io.h>
#include <avr/sleep.h>
#include <avr/eeprom.h>
#include <avr/wdt.h>

#define BUTTON_PIN PB4
#define PORT_PULLUP_MASK _BV(PB4)
#define PORT_OUTPUT_MASK 0

#define BTN_DEBOUNCE_DELAY 10

struct OutMask {
  uint8_t pullup;
  uint8_t output;
};

uint8_t cur_out_mask = 0;
volatile OutMask out_masks[3] = {
  { PORT_PULLUP_MASK | 4, PORT_OUTPUT_MASK },
  { PORT_PULLUP_MASK | 2, PORT_OUTPUT_MASK },
  { PORT_PULLUP_MASK | 1, PORT_OUTPUT_MASK },
};

void set_led_masks(uint8_t bits)
{
  out_masks[0].output = PORT_OUTPUT_MASK | (bits & 3);
  out_masks[1].output = PORT_OUTPUT_MASK | ((bits >> 2) & 1) | ((bits >> 1) & 4);
  out_masks[2].output = PORT_OUTPUT_MASK | ((bits >> 3) & 6);
}

volatile bool btn_down = false, btn_pressed = false, btn_released = false;
volatile uint8_t btn_up_timer = 0;

ISR(TIM0_COMPA_vect)
{
  DDRB = PORT_OUTPUT_MASK;
}

ISR(TIMER0_OVF_vect)
{
  DDRB = PORT_OUTPUT_MASK;
  volatile OutMask &p = out_masks[cur_out_mask];
  PORTB = p.pullup;
  DDRB = p.output;
  if (++cur_out_mask >= 3)
    cur_out_mask = 0;

  uint8_t up = btn_up_timer;
  if (up) {
    if (++up >= BTN_DEBOUNCE_DELAY)
    {
      up = 0;
      btn_down = false;
      btn_released = true;
    }
    btn_up_timer = up;
  }
}

volatile bool step = false;

ISR(WDT_vect)
{
  //step = true;
}

ISR(PCINT0_vect)
{
  if (PINB & _BV(BUTTON_PIN))
  {
    if (btn_down) {
      btn_up_timer = 1;
    }
  }
  else
  {
    if (!btn_down)
    {
      btn_pressed = true;
      btn_down = true;
    }
    btn_up_timer = 0;
  }
}

int main(void)
{
  // Clock divider = 64 (125 kHz)
  CLKPR = _BV(CLKPCE);
  CLKPR = 6;

  // Set up the watchdog timer
  wdt_reset();
  MCUSR = 0x00; // Clear WDRF in MCUSR
  WDTCR = _BV(WDCE) | _BV(WDE);
  WDTCR = _BV(WDIE) | _BV(WDP2) | _BV(WDP0); // 0.5 sec interrupt

  // Set up the timer to fire ~1024 times a second
  TCCR0A = 0; // Normal mode
  TCCR0B = _BV(CS00); // clock frequency
  OCR0A = 128; // Output compare
  OCR0B = 0;
  TCNT0 = 0; // Set counter 0 to zero
  TIMSK = _BV(OCIE0A) | _BV(TOIE0); // Enable overflow and counter A interrupts

  // Set up button change interrupt
  GIMSK |= _BV(PCIE); 
  PCMSK |= _BV(PCINT4);

  // First 4 pins to output, PB4 pull up
  PORTB = PORT_PULLUP_MASK;
  DDRB = PORT_OUTPUT_MASK;

  set_led_masks(0b101010);

  sei(); //Enable global interrupts
  
  // Sleep in the main loop - the work happens in the timer interrupt
  MCUCR = _BV(SE);

  uint8_t idx = 0;
  step = true;

  for (;;)
  {
    if (btn_pressed || btn_released)
    {
      step = true;
      btn_pressed = btn_released = false;
    }

    if (step)
    {
      set_led_masks(1 << idx);
      if (++idx >= 6)
        idx = 0;
      step = false;
    }

    sleep_cpu();
  }
}