#include <WiFi.h>
#include <WiFiUdp.h>
#include <WebServer.h>
#include <ESPmDNS.h>
#include <Preferences.h>
#include <Adafruit_NeoPixel.h>

#define STRIP1_PIN 18
#define STRIP2_PIN 19
#define MAX_LEDS   150
#define UDP_PORT   6969
#define WATCHDOG_TIMEOUT_MS 2000
#define MAX_PACKET_BYTES (4 + MAX_LEDS * 3 + MAX_LEDS * 3)

// ---------------------------------------------------------------------
// All type definitions live up here, before any function -- Arduino's
// build system auto-generates function prototypes and inserts them
// right after the #includes, ahead of anything else in the file. If a
// prototype references a type (enum/struct) that's defined further
// down, that generated prototype fails to compile even though the
// real function definition textually comes after the type. Keeping
// every enum/struct here avoids that class of error entirely.
// ---------------------------------------------------------------------

// MODE_COLOR (static on, adjustable color/brightness) and MODE_OFF
// (genuinely dark) are deliberately separate states -- see
// handleOn()/handleColor()/handleOff() below.
enum Mode { MODE_TEST_CYCLE, MODE_MUSIC_SYNC, MODE_COLOR, MODE_OFF };

struct ColorStep {
  uint8_t r, g, b;
  unsigned long durationMs;
  const char *label;
};

// Wire format for a music-sync frame from strip_music_sync.py on the
// Jetson: <sequence:u16><len1:u8><len2:u8> then len1*3 bytes of RGB
// for strip1, then len2*3 bytes of RGB for strip2. This board does no
// animation math at all for this path -- every render function
// (VU meter, mirror wave, neon racer, fire, ...) lives in Python on
// the Jetson (strip_animations.py) precisely so iterating on them is
// a service restart, not a firmware reflash + physically swapping
// this board between the Jetson and the strips. This board just
// copies whatever pixels arrive straight into the strips and shows
// them.
#pragma pack(push, 1)
struct PixelFrameHeader {
  uint16_t sequence;
  uint8_t len1;
  uint8_t len2;
};
#pragma pack(pop)

// ---------------------------------------------------------------------
// Globals
// ---------------------------------------------------------------------

const char *WIFI_SSID = "VISHAL";
const char *WIFI_PASSWORD = "9971978446@2";
const char *HOSTNAME = "q-strip-test";

Adafruit_NeoPixel strip1(MAX_LEDS, STRIP1_PIN, NEO_GRB + NEO_KHZ800);
Adafruit_NeoPixel strip2(MAX_LEDS, STRIP2_PIN, NEO_GRB + NEO_KHZ800);
WebServer server(80);
Preferences prefs;
WiFiUDP udp;

Mode currentMode = MODE_TEST_CYCLE;

unsigned long lastStepTime = 0;
int stepIndex = 0;

// How many LEDs the standalone test cycle (MODE_TEST_CYCLE) and the
// static-color mode (MODE_COLOR) address on each strip -- adjustable
// at runtime from the dashboard so it can be tuned without a reflash.
// The music-sync pixel-frame path doesn't consult this at all, it
// trusts each packet's own header instead, since the Jetson already
// knows the real lengths. Persisted to flash (Preferences/NVS) so a
// power-cycle doesn't drop it.
int len1 = 30;
int len2 = 15;

// Last-set solid color + brightness for MODE_COLOR (voice/dashboard
// "turn on"/"set color"/"set brightness" commands) -- persisted so
// "turn on" after a power-cycle resumes whatever was last showing
// instead of always resetting to white.
uint8_t currentR = 255;
uint8_t currentG = 255;
uint8_t currentB = 255;
uint8_t currentBrightness = 80;

