#include <Arduino.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <EEPROM.h>
#include "Button_Controler.h"
#include "Motor_Controler.h"
#include "SettingsAndConstants.h"

Button_Controler b_con(LEFT_B_PIN, MIDDLE_B_PIN, RIGHT_B_PIN);
Adafruit_SSD1306 display(128, 64, &Wire, -1);
Motor_Controler motor_control(MOSI_PIN, MISO_PIN, SCK_PIN,
                              EN_S_PIN, DIR_S_PIN, STEP_S_PIN, CS_S_PIN, step_microsteps, STEP_MOTOR_CURRENT, STEP_S_CHOP, TAPE_MOTOR_FULL_STEP_ANGLE, TOOTHS_PER_TAPE_WHEEL,
                              EN_W_PIN, DIR_W_PIN, STEP_W_PIN, CS_W_PIN, width_microsteps, WIDTH_MOTOR_CURRENT, WIDTH_S_CHOP, WIDTH_MOTOR_FULL_STEP_ANGLE, WIDTH_LEAD_PER_ROTATE,
                              RSENSE, STEP_H_CURRENT_COF, WIDTH_H_CURRENT_COF,
                              HOLE_SENSOR_PIN, HOME_SENSOR_PIN, initPosition);
                              
EEPROMClass WIDTH("width", sizeof(int));
EEPROMClass SPEED("speed", sizeof(int));
EEPROMClass PITCH("pitch", sizeof(int));

volatile ScreenscreenStatus screenStat = SHOW_SET_ROUTINE;
volatile RuntimeStatus runtimeStat = OPERATING;

unsigned int sel_speed;
unsigned int sel_pitch;
unsigned int sel_width;

unsigned int refreshes = 0;
unsigned int refresh_counter = 0;
unsigned int menu_selector_index;
bool step_ret;
String menu_name;
RuntimeStatus actSemStat = NONE;

TaskHandle_t secondaryLoopTask;

void SecondaryLoop(void *param)
{
  while(true)
  {
    b_con.Refresh();
    delay(10);

    // Switch semafor
    if (actSemStat != runtimeStat)
    {
      switch (runtimeStat)
      {
        case OPERATING:
          digitalWrite(LED1, LOW);
          digitalWrite(LED2, HIGH);
          break;

        case SETTING:
          digitalWrite(LED1, LOW);
          digitalWrite(LED2, LOW);
          break;

        default:
          digitalWrite(LED1, HIGH);
          digitalWrite(LED2, LOW);
          break;
      }
      actSemStat = runtimeStat;
    }

    #ifdef SIM_STEPS
    if (runtimeStat == OPERATING)
    {
      digitalWrite(FEEDBACK_PIN, LOW);
      step_ret = motor_control.Step(true, step_speed_arr[sel_speed], pitch_arr[sel_pitch]);
      if(!step_ret)
      {
        runtimeStat = ERROR_ST;
        screenStat = SHOW_SET_ROUTINE;
      }
      else
      {
        digitalWrite(FEEDBACK_PIN, HIGH);
      }
      delay(200); 
    }
    #else
    if (runtimeStat == OPERATING)
    {
      // Wait for trigger signal to step
      if (digitalRead(TRIGGER_PIN) == HIGH)
      {
        digitalWrite(FEEDBACK_PIN, LOW);
        delay(50);
        while(digitalRead(TRIGGER_PIN) == HIGH);
        
        step_ret = motor_control.Step(true, step_speed_arr[sel_speed], pitch_arr[sel_pitch]);
        if(!step_ret)
        {
          runtimeStat = ERROR_ST;
          screenStat = SHOW_SET_ROUTINE;
        }
        else
        {
          digitalWrite(FEEDBACK_PIN, HIGH);
        }
        delay(50);
      }
    }
    #endif
  }
}

