#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include "esp_sleep.h"
#include <AD9833.h>
#include <SPI.h>
#include <string>
#include <cstring>
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1
#define SCREEN_ADDRESS 0x3C  // I2C address

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
// flags for interrupts
volatile bool resetFlag = false;
volatile bool menuFlag = false;
volatile bool selFlag = false;
RTC_DATA_ATTR bool sleepFlag = false;

void IRAM_ATTR _reset() {
  resetFlag = true;
}

void IRAM_ATTR _menu() {
  menuFlag = true;
}

void IRAM_ATTR _sel() {
  selFlag = true;
}
void IRAM_ATTR _sleep() { sleepFlag = true; }

// Define button pins
const int wave_selector_btn = 33;
const int reset_btn = 32;
const int menu_btn = 25;
const int sel_pin = 13;  // Added sel pin to pin 4
const int power_btn = 33;

// AD9833 SPI Configuration
#define FSYNC 5  // Chip select pin for AD9833 (can be any GPIO)
AD9833 signalGenerator(FSYNC);

bool blink_state = true;


// State-related variables
int cursorPosition = 0; // Tracks position of cursor on frequency
int currentScreen = 0;  // 0 = control screen, 1 = specs, 2 = error
float frequency = 200.0000; // Default frequency in Hz
float voltage = 0.0;      // Voltage reading
portMUX_TYPE myMutex = portMUX_INITIALIZER_UNLOCKED;
String waveType = "Square"; // Default wave type
String freqStr = "000200.0";
int pointy = 6;
int freqUnitIndex = 0; // Current frequency unit: 0 = Hz, 1 = kHz, 2 = MHz

const int DAC_PIN = 26;
const int TABLE_SIZE = 256;

float sinTable[TABLE_SIZE];
float sawTable[TABLE_SIZE];
float triTable[TABLE_SIZE];
//////////////////////////////////////////////////////////
//FUNCTION GENERATOR
//////////////////////////////////////////////////////////

class FunctionGenerator {
  public:
    // Initialize the function generator
    void init() {
      pinMode(power_btn, INPUT_PULLUP);
    esp_sleep_wakeup_cause_t wakeup_reason = esp_sleep_get_wakeup_cause();
    if (wakeup_reason == ESP_SLEEP_WAKEUP_EXT0) {
      sleepFlag = false;
    } else if(sleepFlag){
      sleep();
    }
      Wire.begin();  // Initialize the Wire library for I2C communication
    SPI.begin();

      pinMode(DAC_PIN, OUTPUT);

      if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
        Serial.begin(9600);
        Serial.println(F("SSD1306 allocation failed"));
      }
      display.clearDisplay();
      display.setTextSize(1);
      display.setTextColor(SSD1306_WHITE);  // Set text color to white
      display.setCursor(0, 0);
      display.println("Initializing...");
      display.display();
      pinMode(4, INPUT); // Joystick X-axis (Horizontal on pin 5)
      pinMode(0, INPUT); // Joystick Y-axis (Vertical on pin 6)