ColorStep steps[] = {
  {255, 0, 0, 1500, "RED"},
  {0, 255, 0, 1500, "GREEN"},
  {0, 0, 255, 1500, "BLUE"},
  {255, 255, 255, 1500, "WHITE"},
  {0, 0, 0, 800, "OFF"},
};
const int numSteps = sizeof(steps) / sizeof(steps[0]);

unsigned long lastMusicPacketMs = 0;
unsigned long lastBlankMs = 0;
uint8_t udpBuffer[MAX_PACKET_BYTES];

// ---------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------

void saveLengths() {
  prefs.putInt("len1", len1);
  prefs.putInt("len2", len2);
}

void saveColorState() {
  prefs.putUChar("curR", currentR);
  prefs.putUChar("curG", currentG);
  prefs.putUChar("curB", currentB);
  prefs.putUChar("bright", currentBrightness);
}

void setSolidColor(uint8_t r, uint8_t g, uint8_t b) {
  for (int i = 0; i < MAX_LEDS; i++) {
    strip1.setPixelColor(i, i < len1 ? strip1.Color(r, g, b) : 0);
    strip2.setPixelColor(i, i < len2 ? strip2.Color(r, g, b) : 0);
  }
  strip1.show();
  strip2.show();
}

void applyCurrentColor() {
  setSolidColor(currentR, currentG, currentB);
}

// ---------------------------------------------------------------------
// Test cycle -- the original bring-up tool, unchanged behavior, now
// just one of several top-level modes instead of the only thing this
// firmware could do.
// ---------------------------------------------------------------------

void runTestCycleFrame() {
  unsigned long now = millis();

  if (now - lastStepTime >= steps[stepIndex].durationMs) {
    lastStepTime = now;
    ColorStep &s = steps[stepIndex];
    Serial.print("Showing: ");
    Serial.println(s.label);
    setSolidColor(s.r, s.g, s.b);
    stepIndex = (stepIndex + 1) % numSteps;
  }
}

// ---------------------------------------------------------------------
// Music sync -- pure pixel-frame playback, no animation logic here.
// ---------------------------------------------------------------------

void pollPixelFrame() {
  int packetSize = udp.parsePacket();

  if (packetSize < (int)sizeof(PixelFrameHeader)) {
    return;
  }

  if (packetSize > MAX_PACKET_BYTES) {
    packetSize = MAX_PACKET_BYTES;
  }

  int readBytes = udp.read(udpBuffer, packetSize);

  if (readBytes < (int)sizeof(PixelFrameHeader)) {
    return;
  }

  PixelFrameHeader header;
  memcpy(&header, udpBuffer, sizeof(header));

  int n1 = min((int)header.len1, MAX_LEDS);
  int n2 = min((int)header.len2, MAX_LEDS);
  int expectedBytes = sizeof(PixelFrameHeader) + n1 * 3 + n2 * 3;

  if (readBytes < expectedBytes) {
    return;  // malformed/truncated -- ignore rather than render garbage
  }

  int offset = sizeof(PixelFrameHeader);

  for (int i = 0; i < n1; i++) {
    strip1.setPixelColor(i, strip1.Color(udpBuffer[offset], udpBuffer[offset + 1], udpBuffer[offset + 2]));
    offset += 3;
  }
  for (int i = n1; i < MAX_LEDS; i++) {
    strip1.setPixelColor(i, 0);
  }

  for (int i = 0; i < n2; i++) {
    strip2.setPixelColor(i, strip2.Color(udpBuffer[offset], udpBuffer[offset + 1], udpBuffer[offset + 2]));
    offset += 3;
  }
  for (int i = n2; i < MAX_LEDS; i++) {
    strip2.setPixelColor(i, 0);
  }

  strip1.show();
  strip2.show();
  lastMusicPacketMs = millis();
}

// ---------------------------------------------------------------------
// HTTP control plane -- bring-up endpoints, music-sync mode switch,
// and on/off/brightness/color for voice/dashboard control.
// ---------------------------------------------------------------------