void setup() {
  // Start serial
  #ifdef SERIAL_COM
  Serial.begin(115200);
  while(!Serial);
  #endif
  
  // Init semafor
  pinMode(LED1, OUTPUT);
  pinMode(LED2, OUTPUT);
  digitalWrite(LED1, HIGH);
  digitalWrite(LED2, HIGH);
  
  // Init display
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { for(;;); }
  delay(150);

  // Setup display
  display.clearDisplay();
  display.display();
  display.setTextColor(SSD1306_WHITE);
  if (ROTATE_DISPLAY) display.setRotation(2);

  // Show logo and version
  display.setTextSize(1);
  display.drawBitmap(0, 0, eps_logo, 128, 64, 1);
  display.setCursor(0, 0);
  display.print(VERSION);
  display.display();
  display.setTextSize(2);

  // Test connection to drivers
  #ifndef DISABLE_DRIVER_CON_TEST
  while(!motor_control.Test_Con())
  {
    Serial.println("Cant connect to drivers!");
    Serial.println("Waiting...");
    delay(1000);
  }
  #endif

  // Init EEPROM
  if (!WIDTH.begin(WIDTH.length()))
  {
    Serial.println("Failed to init WIDTH");
    Serial.println("Restarting...");
    delay(1000);
    ESP.restart();
  }
  if (!SPEED.begin(SPEED.length()))
  {
    Serial.println("Failed to init SPEED");
    Serial.println("Restaring...");
    delay(1000);
    ESP.restart();
  }
  if (!PITCH.begin(PITCH.length()))
  {
    Serial.println("Failed to init PITCH");
    Serial.println("Restarting...");
    delay(1000);
    ESP.restart();
  }

  // Load saved values
  WIDTH.get(0, sel_width);
  SPEED.get(0, sel_speed);
  PITCH.get(0, sel_pitch);
  if (sel_width >= ((sizeof(width_arr) / sizeof(width_arr[0]))) || sel_speed >= ((sizeof(step_speed_arr) / sizeof(step_speed_arr[0]))) || sel_pitch >= ((sizeof(pitch_arr) / sizeof(pitch_arr[0]))))
  {
    Serial.println("Initializing EEPROM");
    sel_width = 0;
    sel_speed = 0;
    sel_pitch = 0;
    WIDTH.put(0, sel_width);
    SPEED.put(0, sel_speed);
    PITCH.put(0, sel_pitch);
    WIDTH.commit();
    SPEED.commit();
    PITCH.commit();
  }

  Serial.println("\nSaved settings:");
  Serial.print("Selected width index: ");
  Serial.println(sel_width);
  Serial.print("Selected speed index: ");
  Serial.println(sel_speed);
  Serial.print("Selected pitch index: ");
  Serial.println(sel_pitch);

  // Show motor parameters
  motor_control.ShowParams();

  // Init interface
  pinMode(TRIGGER_PIN, INPUT);
  pinMode(FEEDBACK_PIN, OUTPUT);

  // Start communication task
  xTaskCreatePinnedToCore(SecondaryLoop, "SecLoop", 10000, NULL, 1, &secondaryLoopTask, 0);
  delay(1000);
}

