#include <WiFi.h>
#include <WiFiUdp.h>
#include <FastLED.h>
#include <cstring>   // memcmp, memset

// ====================== WiFi ======================
const char* ssid     = "YOUR_WIFI_SSID";
const char* password = "YOUR_WIFI_PASSWORD";

// ====================== LED LINES ======================
// Up to ~8 lines recommended on ESP32 (RMT)
#define NUM_LINES 4

// Hard-code your line pins (compile-time constants required by FastLED)
#define LINE0_PIN 2
#define LINE1_PIN 4
#define LINE2_PIN 5
#define LINE3_PIN 18
// Add LINE4_PIN... if you increase NUM_LINES, and mirror in addLeds below.

// Per-line pixel counts (edit to match your install)
const uint16_t LINE_LENGTHS[NUM_LINES] = {
  300, 300, 300, 300  // total = 1200 (example)
};

#define LED_TYPE     WS2812B   // NeoPixel strips
#define COLOR_ORDER  GRB
#define BRIGHTNESS   255

// Per-line pixel buffers (allocated in setup)
CRGB*     lineLeds[NUM_LINES];           // pointers to per-line arrays
uint32_t  lineOffsets[NUM_LINES + 1];    // cumulative virtual offsets (last = total)
uint32_t  VIRTUAL_TOTAL = 0;             // sum of LINE_LENGTHS

// ====================== E1.31 / sACN ======================
#define E131_PORT           5568
#define START_UNIVERSE      1      // first universe to listen to
#define UNIVERSE_COUNT      8      // how many universes (contiguous)
#define DMX_CHANS_PER_UNI   510    // 510 typical for 170 RGB pixels
#define DMX_START_ADDRESS   1      // 1..512 offset inside each universe
#define USE_MULTICAST       1      // 1 = multicast (recommended), 0 = unicast

static const uint16_t PIXELS_PER_UNI = DMX_CHANS_PER_UNI / 3;
static const uint32_t FRAME_MASK_WANTED = (UNIVERSE_COUNT >= 32) ? 0xFFFFFFFFu : ((1u << UNIVERSE_COUNT) - 1u);

// sACN buffer (full DMX payload + headers; 638 is safe)
static const int MAX_E131_PACKET = 115 + 638;
uint8_t packetBuffer[MAX_E131_PACKET];

#if USE_MULTICAST
WiFiUDP udpSockets[UNIVERSE_COUNT];
#else
WiFiUDP udp;
#endif

// ---- E1.31 constants / offsets ----
static const uint8_t ACN_PID[12] = { 0x41,0x53,0x43,0x2d,0x45,0x31,0x2e,0x31,0x37,0x00,0x00,0x00 }; // "ASC-E1.17\0\0\0"
enum : uint32_t {
  VECTOR_ROOT_E131_DATA   = 0x00000004,
  VECTOR_E131_DATA_PACKET = 0x00000002,
};
enum : uint8_t {
  VECTOR_DMP_SET_PROPERTY = 0x02
};

// Sequence/frame sync (update once a full frame is in)
static uint8_t  currentSeq = 0;
static bool     haveSeq = false;
static uint32_t receivedMask = 0;

// ====================== Fallback / Effects ======================
#define WIFI_CONNECT_TIMEOUT_MS  12000   // give Wi-Fi this long on boot
#define SACN_TIMEOUT_MS           2000   // if no sACN in this long, enter fallback
#define EFFECT_FRAME_MS             18   // ~55 FPS effects engine
#define EFFECT_DWELL_MS          12000   // how long to keep one effect before randomizing

bool     fallbackActive = false;
uint32_t lastPacketMs   = 0;
uint32_t lastEffectMs   = 0;
uint32_t effectStartMs  = 0;
uint8_t  currentEffect  = 0;  // which effect we’re running
uint16_t effectPhase    = 0;  // animation phase / step
uint8_t  baseHue        = 0;  // for rainbow-ish effects

// ====================== Helpers ======================
IPAddress universeToMulticastIP(uint16_t universe) {
  // Per E1.31: 239.255.<hi-byte>.<lo-byte>
  return IPAddress(239, 255, (universe >> 8) & 0xFF, universe & 0xFF);
}

static inline uint16_t readBE16(const uint8_t* p) {
  return (uint16_t(p[0]) << 8) | uint16_t(p[1]);
}
static inline uint32_t readBE32(const uint8_t* p) {
  return (uint32_t(p[0]) << 24) | (uint32_t(p[1]) << 16) | (uint32_t(p[2]) << 8) | uint32_t(p[3]);
}

bool checkACNPID(const uint8_t* buf, int len) {
  if (len < 16) return false;
  return memcmp(buf + 4, ACN_PID, 12) == 0;
}