void handleStatus() {
  String json = "{\"online\":true,\"mode\":\"";

  switch (currentMode) {
    case MODE_TEST_CYCLE: json += "test_cycle"; break;
    case MODE_MUSIC_SYNC: json += "music_sync"; break;
    case MODE_COLOR: json += "color"; break;
    default: json += "off"; break;
  }

  json += "\",\"len1\":" + String(len1);
  json += ",\"len2\":" + String(len2);
  json += ",\"r\":" + String(currentR);
  json += ",\"g\":" + String(currentG);
  json += ",\"b\":" + String(currentB);
  json += ",\"brightness\":" + String(currentBrightness);
  json += ",\"ip\":\"" + WiFi.localIP().toString() + "\"}";
  server.send(200, "application/json", json);
}

void handleOff() {
  currentMode = MODE_OFF;
  setSolidColor(0, 0, 0);
  server.send(200, "application/json", "{\"ok\":true}");
}

void handleTest() {
  currentMode = MODE_TEST_CYCLE;
  stepIndex = 0;
  lastStepTime = 0;
  server.send(200, "application/json", "{\"ok\":true}");
}

void handleMusicSync() {
  currentMode = MODE_MUSIC_SYNC;
  lastMusicPacketMs = millis();  // don't let the watchdog trip before the first real packet arrives
  server.send(200, "application/json", "{\"ok\":true}");
}

// Turns the wall light on -- resumes the last color/brightness if no
// r/g/b given (a plain voice "turn on the wall light"), or updates
// the color first if provided (e.g. "turn on the wall light in blue").
void handleOn() {
  if (server.hasArg("r") && server.hasArg("g") && server.hasArg("b")) {
    currentR = constrain(server.arg("r").toInt(), 0, 255);
    currentG = constrain(server.arg("g").toInt(), 0, 255);
    currentB = constrain(server.arg("b").toInt(), 0, 255);
  }

  currentMode = MODE_COLOR;
  saveColorState();
  applyCurrentColor();
  server.send(200, "application/json", "{\"ok\":true}");
}

void handleColor() {
  if (!server.hasArg("r") || !server.hasArg("g") || !server.hasArg("b")) {
    server.send(400, "application/json", "{\"error\":\"need r,g,b query params\"}");
    return;
  }

  currentR = constrain(server.arg("r").toInt(), 0, 255);
  currentG = constrain(server.arg("g").toInt(), 0, 255);
  currentB = constrain(server.arg("b").toInt(), 0, 255);
  currentMode = MODE_COLOR;
  saveColorState();
  applyCurrentColor();
  server.send(200, "application/json", "{\"ok\":true}");
}

void handleBrightness() {
  if (!server.hasArg("value")) {
    server.send(400, "application/json", "{\"error\":\"need value param (0-255)\"}");
    return;
  }

  currentBrightness = constrain(server.arg("value").toInt(), 0, 255);
  strip1.setBrightness(currentBrightness);
  strip2.setBrightness(currentBrightness);
  saveColorState();

  if (currentMode == MODE_COLOR) {
    applyCurrentColor();  // re-show so the brightness change is visible immediately
  }

  server.send(200, "application/json", "{\"ok\":true,\"brightness\":" + String(currentBrightness) + "}");
}

void handleConfig() {
  if (server.hasArg("len1")) {
    int v = server.arg("len1").toInt();
    len1 = constrain(v, 0, MAX_LEDS);
  }
  if (server.hasArg("len2")) {
    int v = server.arg("len2").toInt();
    len2 = constrain(v, 0, MAX_LEDS);
  }
  saveLengths();
  server.send(200, "application/json", "{\"ok\":true,\"len1\":" + String(len1) + ",\"len2\":" + String(len2) + "}");
}

void handleReboot() {
  server.send(200, "application/json", "{\"ok\":true}");
  delay(200);
  ESP.restart();
}