      pinMode(reset_btn, INPUT_PULLUP);
      pinMode(menu_btn, INPUT_PULLUP);
      pinMode(sel_pin, INPUT_PULLUP);  // Set sel pin to input with pullup
    attachInterrupt(digitalPinToInterrupt(power_btn), _sleep, FALLING);
      attachInterrupt(digitalPinToInterrupt(reset_btn), _reset, FALLING);
      attachInterrupt(digitalPinToInterrupt(menu_btn), _menu, FALLING);
      attachInterrupt(digitalPinToInterrupt(sel_pin), _sel, FALLING);
      // signalGenerator.begin();
      displayControlScreen();
    }
    //////////////////////////////////////////////////////////
    //DISPLAY
    //////////////////////////////////////////////////////////

    void displayControlScreen() {
      display.clearDisplay();

      display.setCursor(10, 20);
      display.print("Wave: ");
      portENTER_CRITICAL(&myMutex);
      display.println(waveType);
      portEXIT_CRITICAL(&myMutex);
      display.setCursor(10, 35);
      display.print("Freq: ");
      display.print(freqStr);
      display.print(getFrequencyUnit());
      display.println("Hz");

      display.setCursor(10, 50);
      display.print("Amp: ");
      display.print(voltage);
      display.println("V");

      if (blink_state) {
        display.setCursor(46 + (cursorPosition) * 6, 37);
        display.print("_");
        // delay(200);
      }

      display.display();
    }

    void displaySpecs() {
      display.clearDisplay();

      display.setCursor(5, 20);
      display.println("Voltage Limit: 0-20V");

      display.setCursor(5, 35);
      display.println("Freq: 1Hz-400kHz");

      display.display();
    }

    void displayError() {
      display.clearDisplay();

      display.setCursor(5, 20);
      display.println("Voltage Err: +/-0.1V");

      display.setCursor(5, 35);
      display.println("Freq Err: +/-0.01%");

      display.display();
    }

    //////////////////////////////////////////////////////////
    //BUTTON
    //////////////////////////////////////////////////////////
    void sleep(){
      vTaskDelay(200);
      while(digitalRead(power_btn) == LOW) {
    delay(10);
  }
      display.clearDisplay();
      display.display();
      esp_sleep_enable_ext0_wakeup((gpio_num_t)power_btn, LOW);
      esp_deep_sleep_start();
    }
    void updateFrequency(int change) {
      // if (cursorPosition != 10) {
      //   int x = (int(freqStr[cursorPosition - 5] - '0') + change) % 10;
      //   if (x == -1) x = 9;
      //   freqStr[cursorPosition - 5] = '0' + x;
      // } else {
      //   freqUnitIndex = !freqUnitIndex;
      // }
      float freq = frequency;
      if (cursorPosition == 8){
        freqUnitIndex = !freqUnitIndex;
        if(freqUnitIndex) freq *= 1000;
        else freq /= 1000; 
      }
      else{
      freq = frequency + change*pow(10,pointy-cursorPosition-1*(pointy>cursorPosition)+freqUnitIndex*3);
      }
      if(freq<=400000 && freq>=0.1) {
        frequency = freq;
        if(frequency >= 1000){
          if(!freqUnitIndex && cursorPosition<6 && cursorPosition>2) cursorPosition++;
          freqUnitIndex = 1;
        }
        if(frequency < 1000){
          if(freqUnitIndex && cursorPosition<6 && cursorPosition>2) cursorPosition--;
          freqUnitIndex = 0;
        }
        setSignal();
      }
    }

    // Handle wave selection


    // Get frequency unit as string
    String getFrequencyUnit() {
      if (freqUnitIndex == 0) return " ";
      if (freqUnitIndex == 1) return "k";

      return " ";
    }

    // Reset to initial state
    void reset() {
      frequency = 200.0000;
      String freqStr = "000200.0";
      pointy = 6;
      voltage = 0.0;
      portENTER_CRITICAL(&myMutex);
      waveType = "Square";
      portEXIT_CRITICAL(&myMutex);
      cursorPosition = 0;
      currentScreen = 0;
      displayControlScreen();

      // 001.0000
      // 010.0000
      // 100.0000
      // 000100.0
      // 000010.0
      // 000001.0
      // 000000.1
      setSignal();
    }
    void menu() {
      currentScreen = (currentScreen + 1) % 3; // Cycle through screens
      if (currentScreen == 0) {
        displayControlScreen();
      } else if (currentScreen == 1) {
        displaySpecs();
      } else {
        displayError();
      }
    }
    void sel() {
      if (currentScreen) return;
      portENTER_CRITICAL(&myMutex);
      if (waveType == "Square") {
        waveType = "Triangle";
      } else if (waveType == "Triangle") {
        waveType = "Sine";
      } else {
        waveType = "Square";
      }
      portEXIT_CRITICAL(&myMutex);
      setSignal();
      displayControlScreen();

    }



    //////////////////////////////////////////////////////////
    //Signal
    //////////////////////////////////////////////////////////


    void generate_sine(float frequency) {
      int delay_us = 1000000 / (frequency * TABLE_SIZE);
      for (int i = 0; i < TABLE_SIZE; i++) {
        int dacValue = map(sinTable[i] * 1000, -1000, 1000, 0, 255);
        dacWrite(DAC_PIN, dacValue);
        delayMicroseconds(delay_us);
      }
    }

    void generate_sawtooth(float frequency) {
      int delay_us = 1000000 / (frequency * TABLE_SIZE);
      for (int i = 0; i < TABLE_SIZE; i++) {
        int dacValue = map(sawTable[i] * 1000, 0, 1000, 0, 255);
        dacWrite(DAC_PIN, dacValue);
        delayMicroseconds(delay_us);
      }
    }

    void generate_triangle(float frequency) {
      int delay_us = 1000000 / (frequency * TABLE_SIZE);
      for (int i = 0; i < TABLE_SIZE; i++) {
        dacWrite(DAC_PIN, triTable[i]);
        delayMicroseconds(delay_us);
      }
    }

    void initialise_signals() {
      for (int i = 0; i < TABLE_SIZE; i++) {
        sinTable[i] = sin((float)i / TABLE_SIZE * 2 * PI);
        sawTable[i] = (float)i / TABLE_SIZE;
        triTable[i] = (i < TABLE_SIZE / 2) ? map(i, 0, TABLE_SIZE / 2 - 1, 0, 255) : map(i, TABLE_SIZE / 2, TABLE_SIZE - 1, 255, 0);
      }
    }

    void setSignal() {
      printf("%.2f\n",frequency);
    portENTER_CRITICAL(&myMutex);
    String type = waveType;
    portEXIT_CRITICAL(&myMutex);
    std::string tempstr;
    if(freqUnitIndex) tempstr = std::to_string(frequency/1000.0);
    else tempstr = std::to_string(frequency);
    if(frequency >= 100000){
        pointy = 3;
        for(int i = 0; i < 8;i++)
        {
            freqStr[i] = tempstr[i];
        }
    }
    else if(frequency <10000 && frequency >=1000){
        freqStr[0] = '0';
        freqStr[1]='0';
        for(int i =0;i<6;i++){
            freqStr[i+2] = tempstr[i];
        }
        pointy = 3;
    }
    else if(frequency >=10000 && frequency <100000){
        freqStr[0] = '0';
        for(int i =0;i<7;i++){
            freqStr[i+1] = tempstr[i];
        }
        pointy = 3;
    }
    else
    {
        pointy = 6;
        freqStr[0] = '0';
        freqStr[1] = '0';
        freqStr[2] = '0';
        if(frequency >=100){
        for(int i =0;i<5;i++){
            freqStr[i+3] = tempstr[i];
        }}
        else if(frequency < 100 && frequency >= 10)
        {
          freqStr[3]='0';
          for(int i =0;i<4;i++){
            freqStr[i+4] = tempstr[i];
        } 
        }
        else if(frequency < 10 && frequency >= 1)
        {
          freqStr[3]='0';
          freqStr[4] ='0';
          for(int i =0;i<3;i++){
            freqStr[i+5] = tempstr[i];
        } 
        }
        else if(frequency < 1 && frequency >= 0.1)
        {
          printf("hi");
          freqStr[3]='0';
          freqStr[4] = '0';
          for(int i =0;i<3;i++){
            freqStr[i+5] = tempstr[i];
        } 
        }

    }
    // strncpy(freqStr, tempstr.c_str(), sizeof(freqStr) - 1);
    // freqStr[sizeof(freqStr) - 1] = '\0'; 
    // if (type == "Sine") {
    //     generate_sine(frequency);
    // } else if (type == "Square") {
    //     generate_sawtooth(frequency);
    // } else if (type == "Triangle") {
    //     generate_triangle(frequency);
    // }
}
void generateWave(){
  if (waveType == "Sine") {
        generate_sine(frequency);
    } else if (waveType == "Square") {
        generate_sawtooth(frequency);
    } else if (waveType == "Triangle") {
        generate_triangle(frequency);
    }
}
};