// Parse sACN: extract universe, sequence, pointer to DMX slots and count (no start code)
bool parseE131(const uint8_t* buf, int len,
               uint16_t& outUniverse, uint8_t& outSeq,
               const uint8_t*& outSlots, uint16_t& outSlotCount) {
  if (len < 116) return false;
  if (!checkACNPID(buf, len)) return false;

  uint32_t rootVector = readBE32(buf + 18);
  if (rootVector != VECTOR_ROOT_E131_DATA) return false;

  uint32_t framingVector = readBE32(buf + 40);
  if (framingVector != VECTOR_E131_DATA_PACKET) return false;

  // Framing layer fields
  outSeq      = buf[111];
  outUniverse = readBE16(buf + 113);

  // DMP
  if (buf[117] != VECTOR_DMP_SET_PROPERTY) return false;

  uint16_t propValCount = readBE16(buf + 123); // includes start code
  const int propStart = 125;
  if (propValCount < 1 || (propStart + propValCount) > len) return false;

  const uint8_t* propValues = buf + propStart;
  const uint8_t startCode = propValues[0];
  if (startCode != 0x00) return false; // only DMX data

  outSlots     = propValues + 1;        // DMX slot 1
  outSlotCount = propValCount - 1;

  return true;
}

// Virtual index -> (line, pixel)
bool virtToLinePixel(uint32_t vIndex, uint16_t& outLine, uint16_t& outPixel) {
  if (vIndex >= VIRTUAL_TOTAL) return false;
  uint16_t line = 0;
  while ((line + 1) <= NUM_LINES && vIndex >= lineOffsets[line + 1]) line++;
  outLine = line;
  outPixel = (uint16_t)(vIndex - lineOffsets[line]);
  return true;
}

void setVirtualPixel(uint32_t vIndex, const CRGB &c) {
  uint16_t line, pix;
  if (virtToLinePixel(vIndex, line, pix)) {
    lineLeds[line][pix] = c;
  }
}

void fillAll(const CRGB &c) {
  for (uint16_t i = 0; i < NUM_LINES; i++) {
    for (uint16_t p = 0; p < LINE_LENGTHS[i]; p++) lineLeds[i][p] = c;
  }
}

void mapUniverseSlots(uint16_t universe, const uint8_t* slots, uint16_t slotCount) {
  if (universe < START_UNIVERSE || universe >= (START_UNIVERSE + UNIVERSE_COUNT)) return;

  // DMX start address within this universe (1-based -> 0-based)
  uint16_t slotOffset = (DMX_START_ADDRESS > 1) ? (DMX_START_ADDRESS - 1) : 0;
  if (slotOffset >= slotCount) return;
  slots     += slotOffset;
  slotCount -= slotOffset;

  uint16_t uniIndex = universe - START_UNIVERSE;
  uint32_t startPixelVirt = (uint32_t)uniIndex * PIXELS_PER_UNI;  // where this universe starts in the virtual strip

  uint32_t pixelsFromDMX = slotCount / 3;
  for (uint32_t i = 0; i < pixelsFromDMX; i++) {
    uint32_t vPix = startPixelVirt + i;
    if (vPix >= VIRTUAL_TOTAL) break;
    uint32_t base = i * 3;

    uint8_t r = slots[base + 0];
    uint8_t g = slots[base + 1];
    uint8_t b = slots[base + 2];
    setVirtualPixel(vPix, CRGB(r,g,b));
  }
}

void noteUniverseReceived(uint8_t seq, uint16_t universe) {
  if (universe < START_UNIVERSE || universe >= (START_UNIVERSE + UNIVERSE_COUNT)) return;
  uint16_t idx = universe - START_UNIVERSE;

  // any valid packet exits fallback immediately
  fallbackActive = false;
  lastPacketMs   = millis();

  if (!haveSeq || seq != currentSeq) {
    currentSeq = seq;
    haveSeq = true;
    receivedMask = 0;
  }

  receivedMask |= (1u << idx);

  if ((receivedMask & FRAME_MASK_WANTED) == FRAME_MASK_WANTED) {
    FastLED.show();        // push all lines via RMT
    receivedMask = 0;
  }
}

void showStartupChase() {
  // Quick visual check (limit first ~300 pixels per line)
  for (uint16_t pass = 0; pass < 3; pass++) {
    CRGB col = (pass == 0) ? CRGB::Red : (pass == 1) ? CRGB::Green : CRGB::Blue;
    for (uint16_t i = 0; i < NUM_LINES; i++) {
      uint16_t count = min<uint16_t>(LINE_LENGTHS[i], 300);
      for (uint16_t p = 0; p < count; p++) lineLeds[i][p] = col;
    }
    FastLED.show(); delay(200);
  }
  fillAll(CRGB::Black);
  FastLED.show();
}