void setup() {
  Serial.begin(115200);
  delay(300);
  Serial.println();
  Serial.println("Booting strip firmware...");

  prefs.begin("striptest", false);
  len1 = prefs.getInt("len1", len1);
  len2 = prefs.getInt("len2", len2);
  currentR = prefs.getUChar("curR", currentR);
  currentG = prefs.getUChar("curG", currentG);
  currentB = prefs.getUChar("curB", currentB);
  currentBrightness = prefs.getUChar("bright", currentBrightness);
  Serial.print("Loaded saved lengths: len1=");
  Serial.print(len1);
  Serial.print(" len2=");
  Serial.print(len2);
  Serial.print(" color=(");
  Serial.print(currentR);
  Serial.print(",");
  Serial.print(currentG);
  Serial.print(",");
  Serial.print(currentB);
  Serial.print(") brightness=");
  Serial.println(currentBrightness);

  strip1.begin();
  strip2.begin();
  strip1.setBrightness(currentBrightness);
  strip2.setBrightness(currentBrightness);
  setSolidColor(0, 0, 0);

  WiFi.mode(WIFI_STA);
  WiFi.setHostname(HOSTNAME);
  WiFi.begin(WIFI_SSID, WIFI_PASSWORD);

  Serial.print("Connecting to WiFi");
  unsigned long wifiStart = millis();
  while (WiFi.status() != WL_CONNECTED && millis() - wifiStart < 20000) {
    delay(300);
    Serial.print(".");
  }
  Serial.println();

  if (WiFi.status() == WL_CONNECTED) {
    Serial.print("WiFi connected. IP address: ");
    Serial.println(WiFi.localIP());

    if (MDNS.begin(HOSTNAME)) {
      Serial.print("mDNS responder started: http://");
      Serial.print(HOSTNAME);
      Serial.println(".local");
      MDNS.addService("http", "tcp", 80);
    }

    udp.begin(UDP_PORT);
    Serial.print("UDP pixel-frame listener started on port ");
    Serial.println(UDP_PORT);
  } else {
    Serial.println("WiFi connect FAILED - continuing offline, test cycle will still run.");
  }

  server.on("/status", HTTP_GET, handleStatus);
  server.on("/off", HTTP_POST, handleOff);
  server.on("/on", HTTP_POST, handleOn);
  server.on("/test", HTTP_POST, handleTest);
  server.on("/music-sync", HTTP_POST, handleMusicSync);
  server.on("/color", HTTP_POST, handleColor);
  server.on("/brightness", HTTP_POST, handleBrightness);
  server.on("/config", HTTP_POST, handleConfig);
  server.on("/reboot", HTTP_POST, handleReboot);
  server.begin();
  Serial.println("HTTP server started.");
}

void loop() {
  server.handleClient();

  switch (currentMode) {
    case MODE_TEST_CYCLE:
      runTestCycleFrame();
      break;

    case MODE_MUSIC_SYNC:
      pollPixelFrame();

      // A data gap goes dark but deliberately does NOT leave
      // MODE_MUSIC_SYNC -- it keeps polling, so the moment a fresh
      // frame arrives it resumes automatically with no need to
      // re-POST /music-sync. Only an explicit /off or /test call
      // actually changes currentMode -- that's the sticky,
      // user-controlled state; a transient stream dropout is not.
      if (millis() - lastMusicPacketMs > WATCHDOG_TIMEOUT_MS) {
        // Non-blocking rate limit -- avoid re-blanking (and the
        // strip1.show()/strip2.show() cost that comes with it) every
        // single loop() iteration while genuinely disconnected,
        // without stalling HTTP handling via a blocking delay().
        if (millis() - lastBlankMs >= 200) {
          lastBlankMs = millis();
          setSolidColor(0, 0, 0);
        }
      }
      break;

    case MODE_COLOR:
    case MODE_OFF:
    default:
      break;  // static state, already rendered when it was set
  }
}