void loop() {
  if (refresh_counter > REFRESH_LOOPS)
  {
    switch (screenStat)
    {
      // Main screen waiting routine
      case RUNNING:      
        running_routine();
        break;

      // Menu items selector
      case MAIN_MENU:
        display.clearDisplay();
        print_top_stats();
        
        display.setCursor(menu_name_margins[menu_selector_index], 25);
        menu_name = menu_names[menu_selector_index];
        display.println(menu_name);
        
        display.setCursor(0, 48);
        display.print("<");

        // Special selector for toggle motor
        if (menu_name == "Togg Motor")
        {
          bool motor_disabled = digitalRead(EN_S_PIN);
          if (motor_disabled)
          {
            display.setCursor(29, 48);
            display.print("Enable");
          }
          else
          {
            display.setCursor(25, 48);
            display.print("Disable");
          }
        }
        else
        {
          display.setCursor(28, 48);
          display.print("Select");
        }
        
        display.setCursor(116, 48);
        display.print(">");
        display.display();

        if (b_con.GetLeftState())
        {
          refreshes = 0;
          if (menu_selector_index > 0) menu_selector_index--;
          else menu_selector_index = ((sizeof(menu_names) / sizeof(menu_names[0])) - 1);
          b_con.Clear();
        }

        if (b_con.GetRightState())
        {
          refreshes = 0;
          if (menu_selector_index < ((sizeof(menu_names) / sizeof(menu_names[0])) - 1)) menu_selector_index++;
          else menu_selector_index = 0;
          b_con.Clear();
        }

        if (b_con.GetMidState())
        {
          refreshes = 0;
          
          // Main menu selector
          if (menu_name == "Exit")
          {
            screenStat = SHOW_SET_ROUTINE;
          }
          else if (menu_name == "Step")
          {
            motor_control.Step(true, step_speed_arr[sel_speed], pitch_arr[sel_pitch]);
          }
          else if (menu_name == "Stepback")
          {
            if(!motor_control.Step(false, step_speed_arr[sel_speed], pitch_arr[sel_pitch]))
            {
              runtimeStat = ERROR_ST;
              screenStat = SHOW_SET_ROUTINE;
            }
          }
          else if (menu_name == "Speed")
          {
            screenStat = MENU_SPEED;
          }
          else if (menu_name == "Pitch")
          {
            screenStat = MENU_PITCH;
          }
          else if (menu_name == "Width")
          {
            screenStat = MENU_WIDTH;
          }
          else if (menu_name == "JOG Tape")
          {
            screenStat = TRIM_STEP;
          }
          else if (menu_name == "JOG Width")
          {
            screenStat = TRIM_W;
          }
          else if (menu_name == "Home W")
          {
            if (!motor_control.WasHomed()) screenStat = HOME_WIDTH;
          }
          else if (menu_name == "Togg Motor")
          {
            motor_control.SwitchMotorSleep();
          }
          else if (menu_name == "Rollout T")
          {
            screenStat = ROLLOUT_POPUP;
          }
          
          b_con.Clear();
        }

        // Return to main screen after some time
        if (refreshes > REFRESHES_BEFORE_RETURN2MAINSCREEN)
        {
          refreshes = 0;
          screenStat = SHOW_SET_ROUTINE;
        }
        
        refreshes++;
        break;

      case MENU_SPEED:
        // Setting speed of steps
        speed_change_menu();
        break;

      case MENU_PITCH:
        // Setting pitch
        pitch_change_menu();
        break;

      case MENU_WIDTH:
        // Setting of width in memory
        width_change_menu();
        break;

      case WIDTH_SETTING_POPUP:
        // Homing and setting width
        width_setting_popup();
        break;

      case SHOW_SET_ROUTINE:
        // Show setting and ready for operate and set feedback signal
        if (runtimeStat == SETTING || runtimeStat == OPERATING)
        {
          runtimeStat = OPERATING;
          digitalWrite(FEEDBACK_PIN, HIGH);
        }
        show_screenStats();
        refreshes = 0;
        break;

      case TRIM_STEP:
        // Trim step rountine
        step_trim_routine();
        break;

      case TRIM_W:
        // Trim width routine
        width_trim_routine();
        break;

      case HOME_WIDTH:
        // Popup before homing
        home_width_popup();
        break;

      case ROLLOUT_POPUP:
        // Popup before tape rollout
        rollout_popup();
        break;

      case OFF_SCREEN:
        // Screen shutoff
        if (b_con.GetLeftState() || b_con.GetMidState() || b_con.GetRightState())
        {
          screenStat = SHOW_SET_ROUTINE;
          Serial.println("Screen: ON");
          b_con.Clear();
        }
        break;

      default:
        // Default state for some error purpose - back to main screen
        screenStat = SHOW_SET_ROUTINE;
        break;
    }
    refresh_counter = 0;
  }

  refresh_counter++;
}

void print_top_stats()
{
  display.setCursor(4, 0);
  display.print("P");
  display.print(pitch_arr[sel_pitch] * 4);
  display.print(" W");
  display.print((int)(width_arr[sel_width]));
  display.print(" S");
  display.print(sel_speed + 1);
  display.drawLine(0, 16, display.width(), 16, SSD1306_WHITE);
}

void full_screen_stats()
{
  display.setCursor(0,0);
  display.print("Pitch: ");
  display.println(pitch_arr[sel_pitch] * 4);
  display.print("Width: ");
  display.println((int)(width_arr[sel_width]));
  display.print("Speed: ");
  display.println(sel_speed + 1);
}

void show_screenStats()
{
  display.clearDisplay();

  display.drawLine(0, 47, display.width(), 47, SSD1306_WHITE);

  if (runtimeStat != OPERATING)
  {
    display.setCursor(68, 50);
    display.println("CLEAR");
    
    switch (runtimeStat)
    {
      case ERROR_ST:
        display.setCursor(18, 25);
        display.println("ERROR S!");
        break;

      case ERROR_HM:
        display.setCursor(18, 25);
        display.println("ERROR H!");
        break;

      default:
        break;
    } 
  }
  else
  {
    display.setCursor(0, 50);
    display.print("MENU");
    full_screen_stats();
  }

  display.display();
        
  screenStat = RUNNING;
}

