#pragma once

#include "wled.h"
#include <Adafruit_PWMServoDriver.h>
#include <Wire.h>

class Pca9685Usermod : public Usermod {
  public:
    // Configurable parameters
    bool enabled = true;
    uint8_t i2c_addr = 0x40;
    uint16_t pwm_freq = 50;
    
    // Scaling & Mirroring config
    uint16_t pixelOffset = 0; // Start reading from this pixel index
    bool mirrorRgb = true;    // Mirror WLED pixels to PCA RGB channels
    bool mirrorWhite = true;  // Mirror WLED pixels to PCA White channels

    // Runtime variables
    Adafruit_PWMServoDriver *pwm = nullptr;
    bool initDone = false;

    // Strings for JSON keys to save flash
    static const char _name[];
    static const char _enabled[];
    static const char _i2c_addr[];
    static const char _pwm_freq[];
    static const char _pixelOffset[];
    static const char _mirrorRgb[];
    static const char _mirrorWhite[];

    void setup() {
      if (pwm == nullptr) {
        pwm = new Adafruit_PWMServoDriver(i2c_addr);
      }
      
      // Initialize I2C if not already done by WLED
      if (i2c_scl >= 0 && i2c_sda >= 0) {
        if (!Wire.getClock()) { // Check if Wire is already initialized
             Wire.begin(i2c_sda, i2c_scl);
        }
      } else {
        enabled = false;
        return;
      }
      
      initPca9685();
    }

    void initPca9685() {
      if (!enabled) return;
      
      if (pwm) {
        pwm->begin();
        pwm->setOscillatorFrequency(27000000);
        pwm->setPWMFreq(pwm_freq);  // This is the maximum PWM frequency
        initDone = true;
        DEBUG_PRINTLN(F("PCA9685 init done"));
      }
    }

    void loop() {
      if (!enabled || !initDone || strip.isUpdating()) return;

      // Mirror WLED pixels to PCA9685
      // Uses pixelOffset to determine where to read from in the WLED strip
      
      // RGB Channels (0-11)
      if (mirrorRgb) {
          // RGB 1 (Channels 0,1,2) <- Pixel Offset + 0
          uint32_t c = strip.getPixelColor(pixelOffset + 0);
          setPwmValue(0, (uint16_t)(((c >> 16) & 0xFF) * 16.06)); // R
          setPwmValue(1, (uint16_t)(((c >> 8) & 0xFF) * 16.06));  // G
          setPwmValue(2, (uint16_t)((c & 0xFF) * 16.06));         // B

          // RGB 2 (Channels 3,4,5) <- Pixel Offset + 1
          c = strip.getPixelColor(pixelOffset + 1);
          setPwmValue(3, (uint16_t)(((c >> 16) & 0xFF) * 16.06));
          setPwmValue(4, (uint16_t)(((c >> 8) & 0xFF) * 16.06));
          setPwmValue(5, (uint16_t)((c & 0xFF) * 16.06));

          // RGB 3 (Channels 6,7,8) <- Pixel Offset + 2
          c = strip.getPixelColor(pixelOffset + 2);
          setPwmValue(6, (uint16_t)(((c >> 16) & 0xFF) * 16.06));
          setPwmValue(7, (uint16_t)(((c >> 8) & 0xFF) * 16.06));
          setPwmValue(8, (uint16_t)((c & 0xFF) * 16.06));

          // RGB 4 (Channels 9,10,11) <- Pixel Offset + 3
          c = strip.getPixelColor(pixelOffset + 3);
          setPwmValue(9, (uint16_t)(((c >> 16) & 0xFF) * 16.06));
          setPwmValue(10, (uint16_t)(((c >> 8) & 0xFF) * 16.06));
          setPwmValue(11, (uint16_t)((c & 0xFF) * 16.06));
      }

      // White Channels (12, 13)
      if (mirrorWhite) {
          // White 1 (Channel 12) <- Pixel Offset + 4
          uint32_t c = strip.getPixelColor(pixelOffset + 4);
          uint8_t w = (c >> 24) & 0xFF; // Try to get White component
          if (w == 0) {
            uint8_t r = (c >> 16) & 0xFF;
            uint8_t g = (c >> 8) & 0xFF;
            uint8_t b = c & 0xFF;
            w = max(r, max(g, b)); 
          }
          setPwmValue(12, (uint16_t)(w * 16.06));

          // White 2 (Channel 13) <- Pixel Offset + 5
          c = strip.getPixelColor(pixelOffset + 5);
          w = (c >> 24) & 0xFF; // Try to get White component
          if (w == 0) {
            uint8_t r = (c >> 16) & 0xFF;
            uint8_t g = (c >> 8) & 0xFF;
            uint8_t b = c & 0xFF;
            w = max(r, max(g, b)); 
          }
          setPwmValue(13, (uint16_t)(w * 16.06));
      }
    }

    // Helper to set PWM value safely
    void setPwmValue(int channel, uint16_t value) {
        if (channel < 0 || channel > 15) return;
        if (value > 4095) value = 4096; 
        
        if (value == 4096) {
            pwm->setPWM(channel, 4096, 0); 
        } else if (value == 0) {
            pwm->setPWM(channel, 0, 4096);
        } else {
            pwm->setPWM(channel, 0, value);
        }
    }

