#define LED_COUNT 10
#define LED_BRIGHTNESS 100
#define LED_PIN 12

#define WIFI_SSID ""
#define WIFI_PASS ""

#define OBS_HOST ""
#define OBS_PORT 4444
//#define OBS_PASS ""

#define OBS_SOURCE "ATEM"

#define FASTLED_ESP8266_RAW_PIN_ORDER

// END CONFIGURATION

#include <Arduino.h>
#include <Hash.h>
#include <ESP8266WiFi.h>
#include <WebSocketsClient.h>
#include <ArduinoJson.h>
#include <Base64.h>
#include <FastLED.h>

CRGB leds[LED_COUNT];
WebSocketsClient webSocket;

#ifdef OBS_PASS
#include <SHA256.h>
#define HASH_SIZE 32

SHA256 sha256;
#endif

bool is_currently_live = false;
bool is_currently_preview = false;

void set_program() {
    Serial.println("[Tally] PROGRAM");
    fill_solid(leds, LED_COUNT, CRGB::Red);
    FastLED.show();
}

void set_preview() {
    Serial.println("[Tally] PREVIEW");
    fill_solid(leds, LED_COUNT, CRGB::Green);
    FastLED.show();
}

void set_idle() {
    Serial.println("[Tally] IDLE");
    fill_solid(leds, LED_COUNT, CRGB::Black);
    FastLED.show();
}

void set_error() {
    Serial.println("[Tally] ERROR");
    fill_solid(leds, LED_COUNT, CRGB::Purple);
    FastLED.show();
}

