/*
 * PiPower
 * Copyright 2017 (c) Mark van Renswoude
 * 
 * Yet another power supply for the Raspberry Pi.
 * Designed to run on an ATTiny85 with a self-locking button.
 * 
 * 
 * 
 * Known "bug": if you turn the power button back on while it's shutting
 * down, it may not power cycle the Pi as it'll only turn the relay off
 * for the length of the button debounce time. Since there is a simple
 * workaround (turn the button off, wait a few seconds and turn it back
 * on again) I didn't bother to add the logic for this situation.
 */
#include <avr/sleep.h>


/*
 * Configuration
 */
// The button needs to be on a pin that supports interrupts
const uint8_t PinButton = 0;
const uint8_t PinButtonInterrupt = PCINT0;

// The LED needs to be on a pin that supports PWM
const uint8_t PinLED = 1;

// These just need digital I/O
const uint8_t PinRelay = 2;
const uint8_t PinShutdown = 3;
const uint8_t PinPiState = 4;


// The debounce time used to filter button noise
const unsigned long ButtonDebounceTime = 300;

// How long to wait for a signal from the Pi before alerting the user
const unsigned long BootProblemTime = 5000;

// How long to wait after the Pi signal turns off to turn off the power, to prevent false positives during a reboot
const unsigned long PowerOffTime = 5000;


// How long the LED stays on or off during blinking
const unsigned long LEDBlinkTime = 750;

// How long it takes to fade from on to off or back
const unsigned long LEDFadeTime = 750;



enum LEDState
{
  Off,
  On,
  Blinking,
  Fading
};

struct ButtonState
{
  bool pressed;
  bool debouncing;
};

typedef void (* StateHandler) ();


// Forward declarations
void gotoStateIdle();
void gotoStateBoot();
void gotoStateBootProblem();
void gotoStateShutdown();

ButtonState getButtonState();
void checkPiShutdown();
inline void checkLEDAnimation();

inline LEDState getLEDState();
void setLEDState(LEDState newState);
void setRelayState(uint8_t newState);
void setShutdownState(uint8_t newState);


// State variables
StateHandler currentState;
unsigned long currentTime;


void setup() 
{
  pinMode(PinLED, OUTPUT);
  pinMode(PinRelay, OUTPUT);
  pinMode(PinShutdown, OUTPUT);
  pinMode(PinPiState, INPUT);

  digitalWrite(PinLED, LOW);
  digitalWrite(PinRelay, LOW);
  digitalWrite(PinShutdown, LOW);

  gotoStateIdle();

  // Blink two times to indicate we've got power
  digitalWrite(PinLED, HIGH); delay(250); digitalWrite(PinLED, LOW); delay(250);
  digitalWrite(PinLED, HIGH); delay(250); digitalWrite(PinLED, LOW); delay(250);
}


void loop() 
{
  currentTime = millis();
  currentState();

  checkLEDAnimation();
}


/*
 * Idle
 * 
 * Power button is not pressed, no power going to the Pi.
 * 
 * LED:             Off
 * Relay:           Off
 * Shutdown signal: Off
 */
void stateIdle()
{
  ButtonState buttonState = getButtonState();
  
  if (buttonState.pressed)
  {
    gotoStateBoot();
  }
  else if (!buttonState.debouncing)
  {
    // Go into sleep mode until the button interrupt wakes us up again
    sleep();
  }  
}

void gotoStateIdle()
{
  setLEDState(Off);
  setRelayState(LOW);
  setShutdownState(LOW);

  currentState = stateIdle;
}


/*
 * Boot
 * 
 * Power going to the Pi, waiting for confirmation.
 * 
 * LED:             On
 * Relay:           On
 * Shutdown signal: Off
 */
unsigned long bootStart;

void stateBoot()
{
  if (digitalRead(PinPiState) == HIGH)
    gotoStateOn();
  else if (currentTime - bootStart >= BootProblemTime)
    gotoStateBootProblem();
}

void gotoStateBoot()
{
  setLEDState(On);
  setRelayState(HIGH);
  setShutdownState(LOW);

  bootStart = currentTime;
  currentState = stateBoot;
}


/*
 * Boot problem
 * 
 * No confirmation of a succesful boot within the timeout,
 * alert the user.
 * 
 * LED:             Blinking
 * Relay:           On
 * Shutdown signal: Off
 */
void stateBootProblem()
{
  if (digitalRead(PinPiState) == HIGH)
    gotoStateOn();
  else
  {
    ButtonState buttonState = getButtonState();
    if (!buttonState.pressed)
      gotoStateIdle();
  }
}

void gotoStateBootProblem()
{
  setLEDState(Blinking);
  setRelayState(HIGH);
  setShutdownState(LOW);

  currentState = stateBootProblem;
}



/*
 * On
 * 
 * Pi is running.
 * 
 * LED:             On
 * Relay:           On
 * Shutdown signal: Off
 */
