library ieee;
use ieee.std_logic_1164.all;
use ieee.numeric_std.all;
use work.router_types.all;
use work.encoder_components.all;

entity arbiter_unit is
generic(
    num_paths_in     :   positive := 6;
    num_paths_out    :   positive := 1;
    lsb_sels_size    :   positive := 3
);
port (
    dirs         :   in std_logic_vector(num_paths_in-1 downto 0);
    packets      :   in t_DATA(num_paths_in-1 downto 0);
    avai_paths   :   in std_logic_vector(num_paths_out-1 downto 0);
    out_paths    :   out t_DATA(num_paths_out-1 downto 0);
    valids_out   :   out std_logic_vector(num_paths_out-1 downto 0);
    updated_dirs :   out std_logic_vector(num_paths_in-1 downto 0)
);
end arbiter_unit;

architecture rtl of arbiter_unit is
    constant NPI                : integer := num_paths_in;
    constant NPO                : positive := num_paths_out;
    constant NUM_ENCODERS       : integer := get_num_encoders(NPI);
    constant ENC_SIZE           : positive := get_encoder_size(NPI);

    type t_N_LSB_SELS is array(natural range<>) of std_logic_vector(lsb_sels_size-1 downto 0); -- size depends on num_paths_out
    type t_GRANT_DIRS is array(natural range<>) of std_logic_vector(NPI-1 downto 0);
    type t_SELECTORS is array(natural range<>) of std_logic_vector(get_dout_size(ENC_SIZE)-1 downto 0);

    component NthLsbDetector is
    generic(
        num_dirs     :   positive := 1;
        n            :   positive := 1
    );
    port (
        dirs         :   in std_logic_vector(num_dirs-1 downto 0);
        grant_dirs   :   out std_logic_vector(num_dirs-1 downto 0)
    );
    end component NthLsbDetector;

    signal selectors         : t_SELECTORS(NUM_ENCODERS*NPO-1 downto 0);
    signal valids            : std_logic_vector(NUM_ENCODERS*NPO-1 downto 0);
    signal packet_index      : integer range 0 to NPI;
    signal n_lsb_sels        : t_N_LSB_SELS(NPO-1 downto 0);
    signal nlsb_grant_dirs   : t_GRANT_DIRS(NPO-1 downto 0);
    signal grant_dirs        : t_GRANT_DIRS(NPO-1 downto 0);
begin
    set_n_lsb_sel: process(avai_paths)
        type t_SUM_AVAI_PATHS is array(natural range<>) of integer range 0 to NPO;
        
        variable sum_avai_paths : t_SUM_AVAI_PATHS(NPO-1 downto 0);
    begin
        sum_avai_paths(0) := 0;
        n_lsb_sels(0) <= (others => '0');
        for i in 1 to NPO-1 loop
            if avai_paths(0) = '1' then
                sum_avai_paths(i) := sum_avai_paths(i-1)+1;
            else
                sum_avai_paths(i) := sum_avai_paths(i-1);
            end if;
            n_lsb_sels(i) <= std_logic_vector(to_unsigned(sum_avai_paths(i), LSB_SELS_SIZE));
        end loop;
    end process;

    g_NLSB_UNIT: for i in 0 to NPO-1 generate
    nlsb_detector: NthLsbDetector
        generic map(num_dirs => NPI, n => i+1)
        port map(dirs => dirs, grant_dirs => nlsb_grant_dirs(i));
    end generate;

    g_grant_dirs: for i in 0 to NPO-1 generate
        grant_dirs(i) <= nlsb_grant_dirs(to_integer(unsigned(n_lsb_sels(i))));
    end generate;

    g_encoders: for i in 0 to NPO-1 generate
    begin
        g_encoder: for e in 0 to NUM_ENCODERS-1 generate
            signal enc_din  :   std_logic_vector(ENC_SIZE-1 downto 0);
        begin
            set_enc_din: process(grant_dirs)
            begin
                if e < NUM_ENCODERS-1 then
                    enc_din <= grant_dirs(i)(ENC_SIZE*(e+1)-1 downto ENC_SIZE*e);
                else
                    enc_din(ENC_SIZE-1 downto NPI-ENC_SIZE*(NUM_ENCODERS-1)) <= (others => '0');
                    enc_din(NPI-(ENC_SIZE*(NUM_ENCODERS-1))-1 downto 0) <= grant_dirs(i)(NPI-1 downto ENC_SIZE*(NUM_ENCODERS-1));
                end if;
            end process;

            encoder: encoder_generic
                generic map (din_size => ENC_SIZE)
                port map (din => enc_din, dout => selectors(i+NPO*e), 
                        valid => valids(i+NPO*e));
        end generate;

        set_outputs: process(selectors, valids, packets, avai_paths)
            variable packet_index   :   integer range 0 to NPI;
            variable valid_sel      :   std_logic;
        begin
            packet_index := 0;
            valid_sel := '0';
            for e in 0 to NUM_ENCODERS-1 loop
                if valids(i+NPO*e) = '1' then
                    packet_index := to_integer(unsigned(selectors(i+NPO*e)))+32*e;
                    valid_sel := '1';
                end if;
            end loop;

            out_paths(i)  <= packets(packet_index);
            valids_out(i) <= avai_paths(i) and valid_sel;
        end process;
    end generate;

    set_updated_dirs: process(dirs, selectors, valids, avai_paths)
        variable updated_dirs_index :   integer range 0 to NPI;
    begin
        updated_dirs <= dirs;
        for i in 0 to num_paths_out-1 loop
            for j in 0 to NUM_ENCODERS-1 loop
                if avai_paths(i) = '1' and valids(i+num_paths_out*j) = '1' then
                    updated_dirs_index := to_integer(unsigned(selectors(i+num_paths_out*j)));
                    updated_dirs(updated_dirs_index+32*j) <= '0';
                end if;
            end loop;
        end loop;
    end process;
end rtl;