    // JSON API to set PWM values
    // usage: {"pca9685":{"0": 2048, "1": 4096}} 
    // values 0-4096. 4096 = fully on, 0 = fully off.
    void readFromJsonState(JsonObject& root) {
      if (!initDone || !enabled) return;

      if (root["pca9685"].is<JsonObject>()) {
        JsonObject pcaJson = root["pca9685"];

        // Handle specific RGB/White control
        // Format: "rgb": [[r,g,b], [r,g,b], [r,g,b], [r,g,b]]
        if (pcaJson["rgb"].is<JsonArray>()) {
            JsonArray rgbArr = pcaJson["rgb"];
            for (int i = 0; i < 4 && i < rgbArr.size(); i++) {
                if (rgbArr[i].is<JsonArray>()) {
                    JsonArray color = rgbArr[i];
                    // Map RGB i to channels
                    // RGB 1: 0,1,2 | RGB 2: 3,4,5 | RGB 3: 6,7,8 | RGB 4: 9,10,11
                    int baseChannel = i * 3;
                    if (color.size() >= 3) {
                        setPwmValue(baseChannel,   (uint16_t)(color[0].as<uint8_t>() * 16.06)); // R 0-255 -> 0-4096
                        setPwmValue(baseChannel+1, (uint16_t)(color[1].as<uint8_t>() * 16.06)); // G
                        setPwmValue(baseChannel+2, (uint16_t)(color[2].as<uint8_t>() * 16.06)); // B
                    }
                }
            }
        }

        // Format: "white": [w1, w2]
        if (pcaJson["white"].is<JsonArray>()) {
            JsonArray whiteArr = pcaJson["white"];
            // White 1: 12 | White 2: 13
            if (whiteArr.size() >= 1) setPwmValue(12, (uint16_t)(whiteArr[0].as<uint8_t>() * 16.06));
            if (whiteArr.size() >= 2) setPwmValue(13, (uint16_t)(whiteArr[1].as<uint8_t>() * 16.06));
        }

        // Handle raw channel control (overrides specific if both present)
        for (JsonPair kv : pcaJson) {
           // Skip "rgb" and "white" keys
           if (strcmp(kv.key().c_str(), "rgb") == 0 || strcmp(kv.key().c_str(), "white") == 0) continue;

           int channel = atoi(kv.key().c_str());
           uint16_t value = kv.value().as<uint16_t>();
           setPwmValue(channel, value);
        }
      }
    }

    void addToConfig(JsonObject& root) {
      JsonObject top = root.createNestedObject(FPSTR(_name));
      top[FPSTR(_enabled)] = enabled;
      top[FPSTR(_i2c_addr)] = i2c_addr;
      top[FPSTR(_pwm_freq)] = pwm_freq;
      top[FPSTR(_pixelOffset)] = pixelOffset;
      top[FPSTR(_mirrorRgb)] = mirrorRgb;
      top[FPSTR(_mirrorWhite)] = mirrorWhite;
    }

    bool readFromConfig(JsonObject& root) {
      JsonObject top = root[FPSTR(_name)];
      if (top.isNull()) {
        return false;
      }
      
      bool configComplete = !top.isNull();
      configComplete &= getJsonValue(top[FPSTR(_enabled)], enabled);
      configComplete &= getJsonValue(top[FPSTR(_i2c_addr)], i2c_addr);
      configComplete &= getJsonValue(top[FPSTR(_pwm_freq)], pwm_freq);
      configComplete &= getJsonValue(top[FPSTR(_pixelOffset)], pixelOffset);
      configComplete &= getJsonValue(top[FPSTR(_mirrorRgb)], mirrorRgb);
      configComplete &= getJsonValue(top[FPSTR(_mirrorWhite)], mirrorWhite);

      if (initDone && !enabled) {
         // If disabled at runtime, maybe we should de-init? 
         // For now just stop updating. 
         // There is no easy "end()" in the library.
      }
      
      if (!initDone && enabled) {
         // Re-init if re-enabled or params changed?
         // Simpler to just assume reboot for major param changes like address
         // But we can update freq
         if (pwm) {
            pwm->setPWMFreq(pwm_freq);
         }
      }

      return configComplete;
    }
    
    uint16_t getId() {
      return USERMOD_ID_UNSPECIFIED; // We don't have a reserved ID, using generic
    }
};

const char Pca9685Usermod::_name[] PROGMEM = "PCA9685";
const char Pca9685Usermod::_enabled[] PROGMEM = "enabled";
const char Pca9685Usermod::_i2c_addr[] PROGMEM = "i2c_addr";
const char Pca9685Usermod::_pwm_freq[] PROGMEM = "pwm_freq";
const char Pca9685Usermod::_pixelOffset[] PROGMEM = "pixelOffset";
const char Pca9685Usermod::_mirrorRgb[] PROGMEM = "mirrorRgb";
const char Pca9685Usermod::_mirrorWhite[] PROGMEM = "mirrorWhite";

static Pca9685Usermod pca9685Usermod;
REGISTER_USERMOD(pca9685Usermod);