void speed_change_menu()
{
  display.clearDisplay();
  display.setCursor(12, 0);
  display.println("Speed set");
  display.drawLine(0, 16, display.width(), 16, SSD1306_WHITE);

  display.setCursor(58, 25);
  display.println(sel_speed + 1);

  if (sel_speed > 0)
  {
    display.setCursor(0, 50);
    display.print("-");
  }
  display.setCursor(51, 50);
  display.print("OK");

  if (sel_speed < ((sizeof(step_speed_arr) / sizeof(step_speed_arr[0])) - 1))
  {
    display.setCursor(116, 48);
    display.println("+");
  }
  
  display.display();
  
  if (b_con.GetMidState())
  {
    screenStat = MAIN_MENU;
    b_con.Clear();
    SPEED.put(0, sel_speed);
    SPEED.commit();
    return;
  }

  if (sel_speed > 0 && b_con.GetLeftState())
  {
    sel_speed--;
    b_con.Clear();
  }

  if (sel_speed < ((sizeof(step_speed_arr) / sizeof(step_speed_arr[0])) - 1) && b_con.GetRightState())
  {
    sel_speed++;
    b_con.Clear();
  }
}

void pitch_change_menu()
{
  display.clearDisplay();
  display.setCursor(12, 0);
  display.println("Pitch set");
  display.drawLine(0, 16, display.width(), 16, SSD1306_WHITE);

  display.setCursor(52, 25);
  display.println(pitch_arr[sel_pitch] * 4);

  if (sel_pitch > 0)
  {
    display.setCursor(0, 50);
    display.print("-");
  }

  display.setCursor(51, 50);
  display.print("OK");

  if (sel_pitch < ((sizeof(pitch_arr) / sizeof(pitch_arr[0])) - 1))
  {
    display.setCursor(116, 48);
    display.println("+");
  }

  display.display();
  
  if (b_con.GetMidState())
  {
    screenStat = MAIN_MENU;
    b_con.Clear();
    PITCH.put(0, sel_pitch);
    PITCH.commit();
    return;
  }

  if (sel_pitch > 0 && b_con.GetLeftState())
  {
    sel_pitch--;
    b_con.Clear();
  }

  if (sel_pitch < ((sizeof(pitch_arr) / sizeof(pitch_arr[0])) - 1) && b_con.GetRightState())
  {
    sel_pitch++;
    b_con.Clear();
  }
}

void width_change_menu()
{
  display.clearDisplay();
  display.setCursor(12, 0);
  display.println("Width set");
  display.drawLine(0, 16, display.width(), 16, SSD1306_WHITE);

  display.setCursor(52, 25);
  display.println((int)width_arr[sel_width]);

  if (sel_width > 0)
  {
    display.setCursor(0, 50);
    display.print("-");
  }

  display.setCursor(51, 50);
  display.print("OK");

  if (sel_width < ((sizeof(width_arr) / sizeof(width_arr[0])) - 1))
  {
    display.setCursor(116, 50);
    display.println("+");
  }

  display.display();

  if (b_con.GetMidState())
  {
    screenStat = WIDTH_SETTING_POPUP; 
    
    b_con.Clear();
    return;
  }

  if (sel_width > 0 && b_con.GetLeftState())
  {
    sel_width--;
    b_con.Clear();
  }

  if (sel_width < ((sizeof(width_arr) / sizeof(width_arr[0])) - 1) && b_con.GetRightState())
  {
    sel_width++;
    b_con.Clear();
  }
}

void width_setting_popup()
{
  display.clearDisplay();
  display.setCursor(17, 0);
  display.println("Execute?");
  display.drawLine(0, 16, display.width(), 16, SSD1306_WHITE);

  display.setCursor(0, 50);
  display.print("NO");

  display.setCursor(92, 50);
  display.println("YES");

  display.display();

  // Left width setting
  if (b_con.GetLeftState())
  {
    b_con.Clear();
    WIDTH.get(0, sel_width);
    screenStat = MAIN_MENU;
  }

  if (b_con.GetRightState())
  {
    b_con.Clear();

    // Save selected width on confirm
    WIDTH.put(0, sel_width);
    WIDTH.commit();

    bool homed = motor_control.WasHomed();
    if (!homed)
    {
      // If not homed go to home width screen
      display.clearDisplay();
      display.setCursor(12, 25);
      display.println("Homing...");
      display.display();

      bool home_ret = motor_control.HomeWidth(home_speed);

      if (!home_ret)
      {
        runtimeStat = ERROR_HM;
        screenStat = SHOW_SET_ROUTINE;
        return;
      }

      display.clearDisplay();
      display.setCursor(6, 25);
      display.println("Setting...");
      display.display();

      motor_control.SetWidth(width_arr[sel_width] + width_offset, width_speed);
      screenStat = MAIN_MENU;
    }
    else
    {
      // If homed execute width setting
      if ((motor_control.GetPosition() - width_offset) != width_arr[sel_width])
      {
        display.clearDisplay();
        display.setCursor(6, 25);
        display.println("Setting...");
        display.display();
      
        motor_control.SetWidth(width_arr[sel_width] + width_offset, width_speed);
      }
      
      screenStat = MAIN_MENU;
    }
  }
}