// ====================== Effects Engine ======================
// 0: RainbowCycle, 1: TheaterChase, 2: Twinkle, 3: LarsonScanner, 4: ColorWipe
void pickRandomEffect() {
  currentEffect = random(0, 5);
  effectPhase   = 0;
  baseHue       = random8();
  effectStartMs = millis();
}

void effectRainbowCycle() {
  // smooth rainbow across virtual strip, animated by baseHue
  for (uint32_t v = 0; v < VIRTUAL_TOTAL; v++) {
    uint8_t hue = baseHue + (v * 256UL / max<uint32_t>(VIRTUAL_TOTAL, 1));
    setVirtualPixel(v, CHSV(hue, 255, 255));
  }
  baseHue += 1;
  FastLED.show();
}

void effectTheaterChase() {
  // 3-phase dotted chase
  uint8_t phase = (effectPhase / 2) % 3;
  CRGB col(CHSV(baseHue, 200, 255));
  fillAll(CRGB::Black);
  for (uint32_t v = phase; v < VIRTUAL_TOTAL; v += 3) {
    setVirtualPixel(v, col);
  }
  effectPhase++;
  baseHue += (effectPhase % 15 == 0); // slowly drift hue
  FastLED.show();
}

void effectTwinkle() {
  // random sparkles with gentle fade
  // fade
  for (uint16_t i = 0; i < NUM_LINES; i++) {
    for (uint16_t p = 0; p < LINE_LENGTHS[i]; p++) {
      lineLeds[i][p].nscale8(220); // fade by ~14%
    }
  }
  // new sparks
  for (uint8_t k = 0; k < 10; k++) {
    uint32_t v = random(VIRTUAL_TOTAL);
    setVirtualPixel(v, CHSV(baseHue + random8(64), 200, 255));
  }
  FastLED.show();
}

void effectLarson() {
  // single scanner dot with trailing fade (Cylon)
  static int32_t pos = 0;
  static int8_t  dir = 1;
  // fade
  for (uint16_t i = 0; i < NUM_LINES; i++) {
    for (uint16_t p = 0; p < LINE_LENGTHS[i]; p++) {
      lineLeds[i][p].nscale8(200);
    }
  }
  // dot
  CRGB col(CHSV(baseHue, 255, 255));
  setVirtualPixel(pos, col);
  FastLED.show();

  // move
  pos += dir;
  if (pos <= 0) { pos = 0; dir = 1; baseHue += 16; }
  if ((uint32_t)pos >= VIRTUAL_TOTAL - 1) { pos = VIRTUAL_TOTAL - 1; dir = -1; baseHue += 16; }
}

void effectColorWipe() {
  // fill from start to end, then change color
  CRGB col(CHSV(baseHue, 240, 255));
  uint32_t step = effectPhase % (VIRTUAL_TOTAL + 20);
  if (step == 0) fillAll(CRGB::Black);
  if (step < VIRTUAL_TOTAL) {
    setVirtualPixel(step, col);
  } else if (step == VIRTUAL_TOTAL) {
    // complete wipe; hold a moment, then pick a new hue
    baseHue += 40;
  }
  effectPhase++;
  FastLED.show();
}

void runEffects() {
  uint32_t now = millis();

  // change effect occasionally
  if ((now - effectStartMs) >= EFFECT_DWELL_MS) {
    pickRandomEffect();
  }

  // frame pacing
  if (now - lastEffectMs < EFFECT_FRAME_MS) return;
  lastEffectMs = now;

  switch (currentEffect) {
    case 0: effectRainbowCycle(); break;
    case 1: effectTheaterChase(); break;
    case 2: effectTwinkle();      break;
    case 3: effectLarson();       break;
    case 4: effectColorWipe();    break;
    default: effectRainbowCycle(); break;
  }
}

