#include <Arduino.h>
#include <ESP8266mDNS.h>
#include <PubSubClient.h>
#include <WiFiManager.h>
#include <FastLED.h>
#include <ESP8266httpUpdate.h>

#define _STR(s) #s
#define STR(s) _STR(s)

const char* mqtt_server = "192.168.1.2"; //MQTT Server IP, your home MQTT server eg Mosquitto on RPi, or some public MQTT
const int mqtt_port     = 1883; //MQTT Server PORT, default is 1883 but can be anything.

const char* mqtt_pingall_sub = "home/pingall";
const char* mqtt_pingall_pub = "home/pingall/response";

#define MQTT_BASE "home/diningroom/desk"
const char* mqtt_device_boot = MQTT_BASE "/device";

const char* mqtt_mode_sub = MQTT_BASE;
const char* mqtt_mode_pub = MQTT_BASE "/status";

const char* mqtt_color_sub = MQTT_BASE "/color";
const char* mqtt_color_pub = MQTT_BASE "/color/status";

const char* mqtt_brightness_sub = MQTT_BASE "/brightness";
const char* mqtt_brightness_pub = MQTT_BASE "/brightness/status";

#define RGB_PIN      2
#define RGB_NUM_LEDS 12
#define RGB_CHIPSET  WS2812B
#define RGB_COLOR_ORDER GRB
#define NUM(a) (sizeof(a) / sizeof(*a))

#define FIRMWARE_VERSION  1
//#define FIRMWARE_URL    ""

WiFiClient espClient;
PubSubClient client(espClient);
char convBuffer[10];
CRGB leds[RGB_NUM_LEDS];

void callback(char *topic, uint8_t *payload, unsigned int length);
void blink();
void checkFirmwareUpdate();

void switchMode(struct mode *mode);
void modeOff();
void modeSolid();
void modeRainbow();
void modeRainbowFast();
void modeStrobo();

struct mode {
  const char *name;
  void (*func)();
};

struct mode modes[] = {
  { "off",         modeOff         },
  { "solid",       modeSolid       },
  { "rainbow",     modeRainbow     },
  { "rainbowfast", modeRainbowFast },
  { "strobo",      modeStrobo      },
};

struct
{
  struct mode *mode = &modes[0];
  uint32_t color = CRGB::Red;
  uint8_t brightness = 204; // 80%
} set;

struct
{
  uint8_t brightness;
  bool idle = false;
} current;

void setup()
{
  Serial.begin(115200);
  pinMode(LED_BUILTIN, OUTPUT);

  FastLED.addLeds<RGB_CHIPSET, RGB_PIN, RGB_COLOR_ORDER>(leds, RGB_NUM_LEDS);
  current.brightness = set.brightness;
  // update leds
  set.mode->func();
  FastLED.show();

  WiFiManager wifiManager;
  wifiManager.setConfigPortalTimeout(600);
  Serial.printf_P(PSTR("Setting up WiFi\n"));
  if (!wifiManager.autoConnect("ESP8266_DR_DESK")) {
    Serial.printf_P(PSTR("Failed to connect and hit timeout\n"));
    delay(5000);
    ESP.restart();
  }

  yield();
  checkFirmwareUpdate();
  yield();

  client.setServer(mqtt_server, mqtt_port);
  client.setCallback(callback);

  digitalWrite(LED_BUILTIN, HIGH);  //Turn off led as default
}

void reconnect()
{
  // Loop until we're reconnected
  while (!client.connected())
  {
    // Create a random client ID
    String clientId = "ESP8266Client-";
    clientId += String(random(0xffff), HEX);

    // Attempt to connect
    if (client.connect(clientId.c_str()))
    {
      // Once connected, publish an announcement...
      client.publish(mqtt_device_boot, "connected");
      // ... and resubscribe
      client.subscribe(mqtt_pingall_sub);

      // publish states
      Serial.println(mqtt_mode_sub);
      client.subscribe(mqtt_mode_sub);
      client.publish(mqtt_mode_pub, set.mode->name, true);

      Serial.println(mqtt_color_sub);
      client.subscribe(mqtt_color_sub);
      itoa(set.color, convBuffer, 10);
      client.publish(mqtt_color_pub, convBuffer, true);

      Serial.println(mqtt_brightness_sub);
      client.subscribe(mqtt_brightness_sub);
      itoa(map(set.brightness, 0, 255, 0, 100), convBuffer, 10);
      client.publish(mqtt_brightness_pub, convBuffer, true);
    }
    else
    {
      // Wait 5 seconds before retrying
      delay(5000);
    }
  }
}

