#include <Arduino.h>
#include "LowPower.h"

// Switch pins
#define Torch_SW A3 // Switch TORCH
#define Led_SW A4   // Switch LED
#define DC_SW A5    // Switch DC
#define USB_SW 13   // Switch USB
#define AC_SW 2     // Switch AC

// Output pins
#define Torch_PIN 3 // TORCH output
#define Led_PIN 4   // LED output
#define DC_PIN 5    // DC output
#define USB_PIN 6   // USB output
#define AC_PIN 7    // AC output

// State variables
bool torchState = false;
bool ledState = false;
bool dcState = false;
bool usbState = false;
bool acState = false;

// Wake-up flags
volatile bool wakeUpFlag = false;
volatile bool torchWakeUp = false;
volatile bool ledWakeUp = false;
volatile bool dcWakeUp = false;
volatile bool usbWakeUp = false;
volatile bool acWakeUp = false;

// Interrupt Service Routines
void acSwitchInterrupt() {
  acWakeUp = true;
  wakeUpFlag = true;
}

// Pin Change Interrupt for PCINT1 (A0-A5)
ISR(PCINT1_vect) {
  static uint8_t lastPortC = 0;
  uint8_t currentPortC = PINC;
  uint8_t changedBits = currentPortC ^ lastPortC;
  
  if (changedBits & (1 << 3)) { // A3 (Torch_SW)
    if (!(currentPortC & (1 << 3))) { // Falling edge (button pressed)
      torchWakeUp = true;
      wakeUpFlag = true;
    }
  }
  if (changedBits & (1 << 4)) { // A4 (Led_SW)
    if (!(currentPortC & (1 << 4))) { // Falling edge (button pressed)
      ledWakeUp = true;
      wakeUpFlag = true;
    }
  }
  if (changedBits & (1 << 5)) { // A5 (DC_SW)
    if (!(currentPortC & (1 << 5))) { // Falling edge (button pressed)
      dcWakeUp = true;
      wakeUpFlag = true;
    }
  }
  
  lastPortC = currentPortC;
}

// Pin Change Interrupt for PCINT0 (D8-D13)
ISR(PCINT0_vect) {
  static uint8_t lastPortB = 0;
  uint8_t currentPortB = PINB;
  uint8_t changedBits = currentPortB ^ lastPortB;
  
  if (changedBits & (1 << 5)) { // Pin 13 (bit 5 of PORTB)
    if (!(currentPortB & (1 << 5))) { // Falling edge (button pressed)
      usbWakeUp = true;
      wakeUpFlag = true;
    }
  }
  
  lastPortB = currentPortB;
}

void handleWakeUpSwitches() {
  wakeUpFlag = false;
  
  if (torchWakeUp) {
    torchWakeUp = false;
    torchState = !torchState;
    digitalWrite(Torch_PIN, torchState);
    Serial.print("WAKE-UP: Torch ");
    Serial.println(torchState ? "ON" : "OFF");
    delay(200);
  }
  
  if (ledWakeUp) {
    ledWakeUp = false;
    ledState = !ledState;
    digitalWrite(Led_PIN, ledState);
    Serial.print("WAKE-UP: LED ");
    Serial.println(ledState ? "ON" : "OFF");
    delay(200);
  }
  
  if (dcWakeUp) {
    dcWakeUp = false;
    dcState = !dcState;
    digitalWrite(DC_PIN, dcState);
    Serial.print("WAKE-UP: DC ");
    Serial.println(dcState ? "ON" : "OFF");
    delay(200);
  }
  
  if (usbWakeUp) {
    usbWakeUp = false;
    usbState = !usbState;
    digitalWrite(USB_PIN, usbState);
    Serial.print("WAKE-UP: USB ");
    Serial.println(usbState ? "ON" : "OFF");
    delay(200);
  }
  
  if (acWakeUp) {
    acWakeUp = false;
    acState = !acState;
    digitalWrite(AC_PIN, acState);
    Serial.print("WAKE-UP: AC ");
    Serial.println(acState ? "ON" : "OFF");
    delay(200);
  }
}

