library ieee;
use ieee.std_logic_1164.all;
use ieee.numeric_std.all;
use ieee.math_real.all;

entity WSDriver is
    generic (
        F_CLK       : natural; -- Board frequency in Hz
        N_LED_MAX   : natural  -- Maximum number of LEDs to instantiate
    );
    port (
        clk         : in std_logic;
        rst_n       : in std_logic;

        -- LED address and LED value
        led_wr_in   : in std_logic;
        led_in      : in std_logic_vector(23 downto 0);
        addr_in     : in std_logic_vector(integer(ceil(log2(real(N_LED_MAX)))) - 1 downto 0);

        -- Write-enable and value of N LEDs to keep active/addressable
        n_wr_in     : in std_logic;
        n_in        : in std_logic_vector(integer(ceil(log2(real(N_LED_MAX)))) - 1 downto 0);

        -- Output 1-bit line for WS2812 strip
        ws_out      : out std_logic;

        -- Output Ready
        ready_out   : out std_logic
    );
end WSDriver;

architecture Behavioral of WSDriver is
    constant F_WS : real := 1.0/((0.3)*10**(-6.0));
    constant T_RES_US : real := 50.0*10**(-6.0);
    constant RES_TICKS : integer := integer(ceil(T_RES_US*F_WS));
    
    constant BITS_PER_LED : integer := 24;

    subtype LED is std_logic_vector(BITS_PER_LED - 1 downto 0);
    type LED_array is array (natural range 0 to N_LED_MAX - 1) of LED;

    type State is (ready, tx, reset);

    signal ws_clk : std_logic;
    signal led_idx_reg, led_idx_next : natural range 0 to N_LED_MAX - 1;
    signal bit_idx_reg, bit_idx_next : natural range 0 to N_LED_MAX - 1;
    signal state_reg, state_next : State;

    constant WS_CLKS_PER_BIT : integer := 4;
    signal ws_cntr_reg, ws_cntr_next : integer range WS_CLKS_PER_BIT - 1 downto 0;
    signal ws_en : std_logic;

    signal leds_reg, leds_next : LED_array;

    -- Register for maintaining the set of driven LEDs
    signal n_leds_reg, n_leds_next : unsigned(integer(ceil(log2(real(N_LED_MAX)))) - 1 downto 0) := (others => '0');

    -- Currently accessed LED
    signal current_led : LED;

    -- Trailing reset counter
    signal reset_cntr_reg, reset_cntr_next : integer range RES_TICKS - 1 downto 0;

    -- Detecting wr_in assertions
    signal wr_in_detec_reg, wr_in_detec_next : std_logic := '0';


begin

    -- Next-state LED and Bit index process
    process (all) begin
        bit_idx_next <= bit_idx_reg;
        led_idx_next <= led_idx_reg;
        reset_cntr_next <= reset_cntr_reg;
        state_next <= state_reg;
        wr_in_detec_next <= wr_in_detec_reg;

        -- Latch wr_in value, to ensure re-transition to tx state if wr_in occured
        -- during a non-ready state.
        if led_wr_in = '1' then
            wr_in_detec_next <= '1';
        end if;

        case state_reg is
            -- STATE READY
            when ready =>
                bit_idx_next <= BITS_PER_LED - 1;
                led_idx_next <= 0;
                
                if n_leds_reg /= 0 and wr_in_detec_reg = '1' then
                    wr_in_detec_next <= '0';
                    state_next <= tx;
                end if;

            -- STATE TX
            when tx =>
                if ws_cntr_reg = (WS_CLKS_PER_BIT - 1) and ws_clk ='1' then
                    if bit_idx_reg = 0 then
                        -- Transition to next LED
                        bit_idx_next <= BITS_PER_LED - 1;
                        led_idx_next <= led_idx_reg + 1;
                    else
                        -- Bit was fully shifted, transition to next bit
                        bit_idx_next <= bit_idx_reg - 1;
                    end if;

                    -- All LEDs fully shifted out (+last bit of last led)? transition to ready
                    if bit_idx_reg = 0 and led_idx_reg + 1 = n_leds_reg then
                        state_next <= reset;
                    end if;
                end if;
                    
            -- STATE TRAILING RESET
            when reset =>
                if ws_clk ='1' then
                    if reset_cntr_reg = RES_TICKS - 1 then
						reset_cntr_next <= 0;
                        state_next <= ready;
                    else
                        reset_cntr_next <= reset_cntr_reg + 1;
                    end if;
                end if;
        end case;
    end process;

    -- Currently accessed LED multiplexer
    current_led <= leds_reg(led_idx_reg) when state_reg = tx else (others => '0');

    -- Next state WS counter
    process (state_reg, ws_cntr_reg, ws_clk) begin
        ws_cntr_next <= ws_cntr_reg;
        if state_reg = ready or state_reg = reset then
            ws_cntr_next <= 0;
        elsif ws_clk = '1' then
            if ws_cntr_reg /= (WS_CLKS_PER_BIT - 1) then
                ws_cntr_next <= ws_cntr_reg + 1;
            else
                ws_cntr_next <= 0;
            end if;
        end if;
    end process;

    -- Next-state LED values
    process(led_wr_in, led_in, leds_reg, addr_in, n_wr_in, n_in) begin
        leds_next <= leds_reg;
        n_leds_next <= n_leds_reg;
		  
        if led_wr_in = '1' then
            leds_next(to_integer(unsigned(addr_in))) <= led_in;
        end if;
		  -- Set N LEDs precedes LED write operation
        if n_wr_in = '1' then
            n_leds_next <= unsigned(n_in);
        end if;
    end process;

    -- WS output clock
    clkgen_ent : entity work.ClkGen
    generic map (
        F_CLK => F_CLK,
        F_OUT => integer(F_WS)
    )
    port map (
        clk => clk,
        rst_n => rst_n,
        clk_o => ws_clk,
        en => ws_en
    );

    -- Only enable ws clock gen when required
    ws_en <= '1' when state_reg /= ready else '0';

    -- Clocking logic
    process(clk) begin
        if rising_edge(clk) then
            if rst_n = '0' then
                ws_cntr_reg <= 0;
                leds_reg <= (others => (others => '0'));
                n_leds_reg <= (others => '0');
                wr_in_detec_reg <= '0';
                bit_idx_reg <= BITS_PER_LED - 1;
                state_reg <= ready;
                led_idx_reg <= 0;
                reset_cntr_reg <= 0;
            else
                ws_cntr_reg <= ws_cntr_next;
                leds_reg <= leds_next;
                n_leds_reg <= n_leds_next;
                wr_in_detec_reg <= wr_in_detec_next;
                bit_idx_reg <= bit_idx_next;
                state_reg <= state_next;
                led_idx_reg <= led_idx_next;
                reset_cntr_reg <= reset_cntr_next;
            end if;
        end if;
    end process;

    -- WS2811 output bit sequence
    process(state_reg, ws_cntr_reg) begin
        ws_out <= '0';
        if state_reg = tx then 
            case ws_cntr_reg is
                when 0 => ws_out <= '1';
                when 1 => ws_out <= current_led(bit_idx_reg);
                when 2 => ws_out <= '0';                        
                when 3 => ws_out <= '0';
            end case;
        end if;
    end process;

    ready_out <= '1' when state_reg = ready else '0';

end Behavioral;