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

#define BUTTON_PIN  4  // PB4
#define LED1_PIN    1  // PB1
#define LED2_PIN    3  // PB3
#define ISR_FLAG    5  // PB5
#define LED_ON_DURATION 1000
#define BLINK_RATE 200

volatile bool wakeFlag = false;
uint32_t switchLedOffNow = 0;

void usiSetup() {
  // Configure USI for SPI Master mode
  USICR = (1 << USIWM0) | (1 << USICS1); // 3-wire mode, software clock
  DDRB |= (1 << PB0) | (1 << PB2); // Set DO (PB0) and USCK (PB2) as outputs
}

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

ISR(PCINT0_vect) {
  static uint8_t ISRCount = 0;
  // Button pin changed state
  GIMSK &= ~(1 << PCIE);  // Disable Pin Change Interrupt (optional, for safety)
  wakeFlag = true;
  digitalWrite(ISR_FLAG, HIGH);  // Not an OUTPUT yet but set HIGH anyway
  USIDR = ++ISRCount;
  USISR = (1 << USIOIF); // Clear counter flag}
}

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() {
  usiSetup();
  // 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
  sleepNow();
}

void loop() {
  if (wakeFlag) {
    wakeFlag = false;
    // PCMSK &= ~(1 << BUTTON_PIN);         // Disable interrupt on BUTTON_PIN
    // pinMode(BUTTON_PIN, OUTPUT);
    // pinMode(LED, OUTPUT);
    // pinMode(LED2, OUTPUT);
    DDRB |= bit(BUTTON_PIN) | bit(LED1_PIN) | bit(LED2_PIN) | bit(ISR_FLAG); //  These pins as OUTPUTs in one go
    digitalWrite(BUTTON_PIN, LOW); // Discharge the capacitor
    // LED1 ON for 1s
    digitalWrite(LED1_PIN, HIGH);
    switchLedOffNow = millis() + LED_ON_DURATION;
    digitalWrite(ISR_FLAG, LOW);    // DEBUG: Helps counting the number of time the ISR is triggered

    // Send SPI data
    while (!(USISR & (1 << USIOIF))) {
      USICR |= (1 << USITC); // Toggle clock
    }
  }

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

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

  blinkIfAlive();
}
// #include <avr/io.h>