void goingToSleep() {
  Serial.println("Going to sleep...");
  Serial.flush();
  
  // Turn off all LEDs before sleep
  digitalWrite(Torch_PIN, LOW);
  digitalWrite(Led_PIN, LOW);
  digitalWrite(DC_PIN, LOW);
  digitalWrite(USB_PIN, LOW);
  digitalWrite(AC_PIN, LOW);

  // Clear wake-up flags
  wakeUpFlag = false;
  torchWakeUp = false;
  ledWakeUp = false;
  dcWakeUp = false;
  usbWakeUp = false;
  acWakeUp = false;

  // Setup interrupts for wake-up
  attachInterrupt(digitalPinToInterrupt(AC_SW), acSwitchInterrupt, FALLING);
  
  PCICR |= ((1 << PCIE1) | (1 << PCIE0));
  PCMSK1 |= ((1 << PCINT11) | (1 << PCINT12) | (1 << PCINT13));  // A3, A4, A5
  PCMSK0 |= (1 << PCINT5);                                       // D13

  // Sleep
  LowPower.powerDown(SLEEP_FOREVER, ADC_OFF, BOD_OFF);
  
  // Cleanup interrupts after wake-up
  detachInterrupt(digitalPinToInterrupt(AC_SW));
  PCICR &= ~((1 << PCIE1) | (1 << PCIE0));
  PCMSK1 = 0;
  PCMSK0 = 0;
  
  Serial.println("*** WOKE UP! ***");
}

void setup() {
  Serial.begin(9600);

  // Setup switch pins with pull-up
  pinMode(Torch_SW, INPUT_PULLUP);
  pinMode(Led_SW, INPUT_PULLUP);
  pinMode(DC_SW, INPUT_PULLUP);
  pinMode(USB_SW, INPUT_PULLUP);
  pinMode(AC_SW, INPUT_PULLUP);

  // Setup LED pins
  pinMode(Torch_PIN, OUTPUT);
  pinMode(Led_PIN, OUTPUT);
  pinMode(DC_PIN, OUTPUT);
  pinMode(USB_PIN, OUTPUT);
  pinMode(AC_PIN, OUTPUT);

  // Turn off all LEDs
  digitalWrite(Torch_PIN, LOW);
  digitalWrite(Led_PIN, LOW);
  digitalWrite(DC_PIN, LOW);
  digitalWrite(USB_PIN, LOW);
  digitalWrite(AC_PIN, LOW);

  Serial.println("=== Simple LED Toggle with Sleep ===");
  Serial.println("Press any switch to toggle LED and wake from sleep");
}

void loop() {
  // Handle wake-up from sleep
  if (wakeUpFlag) {
    handleWakeUpSwitches();
  }

  // Normal operation - check switches
  if (digitalRead(Torch_SW) == LOW) {
    delay(50);
    if (digitalRead(Torch_SW) == LOW) {
      torchState = !torchState;
      digitalWrite(Torch_PIN, torchState);
      while (digitalRead(Torch_SW) == LOW);
      Serial.print("Torch ");
      Serial.println(torchState ? "ON" : "OFF");
    }
  }

  if (digitalRead(Led_SW) == LOW) {
    delay(50);
    if (digitalRead(Led_SW) == LOW) {
      ledState = !ledState;
      digitalWrite(Led_PIN, ledState);
      while (digitalRead(Led_SW) == LOW);
      Serial.print("LED ");
      Serial.println(ledState ? "ON" : "OFF");
    }
  }

  if (digitalRead(DC_SW) == LOW) {
    delay(50);
    if (digitalRead(DC_SW) == LOW) {
      dcState = !dcState;
      digitalWrite(DC_PIN, dcState);
      while (digitalRead(DC_SW) == LOW);
      Serial.print("DC ");
      Serial.println(dcState ? "ON" : "OFF");
    }
  }

  if (digitalRead(USB_SW) == LOW) {
    delay(50);
    if (digitalRead(USB_SW) == LOW) {
      usbState = !usbState;
      digitalWrite(USB_PIN, usbState);
      while (digitalRead(USB_SW) == LOW);
      Serial.print("USB ");
      Serial.println(usbState ? "ON" : "OFF");
    }
  }

  if (digitalRead(AC_SW) == LOW) {
    delay(50);
    if (digitalRead(AC_SW) == LOW) {
      acState = !acState;
      digitalWrite(AC_PIN, acState);
      while (digitalRead(AC_SW) == LOW);
      Serial.print("AC ");
      Serial.println(acState ? "ON" : "OFF");
    }
  }

  // Check if all LEDs are OFF, then sleep
  if (!torchState && !ledState && !dcState && !usbState && !acState) {
    Serial.println("All LEDs OFF - Going to sleep");
    goingToSleep();
  }

  delay(100);
}