void stateOn()
{
  ButtonState buttonState = getButtonState();
  if (!buttonState.pressed)
    gotoStateShutdown();
    
  checkPiShutdown();
}

void gotoStateOn()
{
  setLEDState(On);
  setRelayState(HIGH);
  setShutdownState(LOW);

  currentState = stateOn;
}



/*
 * Shutdown
 * 
 * Signal to the Pi that a shutdown is requested.
 * 
 * LED:             Fading
 * Relay:           On
 * Shutdown signal: On
 */
void stateShutdown()
{
  checkPiShutdown();
}

void gotoStateShutdown()
{
  setLEDState(Fading);
  setRelayState(HIGH);
  setShutdownState(HIGH);

  currentState = stateShutdown;
}



/*
 * Helpers
 */
// https://bigdanzblog.wordpress.com/2014/08/10/attiny85-wake-from-sleep-on-pin-state-change-code-example/
void sleep() 
{
  GIMSK |= _BV(PCIE);                     // Enable Pin Change Interrupts
  PCMSK |= _BV(PinButtonInterrupt);       // Use button as interrupt pin
  ADCSRA &= ~_BV(ADEN);                   // ADC off
  set_sleep_mode(SLEEP_MODE_PWR_DOWN);    // Replaces above statement

  sleep_enable();                         // Sets the Sleep Enable bit in the MCUCR Register (SE BIT)
  sei();                                  // Enable interrupts
  sleep_cpu();                            // Sleep

  cli();                                  // Disable interrupts
  PCMSK &= ~_BV(PinButtonInterrupt);      // Turn off button as interrupt pin
  sleep_disable();                        // Clear SE bit
  ADCSRA |= _BV(ADEN);                    // ADC on

  sei();                                  // Enable interrupts
}

ISR(PCINT0_vect) 
{
    // This is called when the interrupt occurs, but I don't need to do anything in it
}
    

 
uint8_t debouncedButtonState = LOW;
uint8_t lastButtonState = LOW;
unsigned long lastButtonChange = 0;

ButtonState getButtonState()
{
  ButtonState state;
  uint8_t currentButtonState = digitalRead(PinButton);

  if (currentButtonState != lastButtonState)
  {
    lastButtonState = currentButtonState;
    lastButtonChange = currentTime;
  }

  if (currentTime - lastButtonChange >= ButtonDebounceTime)
  {
    if (currentButtonState != debouncedButtonState)
      debouncedButtonState = currentButtonState;

    state.debouncing = false;
  }
  else
    state.debouncing = true;

  state.pressed = (debouncedButtonState == HIGH);
  return state;
}


bool shutdownDetected = false;
unsigned long shutdownStart;

void checkPiShutdown()
{
  if (digitalRead(PinPiState) == LOW)
  {
    if (shutdownDetected)
    {
      if (currentTime - shutdownStart >= PowerOffTime)
        gotoStateIdle();
    }
    else
    {
      shutdownStart = currentTime;
      shutdownDetected = true;
    }
  }
  else
    shutdownDetected = false;
}


LEDState ledState = Off;
unsigned long ledLastAnimation;
bool ledAnimationOn;

inline LEDState getLEDState()
{
  return ledState;  
}

void setLEDState(LEDState newState)
{
  if (newState == ledState)
    return;

  switch(newState)
  {
    case Off:
      digitalWrite(PinLED, LOW);
      break;

    case On:
      digitalWrite(PinLED, HIGH);
      break;

    default:
      ledLastAnimation = currentTime;
      ledAnimationOn = true;
      break;
  }
  ledState = newState;
}


inline void checkLEDAnimation()
{
  switch(getLEDState())
  {
    case Blinking:
      if (currentTime - ledLastAnimation >= LEDBlinkTime)
      {
        ledAnimationOn = !ledAnimationOn;
        ledLastAnimation = currentTime;
      }

      digitalWrite(PinLED, ledAnimationOn ? HIGH : LOW);
      break;
          
    case Fading:
      unsigned long elapsedTime = currentTime - ledLastAnimation;
      while (elapsedTime >= LEDFadeTime)
      {
        ledAnimationOn = !ledAnimationOn;
        ledLastAnimation = currentTime;
        
        elapsedTime -= LEDFadeTime;
      }
      
      byte currentStep = (byte)((elapsedTime * 255) / LEDFadeTime);
      if (!ledAnimationOn)
        currentStep = 255 - currentStep;
        
      analogWrite(PinLED, currentStep);
      break;
  }
}


uint8_t relayState = LOW;

void setRelayState(uint8_t newState)
{
  if (newState == relayState)
    return;

  digitalWrite(PinRelay, newState);
  relayState = newState;
}


uint8_t shutdownState = LOW;

void setShutdownState(uint8_t newState)
{
  if (newState == shutdownState)
    return;

  digitalWrite(PinShutdown, newState);
  shutdownState = newState;  
}