//////////////////////////////////////////////////////////
//GLOBAL
//////////////////////////////////////////////////////////
FunctionGenerator generator;

void setSignalTask(void *pvParameters) {
  while (true) {
    generator.generateWave(); // Call your setSignal function to run continuously
    vTaskDelay(1); // Yield to other tasks (you can adjust delay as per requirement)
  }
}
// Function generator handler
void function_generator_handler() {
  // Handle screen selection using menu button
  if(sleepFlag) {
    generator.sleep();
  }
  if (resetFlag) {
    // Handle reset event
    generator.reset();
    delay(200);
    resetFlag = false;
  }

  if (menuFlag) {
    // Handle menu event
    generator.menu();
    delay(200);
    menuFlag = false;
  }

  if (selFlag) {
    // Handle selection event
    generator.sel();
    delay(200);
    selFlag = false;

  }
  int joyX = analogRead(4); // Read horizontal movement
  int joyY = analogRead(0); // Read vertical movement
  // esp_sleep_wakeup_cause_t wakeup_reason = esp_sleep_get_wakeup_cause();
  // printf("%b\n",wakeup_reason == ESP_SLEEP_WAKEUP_EXT0);

  // Handle cursor movement with joystick
  if (!currentScreen) {
    if (joyX == 4095) { // Move cursor left
      cursorPosition = max(0, cursorPosition - 1);
      if (cursorPosition == pointy)cursorPosition--;
      generator.displayControlScreen();
      delay(200); // Debounce
    } else if (joyX == 0) { // Move cursor right
      cursorPosition = min(8, cursorPosition + 1);
      if (cursorPosition == pointy)cursorPosition++;
      generator.displayControlScreen();
      delay(200); // Debounce
    } else if (joyY == 4095) { // Increase value at cursor position
      generator.updateFrequency(1);
      generator.displayControlScreen();
      delay(200); // Debounce
    } else if (joyY == 0) { // Decrease value at cursor position
      generator.updateFrequency(-1);
      generator.displayControlScreen();
      delay(200); // Debounce
    }

    float new_voltage = (analogRead(12)) * (20.0 / 4095.0); // Scale to 0-20V

    // Update voltage if changed
    if (new_voltage != voltage) {
      voltage = new_voltage;
      generator.displayControlScreen();
    }
  }
}

void setup() {
  generator.init();
  xTaskCreatePinnedToCore(setSignalTask, "setSignalTask", 4096, NULL, 1, NULL, 1);
}

void loop() {
  function_generator_handler();
}