// ====================== Setup/Loop ======================
void setup() {
  Serial.begin(115200);
  delay(50);

  // Seed RNG for random effects
  randomSeed(esp_random());

  // Compute cumulative offsets and total pixels
  VIRTUAL_TOTAL = 0;
  for (uint16_t i = 0; i < NUM_LINES; i++) {
    lineOffsets[i] = VIRTUAL_TOTAL;
    VIRTUAL_TOTAL += LINE_LENGTHS[i];
  }
  lineOffsets[NUM_LINES] = VIRTUAL_TOTAL;

  // Allocate line buffers
  for (uint16_t i = 0; i < NUM_LINES; i++) {
    lineLeds[i] = (CRGB*)malloc(sizeof(CRGB) * LINE_LENGTHS[i]);
    memset(lineLeds[i], 0, sizeof(CRGB) * LINE_LENGTHS[i]);
  }

  // Attach each line to FastLED (UNROLLED: compile-time pin requirements)
#if NUM_LINES >= 1
  FastLED.addLeds<LED_TYPE, LINE0_PIN, COLOR_ORDER>(lineLeds[0], LINE_LENGTHS[0]);
#endif
#if NUM_LINES >= 2
  FastLED.addLeds<LED_TYPE, LINE1_PIN, COLOR_ORDER>(lineLeds[1], LINE_LENGTHS[1]);
#endif
#if NUM_LINES >= 3
  FastLED.addLeds<LED_TYPE, LINE2_PIN, COLOR_ORDER>(lineLeds[2], LINE_LENGTHS[2]);
#endif
#if NUM_LINES >= 4
  FastLED.addLeds<LED_TYPE, LINE3_PIN, COLOR_ORDER>(lineLeds[3], LINE_LENGTHS[3]);
#endif
  FastLED.setBrightness(BRIGHTNESS);
  FastLED.show();

  // WiFi
  WiFi.mode(WIFI_STA);
  WiFi.begin(ssid, password);
  Serial.print("WiFi: connecting");
  uint32_t startWait = millis();
  while (WiFi.status() != WL_CONNECTED && (millis() - startWait) < WIFI_CONNECT_TIMEOUT_MS) {
    delay(300); Serial.print(".");
  }
  Serial.println();

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

    // UDP (ESP32 core v3.x signature)
#if USE_MULTICAST
    for (uint16_t u = 0; u < UNIVERSE_COUNT; u++) {
      uint16_t uni = START_UNIVERSE + u;
      IPAddress group = universeToMulticastIP(uni);
      bool ok = udpSockets[u].beginMulticast(group, E131_PORT);
      Serial.print("Join "); Serial.print(group); Serial.print(": ");
      Serial.println(ok ? "OK" : "FAILED");
    }
#else
    udp.begin(E131_PORT);
    Serial.print("E1.31 unicast listening on port "); Serial.println(E131_PORT);
#endif
    lastPacketMs = millis();    // start grace period before fallback
    fallbackActive = false;
  } else {
    Serial.println("WiFi connect failed -> entering fallback mode.");
    fallbackActive = true;
    pickRandomEffect();
    fillAll(CRGB::Black); FastLED.show();
  }

  showStartupChase();
}

void loop() {
  // If WiFi ever drops, go to fallback
  if (WiFi.status() != WL_CONNECTED) {
    if (!fallbackActive) {
      Serial.println("WiFi lost -> fallback effects.");
      fallbackActive = true;
      pickRandomEffect();
    }
  }

  bool gotAnyPacket = false;

#if USE_MULTICAST
  if (!fallbackActive) {
    for (uint16_t s = 0; s < UNIVERSE_COUNT; s++) {
      int packetSize;
      while ((packetSize = udpSockets[s].parsePacket()) > 0) {
        int toRead = min(packetSize, MAX_E131_PACKET);
        int readLen = udpSockets[s].read(packetBuffer, toRead);
        if (readLen <= 0) continue;

        uint16_t universe = 0;
        uint8_t  seq = 0;
        const uint8_t* slots = nullptr;
        uint16_t slotCount = 0;

        if (parseE131(packetBuffer, readLen, universe, seq, slots, slotCount)) {
          mapUniverseSlots(universe, slots, slotCount);
          noteUniverseReceived(seq, universe);
          gotAnyPacket = true;
        }
      }
    }
  }
#else
  if (!fallbackActive) {
    int packetSize;
    while ((packetSize = udp.parsePacket()) > 0) {
      int toRead = min(packetSize, MAX_E131_PACKET);
      int readLen = udp.read(packetBuffer, toRead);
      if (readLen <= 0) continue;

      uint16_t universe = 0;
      uint8_t  seq = 0;
      const uint8_t* slots = nullptr;
      uint16_t slotCount = 0;

      if (parseE131(packetBuffer, readLen, universe, seq, slots, slotCount)) {
        mapUniverseSlots(universe, slots, slotCount);
        noteUniverseReceived(seq, universe);
        gotAnyPacket = true;
      }
    }
  }
#endif

  // Time-out into fallback if no sACN for a while
  if (!fallbackActive) {
    if (millis() - lastPacketMs > SACN_TIMEOUT_MS) {
      Serial.println("sACN timeout -> fallback effects.");
      fallbackActive = true;
      pickRandomEffect();
    }
  }

  // Run effects if in fallback
  if (fallbackActive) {
    runEffects();
  }

  delay(1); // yield to WiFi
}