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

// --- HARDWARE CONFIGURATION ---
#define PIN_NEOPIXEL 6
#define NEOPIXELS_X 16
#define NEOPIXELS_Y 8
#define NUM_PIXELS   64*2 // 8x8 LEDs in 2 panels

LedMatrix::LedMatrix() 
    : pixels(NUM_PIXELS, PIN_NEOPIXEL, NEO_GRB + NEO_KHZ800)
{
}

void LedMatrix::begin() {
    pixels.setPin(PIN_NEOPIXEL);
    pixels.begin();
    pixels.setBrightness(40);
    pixels.clear();
    pixels.show();
}

void LedMatrix::setBrightness(uint8_t b) {
    pixels.setBrightness(b);
}

void LedMatrix::clear() {
    pixels.clear();
}

void LedMatrix::show() {
    pixels.show();
}

uint32_t LedMatrix::getNoteColor(int note, bool dim) {
    if (note == -1) return 0;
    uint16_t hue = 30000 + (note % 12) * 3628;
    return Adafruit_NeoPixel::ColorHSV(hue, 255, dim ? 10 : 50);
}

int LedMatrix::getPixelIndex(int x, int y) {
    return NEOPIXELS_Y * x + (NEOPIXELS_Y - 1 - y);
}

void LedMatrix::update(const Step sequence[][NUM_STEPS], int playbackStep, bool isPlaying, 
                UIState currentState, bool songModeEnabled, 
                int songRepeatsRemaining, bool sequenceChangeScheduled, PlayMode playMode,
                int selectedTrack, const int* numSteps, const bool* trackMute) {
    pixels.clear();
    
    const uint32_t COLOR_PLAYHEAD = pixels.Color(0, 255, 0);
    const uint32_t COLOR_PLAYHEAD_DIM = pixels.Color(0, 32, 0);
    const uint32_t COLOR_MUTED_PLAYHEAD = pixels.Color(0, 0, 255);
    const uint32_t COLOR_CURSOR = pixels.Color(255, 100, 100);
    const uint32_t COLOR_CURSOR_DIM = pixels.Color(32, 0, 0);

    if(playMode == MODE_POLY) {
        for(int t=0; t<NUM_TRACKS; t++) {
            int currentTrackSteps = numSteps[t];
            for(int s=0; s<NUM_STEPS; s++) {
                int col = s; // Each step is a column

                // --- First row of track pair: Notes ---
                int note_row = t * 2;
                uint32_t note_color = 0;
                if (s < currentTrackSteps) {
                    int note = sequence[t][s].note;
                    if (note != -1) {
                        note_color = getNoteColor(note, !sequence[t][s].accent);
                    }
                }
                pixels.setPixelColor(getPixelIndex(col, note_row), note_color);

                // --- Second row of track pair: Steps & Playhead ---
                int step_row = t * 2 + 1;
                uint32_t step_color = 0; // Off by default for steps > currentTrackSteps
                if (s < currentTrackSteps) {
                    step_color = COLOR_PLAYHEAD_DIM; // It's a valid step
                    if (isPlaying && (s == (playbackStep % currentTrackSteps))) {
                        if (trackMute[t]) {
                            step_color = COLOR_MUTED_PLAYHEAD;
                        } else {
                            step_color = COLOR_PLAYHEAD;
                        }
                    }
                }
                pixels.setPixelColor(getPixelIndex(col, step_row), step_color);
            }
        }
    } else {
        // --- Mono Mode (original) ---
        const Step* trackSequence = sequence[selectedTrack];
        for (int s = 0; s < NUM_STEPS; s++) {
            if (s >= numSteps[selectedTrack]) continue;

            int x = s % NUM_STEPS;
            int yBase = (s / NUM_STEPS) * 2;
            uint32_t color = 0, dimColor = 0;
            bool isCursorHere = (isPlaying && s == (playbackStep % numSteps[selectedTrack]));
            if (trackSequence[s].note != -1) {
                color = getNoteColor(trackSequence[s].note, trackSequence[s].tie);
                dimColor = getNoteColor(trackSequence[s].note, true);
            }
            uint32_t c[4] = {0};
            if (trackSequence[s].note != -1) {
                int octave = trackSequence[s].note / 12;
                if (octave > 4) { c[0] = color; if (trackSequence[s].accent) c[1] = dimColor; }
                else if (octave < 4) { c[2] = color; if (trackSequence[s].accent) c[1] = dimColor; }
                else { c[1] = color; if (trackSequence[s].accent) { c[0] = dimColor; c[2] = dimColor; } }
            }
            uint32_t cursorColor = pixels.Color(0, 0, 50);
            if (isPlaying) {
                 cursorColor = pixels.Color(0, 50, 0);
                 if (songModeEnabled && s >= NUM_STEPS) {
                     int repeats = min(songRepeatsRemaining, NUM_STEPS);
                     if (x >= (NUM_STEPS - repeats)) cursorColor = (songRepeatsRemaining == 1 && x == 7 && (millis()/250)%2) ? pixels.Color(255, 200, 0) : pixels.Color(100, 220, 40);
                 }
            }
            
            if (cursorColor != 0) {
                if (isCursorHere) {
                    for(int i=0; i<4; i++) {
                        if (c[i] == 0) c[i] = cursorColor;
                    }
                } else {
                    uint8_t r = (uint8_t)(cursorColor >> 16), g = (uint8_t)(cursorColor >> 8), b = (uint8_t)cursorColor;
                    c[3] = pixels.Color(r/5, g/5, b/5);
                }
            }
            for(int i=0; i<4; i++) pixels.setPixelColor(getPixelIndex(x, yBase + i), c[i]);
        }

        // --- Overview of all tracks on bottom 4 rows ---
        for(int t=0; t<NUM_TRACKS; t++) {
            int row = 4 + t;
            int currentTrackSteps = numSteps[t];
            for(int s=0; s<NUM_STEPS; s++) {
                if (s >= currentTrackSteps) continue;
                
                uint32_t pixelColor = 0;
                
                // Background for active track
                if (t == selectedTrack) {
                    pixelColor = COLOR_CURSOR_DIM;
                }
                else if (sequence[t][s].note != -1) // Note
                {
                    pixelColor = getNoteColor(sequence[t][s].note, !sequence[t][s].accent);
                }
                
                // Playhead
                if (isPlaying && (s == (playbackStep % currentTrackSteps))) {
                    if (t == selectedTrack) {
                        pixelColor = COLOR_CURSOR;
                    } else if (trackMute[t]) {
                        pixelColor = COLOR_MUTED_PLAYHEAD;
                    } else {
                        pixelColor = COLOR_PLAYHEAD;
                    }
                }
                
                pixels.setPixelColor(getPixelIndex(s, row), pixelColor);
            }
        }
    }
    if (sequenceChangeScheduled && (millis() / 125) % 2) pixels.setPixelColor(NUM_PIXELS - 1, pixels.Color(127, 50, 0));
    pixels.show();
}