#include <avr/sleep.h>
#include <avr/interrupt.h>

#define BUTTON_PIN 2  // PB2
#define LED1_PIN 1    // PB1
#define LED2_PIN 3    // PB3
#define LED_ON_DURATION 1000
#define BLINK_RATE 200

volatile bool wakeFlag = false;
uint32_t switchLedOffNow = 0;

void sleepNow() {
  // Prepare for sleep
  PORTB &= 0b11100000; // PB4-PB0 -> LOW
  DDRB &= ~(bit(BUTTON_PIN) | bit(LED1_PIN) | bit(LED2_PIN)); //  These pins as INPUTs in one go
  set_sleep_mode(SLEEP_MODE_PWR_DOWN);
  sleep_enable();
  sleep_mode(); // now sleeping until interrupt
  sleep_disable();
}

ISR(PCINT0_vect) {
  // Button pin changed state
  wakeFlag = true;
}

void blinkIfAlive() {
  static uint32_t nextBlink = 0;
  if (millis() > nextBlink) {
    digitalWrite(LED2_PIN, !digitalRead(LED2_PIN));
    nextBlink += BLINK_RATE;  // #define BLINK_RATE 200 or whatever
  }
}

void setup() {
  // pinMode(BUTTON_PIN, INPUT);  // button wired to GND
  // pinMode(LED1_PIN, OUTPUT);
  // pinMode(LED2_PIN, OUTPUT);

  // Disable ADC - saves about 324.5uA in sleep mode!
  ADCSRA &= ~bit(ADEN);  // Disable ADC
  ACSR |= bit(ACD);      // Disable analog comparator
  // Set up Pin Change Interrupt on PB2 (Pin 2)
  GIMSK |= (1 << PCIE);               // Enable Pin Change Interrupts
  PCMSK |= (1 << BUTTON_PIN);         // Enable interrupt on BUTTON_PIN
  sleepNow();
}

void loop() {
  if (wakeFlag) {
    wakeFlag = false;
    // pinMode(BUTTON_PIN, OUTPUT);
    // pinMode(LED, OUTPUT);
    // pinMode(LED2, OUTPUT);
    DDRB |= bit(BUTTON_PIN) | bit(LED1_PIN) | bit(LED2_PIN); //  These pins as OUTPUTs in one go
    // LED1 ON for 1s
    digitalWrite(LED1_PIN, HIGH);
    switchLedOffNow = millis() + LED_ON_DURATION;
  }

  if (digitalRead(LED1_PIN) && millis() > switchLedOffNow) {
    // Do whatever has to be done

      digitalWrite(LED1_PIN, LOW);
      sleepNow();
  }

  blinkIfAlive();

}