#include "ch32v003fun.h"

#define WS2812DMA_IMPLEMENTATION
#include "ws2812b_dma_spi_led_driver.h"

#include "ch32v003_GPIO_branchless.h"


#include "colour_utilities.h"
#include "colours.h"
#include "animations.h"
#include "screens.h"

#define NR_LEDS 25



enum State
{
    off,
    idle,
    pride_scroll,
    hehe,
};


uint32_t output[25];
uint8_t frame = 0;
uint8_t btn_cnt = 0;
enum State state = idle;

// Callbacks that you must implement.
uint32_t WS2812BLEDCallback( int ledno )
{
    if (ledno < 25) {
        if (state == off) {
            return WHITE;
        } else {
            return output[ledno]; 
        }
    } else {
        return OFF;
    }
}

int main()
{
    SystemInit();

    // Enable GPIOD (for debugging)
	RCC->APB2PCENR |= RCC_APB2Periph_GPIOD;
	GPIOD->CFGLR &= ~(0xf<<(4*0));
	GPIOD->CFGLR |= (GPIO_Speed_10MHz | GPIO_CNF_OUT_PP)<<(4*0);

	GPIOD->BSHR = 1;	 // Turn on GPIOD0
	WS2812BDMAInit( );

    GPIO_port_enable(GPIO_port_C);
    GPIO_pinMode(GPIOv_from_PORT_PIN(GPIO_port_C, 4), GPIO_pinMode_I_floating, GPIO_Speed_In);

    // Initialize the output array with some colors
    for (int i = 0; i < NR_LEDS; i++) {
        output[i] = BLACK; // Start with all LEDs off
    }

    while (1)
    {
        // count flashbang to 7 to toggle the off state
        if (GPIO_digitalRead(GPIOv_from_PORT_PIN(GPIO_port_C, 4))) {
            btn_cnt++;
            if (btn_cnt >= 2) {
                // cycle through states
                switch (state) {
                    case off:
                        state = pride_scroll;
                        break;
                    case idle:
                        state = pride_scroll;
                        frame = 0; // Reset the frame for pride scroll
                        break;
                    case pride_scroll:
                        state = hehe;
                        frame = 0; // Reset the frame for text scroll
                        break;  
                    case hehe:
                        state = off;
                        break;
                }

                btn_cnt = 0; // Reset the button count
            }
        } else {
            btn_cnt = 0; // Reset the button count if not pressed
        }

        // Update the LEDs based on the current state
        switch (state) {
            case off:
                // handled in the callback
                break;
            case idle:
                // Set all LEDs to idle screen
                for (int i = 0; i < NR_LEDS; i++) {
                    output[i] = idle_96D[i];
                }
                break;
            case pride_scroll:
                // animate the pride flag scrolling
                if (frame >= 10) {
                    frame = 0; // Reset the frame after one full cycle
                }
                HorizontalScroll(queer_trans_flag_96D, queer_bi_flag_96D, frame, output);
                frame++;
                break;
            case hehe:
                // Set all LEDs to hehe screen
                for (int i = 0; i < NR_LEDS; i++) {
                    output[i] = hehe_96D[i];
                }
        }

		GPIOD->BSHR = 1;	 // Turn on GPIOD0
		// Wait for LEDs to totally finish.
		Delay_Ms( 12 );
		GPIOD->BSHR = 1<<16; // Turn it off

		while( WS2812BLEDInUse );

        // Trigger the DMA transfer
        WS2812BDMAStart( NR_LEDS + 1);
        // Wait for a short period before updating the next LED
        Delay_Ms( 400 ); // Adjust the delay as needed 
    }
}