void callback(char *topic, uint8_t *payload, unsigned int length)
{
  uint8_t c_payload[length];
  memcpy(c_payload, payload, length);
  c_payload[length] = '\0';

  if (strcmp(topic, mqtt_pingall_sub) == 0)
  {
    blink();
    client.publish(mqtt_pingall_pub,
      "{\"diningroom_desk\":\"connected\"}");
  }
  else if (strcmp(topic, mqtt_mode_sub) == 0)
  {
    for (uint8_t i = 0; i < NUM(modes); ++i)
    {
      if (strcmp(modes[i].name, (char *)c_payload) != 0)
        continue;
      switchMode(&modes[i]);
      break;
    }
  }
  else if (strcmp(topic, mqtt_color_sub) == 0)
  {
    // switch from OFF to SOLID
    if (set.mode == &modes[0])
      switchMode(&modes[1]);

    set.color = atoi((char *)c_payload);
    current.idle = false;
    itoa(set.color, convBuffer, 10);
    client.publish(mqtt_color_pub, convBuffer, true);
  }
  else if (strcmp(topic, mqtt_brightness_sub) == 0)
  {
    set.brightness = map(atoi((char *)c_payload), 0, 100, 0, 255);
    current.idle = false;
    itoa(map(set.brightness, 0, 255, 0, 100), convBuffer, 10);
    client.publish(mqtt_brightness_pub, convBuffer, true);
  }
}

void blink()
{
  //Blink on received MQTT message
  digitalWrite(LED_BUILTIN, LOW);
  delay(25);
  digitalWrite(LED_BUILTIN, HIGH);
}

void calcBrightness()
{
  #define FADE_STEP 10
  if (current.brightness == set.brightness)
    return;
  int fadeAmount = set.brightness - current.brightness;
  if (abs(fadeAmount) > FADE_STEP)
    fadeAmount = (fadeAmount > 0) ? FADE_STEP : -FADE_STEP;
  current.brightness += fadeAmount;
}

void switchMode(struct mode *mode)
{
  if (set.mode == mode)
    return;
  Serial.print("Switching mode to ");
  Serial.println(mode->name);
  set.mode = mode;
  current.idle = false;
  client.publish(mqtt_mode_pub, set.mode->name, true);
}

void modeOff()
{
  fill_solid(leds, RGB_NUM_LEDS, CRGB::Black);
  current.idle = true;
}

void modeSolid()
{
  fill_solid(leds, RGB_NUM_LEDS, set.color);
  calcBrightness();
  nscale8_video(leds, RGB_NUM_LEDS, current.brightness);
  current.idle = (current.brightness == set.brightness);
}

void modeRainbow()
{
  static CRGBPalette16 palette = RainbowColors_p;
  static uint8_t hue = 0;
  calcBrightness();
  for (int i = 0; i < RGB_NUM_LEDS; ++i)
    leds[i] = ColorFromPalette(palette, hue, current.brightness);
  hue++;
}

void modeRainbowFast()
{
  static CRGBPalette16 palette = RainbowColors_p;
  static uint8_t hue = 0;
  calcBrightness();
  for (int i = 0; i < RGB_NUM_LEDS; ++i)
    leds[i] = ColorFromPalette(palette, hue, current.brightness);
  hue+=5;
}

void modeStrobo()
{
  static bool state = 0;
  if (set.color == CRGB::Black)
    return;
  fill_solid(leds, RGB_NUM_LEDS, (state) ? set.color : CRGB::Black);
  calcBrightness();
  nscale8_video(leds, RGB_NUM_LEDS, current.brightness);
  state = !state;
}

void checkFirmwareUpdate()
{
  BearSSL::WiFiClientSecure update_client;
  update_client.setInsecure();

  client.publish(mqtt_device_boot, "fwupdate running");
  ESPhttpUpdate.setLedPin(LED_BUILTIN, HIGH);
  ESPhttpUpdate.rebootOnUpdate(true);
  t_httpUpdate_return ret = ESPhttpUpdate.update(update_client, FIRMWARE_URL, STR(FIRMWARE_VERSION));
  switch(ret)
  {
    case HTTP_UPDATE_FAILED:
      {
        Serial.printf_P(PSTR("HTTP_UPDATE_FAILED Error (%d): %s\n"),
          ESPhttpUpdate.getLastError(), ESPhttpUpdate.getLastErrorString().c_str());
        String tmp = String("fwupdate error: ") + ESPhttpUpdate.getLastErrorString();
        client.publish(mqtt_device_boot, tmp.c_str());
      }
      break;

    case HTTP_UPDATE_NO_UPDATES:
      Serial.printf_P(PSTR("HTTP_UPDATE_NO_UPDATES\n"));
      client.publish(mqtt_device_boot, "fwupdate noupdates");
      break;

    case HTTP_UPDATE_OK:
      Serial.printf_P(PSTR("HTTP_UPDATE_OK\n"));
      client.publish(mqtt_device_boot, "fwupdate ok");
      break;
  }
}

void loop()
{
  if (!client.connected())
    reconnect();
  client.loop();

  EVERY_N_MILLISECONDS(100)
  {
    if (!current.idle)
    {
      set.mode->func();
      FastLED.show();
    }
    if (!current.idle && current.brightness == 0)
      current.idle = true;
  }
}