#include "pixel_stomp_mux.h"
#include <Arduino.h>

PixelStompMux::PixelStompMux(uint8_t dat, uint8_t ld, uint8_t clk, uint8_t di)
    : pin_dat(dat), pin_ld(ld), pin_clk(clk), pin_di(di),
      last_button_state(0), current_led_state(0) {
}

void PixelStompMux::begin() {
    pinMode(pin_dat, OUTPUT);
    pinMode(pin_ld, OUTPUT);
    pinMode(pin_clk, OUTPUT);
    pinMode(pin_di, INPUT_PULLUP);

    digitalWrite(pin_dat, LOW);
    digitalWrite(pin_ld, LOW);
    digitalWrite(pin_clk, LOW);

    Serial.printf("[MUX] Init DAT=%d LD=%d CLK=%d DI=%d\n",
                  pin_dat, pin_ld, pin_clk, pin_di);

    Serial.println("[MUX] GPIO probe:");
    Serial.printf("  DI (GPIO %d) reads: %d (should be 1 with pullup)\n",
                  pin_di, digitalRead(pin_di));

    for (int i = 0; i < 3; i++) {
        digitalWrite(pin_clk, HIGH);
        delayMicroseconds(10);
        digitalWrite(pin_clk, LOW);
        delayMicroseconds(10);
    }
    Serial.printf("  DI after 3 CLK pulses: %d\n", digitalRead(pin_di));

    digitalWrite(pin_ld, HIGH);
    delayMicroseconds(10);
    digitalWrite(pin_ld, LOW);
    delayMicroseconds(10);
    Serial.printf("  DI after LD pulse: %d\n", digitalRead(pin_di));
}

uint16_t PixelStompMux::read_buttons() {
    pulse_pin(pin_ld, HIGH);
    delayMicroseconds(10);
    pulse_pin(pin_ld, LOW);
    delayMicroseconds(10);

    uint16_t state = shift_in(NUM_INPUTS);
    last_button_state = state;

    return state;
}

bool PixelStompMux::is_button_pressed(uint8_t index) {
    if (index >= NUM_INPUTS) return false;

    uint16_t state = read_buttons();
    return !(state & (1 << index));
}

void PixelStompMux::write_leds(uint16_t led_state) {
    current_led_state = led_state;
    shift_out(led_state, NUM_OUTPUTS);

    pulse_pin(pin_ld, HIGH);
    delayMicroseconds(10);
    pulse_pin(pin_ld, LOW);
}

void PixelStompMux::set_led(uint8_t index, bool on) {
    if (index >= NUM_OUTPUTS) return;

    if (on) {
        current_led_state |= (1 << index);
    } else {
        current_led_state &= ~(1 << index);
    }

    write_leds(current_led_state);
}

void PixelStompMux::clear_all() {
    write_leds(0x0000);
    Serial.println("[MUX] All LEDs off");
}

void PixelStompMux::dump() {
    uint16_t buttons = read_buttons();
    Serial.printf("[MUX] Buttons: 0x%03X LEDs: 0x%03X\n", buttons, current_led_state);

    for (int i = 0; i < NUM_INPUTS; i++) {
        bool pressed = !(buttons & (1 << i));
        bool led_on = current_led_state & (1 << i);
        Serial.printf("  [%d] Button:%s LED:%s\n", i,
                      pressed ? "PRESS" : "off",
                      led_on ? "ON" : "off");
    }
}

void PixelStompMux::shift_out(uint16_t data, uint8_t bits) {
    for (int i = bits - 1; i >= 0; i--) {
        digitalWrite(pin_dat, (data >> i) & 1);
        delayMicroseconds(1);
        digitalWrite(pin_clk, HIGH);
        delayMicroseconds(1);
        digitalWrite(pin_clk, LOW);
        delayMicroseconds(1);
    }
}

uint16_t PixelStompMux::shift_in(uint8_t bits) {
    uint16_t data = 0;

    for (int i = bits - 1; i >= 0; i--) {
        bool bit = digitalRead(pin_di);
        if (bit) {
            data |= (1 << i);
        }
        digitalWrite(pin_clk, HIGH);
        delayMicroseconds(1);
        digitalWrite(pin_clk, LOW);
        delayMicroseconds(1);
    }

    return data;
}

void PixelStompMux::pulse_pin(uint8_t pin, uint8_t level, uint16_t duration_us) {
    digitalWrite(pin, level);
    delayMicroseconds(duration_us);
}