void handleWebSocketEvent(WStype_t type, uint8_t * payload, size_t length) {
    switch(type) {
        case WStype_DISCONNECTED:
            Serial.println("[WS] disconnected");
            /*
                We do NOT set the LEDs to off in here on purpose. Maybe
                just the wifi connection was interrupted or something
                like that. It is important we don't disable the LEDs, in
                case the OBS instance is still running.

                We can't set the LEDs to "error", because if someone
                quits OBS, we will get disconnected.
            */
            break;

        case WStype_CONNECTED:
            Serial.printf("[WS] connected to %s\n", payload);
#ifdef OBS_PASS
            // Find out if we need authentication
            webSocket.sendTXT("{\"request-type\":\"GetAuthRequired\",\"message-id\":\"1\"}");
#else
            webSocket.sendTXT("{\"request-type\":\"GetCurrentScene\",\"message-id\":\"3\"}");
            webSocket.sendTXT("{\"request-type\":\"GetPreviewScene\",\"message-id\":\"4\"}");
#endif
            break;

        case WStype_TEXT: {
            Serial.printf("[WS] %s\n", payload);

            StaticJsonDocument<5000> doc;
            DeserializationError error = deserializeJson(doc, payload);

            if (error) {
                Serial.print("deserializeJson() failed: ");
                Serial.println(error.c_str());
                break;
            }

#ifdef OBS_PASS
            if (doc.containsKey("authRequired")) {
                Serial.println("[OBS] auth requested");

                sha256.reset();
                sha256.update(OBS_PASS, strlen(OBS_PASS));
                const char* salt = doc["salt"];
                sha256.update(salt, strlen(salt));
                char value[HASH_SIZE];
                sha256.finalize(value, HASH_SIZE);

                Serial.print("[OBS] sha256 authentication hash is: ");
                for (size_t i = 0; i < HASH_SIZE; i++) {
                    Serial.print(static_cast<unsigned int>(value[i]), HEX);
                }
                Serial.println();

                int encodedLength = Base64.encodedLength(HASH_SIZE);
                char encodedPassSaltHash[encodedLength];
                Base64.encode(encodedPassSaltHash, value, HASH_SIZE);

                const char* challenge = doc["challenge"];
                sha256.reset();
                sha256.update(encodedPassSaltHash, encodedLength);
                sha256.update(challenge, strlen(challenge));
                sha256.finalize(value, HASH_SIZE);

                Serial.print("[OBS] sha256 challenge hash is: ");
                for (size_t i = 0; i < HASH_SIZE; i++) {
                    Serial.print(static_cast<unsigned int>(value[i]), HEX);
                }
                Serial.println();

                char encodedAuthString[encodedLength];
                Base64.encode(encodedAuthString, value, HASH_SIZE);

                String authRequest = String("{\"request-type\":\"Authenticate\",\"message-id\":\"2\",\"auth\":\"") + encodedAuthString + "\"}";
                webSocket.sendTXT(authRequest);
                break;

            } else if (doc.containsKey("message-id") && doc["message-id"] == "2") {
                if (strcmp(doc["status"], "ok") == 0) {
                    Serial.println("[OBS] authentication successful");
                    webSocket.sendTXT("{\"request-type\":\"GetCurrentScene\",\"message-id\":\"3\"}");
                    webSocket.sendTXT("{\"request-type\":\"GetPreviewScene\",\"message-id\":\"4\"}");
                } else {
                    Serial.println("[OBS] authentication FAILED");
                    set_error();
                }
                break;
            }
#endif
            if (doc.containsKey("sources")) {
                bool my_source_in_current_event = false;

                for (uint8_t i = 0; i < doc["sources"].size(); i++) {
                    if (strcmp(doc["sources"][i]["name"], OBS_SOURCE) == 0) {
                        Serial.println("[OBS] Source found in current event");
                        my_source_in_current_event = true;
                    }
                }

                if ((doc.containsKey("update-type") &&strcmp(doc["update-type"], "SwitchScenes") == 0) || (doc.containsKey("message-id") && strcmp(doc["message-id"], "3") == 0)) {
                    Serial.println("[OBS] program event");
                    if (my_source_in_current_event) {
                        is_currently_live = true;
                    } else {
                        is_currently_live = false;
                    }
                } else if ((doc.containsKey("update-type") &&strcmp(doc["update-type"], "PreviewSceneChanged") == 0) || (doc.containsKey("message-id") && strcmp(doc["message-id"], "4") == 0)) {
                    Serial.println("[OBS] preview event");
                    if (my_source_in_current_event) {
                        is_currently_preview = true;
                    } else {
                        is_currently_preview = false;
                    }
                }
            } else if (doc.containsKey("update-type")) {
                if (strcmp(doc["update-type"], "Exiting") == 0) {
                    Serial.println("[OBS] quit");
                    is_currently_preview = false;
                    is_currently_live = false;
                } else if (strcmp(doc["update-type"], "StudioModeSwitched") == 0 && !doc["new-state"]) {
                    Serial.println("[OBS] studio mode disabled");
                    is_currently_preview = false;
                }
            }

            if (is_currently_live) {
                set_program();
            } else if (is_currently_preview) {
                set_preview();
            } else {
                set_idle();
            }

            break;
        }

        case WStype_PING:
        case WStype_PONG:
            Serial.println("[WS] ping/pong");
            break;

        case WStype_BIN:
        case WStype_ERROR:
        case WStype_FRAGMENT_TEXT_START:
        case WStype_FRAGMENT_BIN_START:
        case WStype_FRAGMENT:
        case WStype_FRAGMENT_FIN:
            break;
    }
}

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

    Serial.println("[Tally] starting up");

    FastLED.addLeds<WS2812B, LED_PIN, GRB>(leds, LED_COUNT);
    FastLED.setBrightness(LED_BRIGHTNESS);

    set_error();
    delay(100);
    set_idle();
    delay(100);
    set_error();
    delay(100);
    set_idle();
    delay(100);
    set_error();

    Serial.printf("[Tally] connecting to wifi ssid: %s\n", WIFI_SSID);

    WiFi.mode(WIFI_STA);
    WiFi.begin(WIFI_SSID, WIFI_PASS);

    while(WiFi.status() != WL_CONNECTED) {
        delay(250);
        Serial.print(".");
    }
    Serial.println();
    Serial.print("[Tally] connected to wifi, ip address ");
    Serial.println(WiFi.localIP());

    set_idle();

    Serial.println("[Tally] connecting to OBS");
    webSocket.begin(OBS_HOST, OBS_PORT, "/");
    webSocket.onEvent(handleWebSocketEvent);
    webSocket.setReconnectInterval(2000);

    Serial.println("[Tally] startup complete");
}

void loop() {
    webSocket.loop();
}