void rollout_popup()
{
  display.clearDisplay();
  display.setCursor(17, 0);
  display.println("Execute?");
  display.drawLine(0, 16, display.width(), 16, SSD1306_WHITE);

  display.setCursor(0, 50);
  display.print("NO");

  display.setCursor(92, 50);
  display.println("YES");

  display.display();

  if (b_con.GetLeftState())
  {
    b_con.Clear();
    screenStat = MAIN_MENU;
  }

  if (b_con.GetRightState())
  {
    b_con.Clear();
    
    display.clearDisplay();
    display.setCursor(6, 25);
    display.println("Rolling...");
    display.display();

    motor_control.RollOutTape(sel_pitch, rollout_speed);
    screenStat = MAIN_MENU;
  }
}

void step_trim_routine()
{
  display.clearDisplay();
  display.setCursor(29, 0);
  display.println("JOG T");
  display.drawLine(0, 16, display.width(), 16, SSD1306_WHITE);

  display.setCursor(0, 50);
  display.print("F");

  display.setCursor(40, 50);
  display.print("EXIT");

  display.setCursor(116, 50);
  display.println("B");

  display.display();

  if (b_con.GetMidState())
  {
    screenStat = MAIN_MENU;
  }

  if (b_con.GetLeftState())
  {
    motor_control.TrimStep(true, step_speed_arr[sel_speed]);
  }

  if (b_con.GetRightState())
  {
    motor_control.TrimStep(false, step_speed_arr[sel_speed]);
  }
  
  b_con.Clear();
}

void width_trim_routine()
{
  display.clearDisplay();
  display.setCursor(29, 0);
  display.println("JOG W");
  display.drawLine(0, 16, display.width(), 16, SSD1306_WHITE);

  display.setCursor(0, 50);
  display.print("-");

  display.setCursor(40, 50);
  display.print("EXIT");

  display.setCursor(116, 50);
  display.println("+");

  display.display();

  if (b_con.GetMidState())
  {
    screenStat = MAIN_MENU;
  }

  if (b_con.GetLeftState())
  {
    motor_control.TrimWidth(false, width_speed);
  }

  if (b_con.GetRightState())
  {
    motor_control.TrimWidth(true, width_speed);
  }
  
  b_con.Clear();
}

void home_width_popup()
{
  display.clearDisplay();
  display.setCursor(17, 0);
  display.println("Execute?");
  display.drawLine(0, 16, display.width(), 16, SSD1306_WHITE);

  display.setCursor(0, 50);
  display.print("NO");

  display.setCursor(92, 50);
  display.println("YES");

  display.display();

  if (b_con.GetLeftState())
  {
    b_con.Clear();
    screenStat = MAIN_MENU;
  }

  if (b_con.GetRightState())
  {
    b_con.Clear();
    
    display.clearDisplay();
    display.setCursor(12, 25);
    display.println("Homing...");
    display.display();

    bool home_ret = motor_control.HomeWidth(home_speed);
    if (!home_ret)
    {
      runtimeStat = ERROR_HM;
      screenStat = SHOW_SET_ROUTINE;
      return;
    }
    
    screenStat = MAIN_MENU;
  }
}

void running_routine()
{
  bool left_button = b_con.GetLeftState();
  bool right_button = b_con.GetRightState();
  bool mid_button = b_con.GetMidState();

  if (left_button)
  {
    menu_selector_index = 1;
    refreshes = 0;
    if (runtimeStat == OPERATING)
    {
      runtimeStat = SETTING;
      digitalWrite(FEEDBACK_PIN, LOW);
      screenStat = MAIN_MENU;
    }

    b_con.Clear();
    return;
  }

  // Clear error
  if (right_button && (runtimeStat != OPERATING || runtimeStat != SETTING))
  {
    runtimeStat = OPERATING;
    screenStat = SHOW_SET_ROUTINE;
    refreshes = 0;
    digitalWrite(FEEDBACK_PIN, HIGH);
    b_con.Clear();
    return;
  }

  // Checking refreshes for screen shutdown
  if (refreshes > REFRESHES_BEFORE_SCREEN_OFF && SCREEN_SAVER)
  {
    screenStat = OFF_SCREEN;
    refreshes = 0;
    display.clearDisplay();
    display.display();
    Serial.println("Screen: Off");
  }
        
  refreshes++;
}
