`include "defines.v"
`define TX_DATA_VALID_CNT 13
module tx_fsm(
    input [31:0] aa_i,
    input [1:0] tx_rate_i,
    input whiten_bypass_i,
    input [7:0] tx_up_us_i,
    input [4:0] tx_down_us_i,
    input [7:0] tx_byte_i,
    input [7:0] tx_payload_length_i,
    input aes_ccm_en_i,
    output tx_fifo_pop_o,
    input start_i,
    output reg tx_en_o,
    output reg tx_data_valid_o,
    output tx_data_o,

    output aes_ccm_start_o,
    output [7:0] unencrypted_byte_o,
    output [7:0] unecnrypted_byte_cnt_o,
    output unencrypted_byte_update_o,
    output reg [7:0] ccm_aad_o,
    input encrypted_byte_valid_i,
    input [31:0] tx_mic_i,
    input [7:0] encrypted_byte_i,

    output crc_data_o,
    output tx_crc_en_o,
    input [23:0] tx_crc_i,

    output reg unwhitened_data_o,
    output tx_whiten_shift_o,
    input whitened_data_i,

    input clk,
    input rst_n
);

reg [2:0] c_st;
reg [2:0] n_st;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) c_st <= `TX_IDLE;
    else c_st <= n_st;
end

reg [4:0] cnt;
wire sym_end;
assign sym_end = (tx_rate_i == 0 && cnt == `CLK_MHZ - 1) || (tx_rate_i == 1 && cnt == `CLK_MHZ/2 - 1);
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) begin
        cnt <= ~0;
        tx_en_o <= 0;
    end
    else if(n_st != `TX_IDLE) begin
        tx_en_o <= 1;
        if(sym_end) cnt <= 0;
        else cnt <= cnt + 1;
    end
    else begin
        tx_en_o <= 0;
        cnt <= 0;
    end
end

always @(posedge clk,negedge rst_n) begin
    if(!rst_n) tx_data_valid_o <= 0;
    else if(c_st != `TX_IDLE && c_st != `TX_UP && c_st != `TX_DOWN && cnt == `TX_DATA_VALID_CNT) tx_data_valid_o <= 1;
    else tx_data_valid_o <= 0;
end

reg [8:0] up_cnt;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) up_cnt <= 0;
    else if(c_st == `TX_UP) begin 
        if(sym_end) up_cnt <= up_cnt + 1; 
    end
    else up_cnt <= 0;
end

reg [3:0] preamble_cnt;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) preamble_cnt <= 0;
    else if(c_st == `TX_PREAMBLE) begin 
        if(sym_end) preamble_cnt <= preamble_cnt + 1; 
    end
    else preamble_cnt <= 0;
end

reg [4:0] aa_cnt;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) aa_cnt <= 0;
    else if(c_st == `TX_SYNC_WORD) begin 
        if(sym_end) aa_cnt <= aa_cnt + 1;
    end
    else aa_cnt <= 0;
end

reg [11:0] pdu_cnt;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) pdu_cnt <= 0;
    else if(c_st == `TX_PDU) begin
        if(sym_end) pdu_cnt <= pdu_cnt + 1;
    end
    else pdu_cnt <= 0;
end

always @(posedge clk,negedge rst_n) begin
    if(!rst_n) ccm_aad_o <= 0;
    else if(c_st == `TX_PDU && pdu_cnt == 0) ccm_aad_o <= tx_byte_i;
end

reg [4:0] crc_cnt;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) crc_cnt <= 0;
    else if(c_st == `TX_CRC) begin
        if(sym_end) crc_cnt <= crc_cnt + 1;
    end
    else crc_cnt <= 0;
end

reg [4:0] down_cnt;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) down_cnt <= 0;
    else if(c_st == `TX_DOWN) begin
        if(sym_end) down_cnt <= down_cnt + 1;
    end
    else down_cnt <= 0;
end

wire [5:0] mic_cnt;
wire on_mic_tx;
wire [8:0] tx_pdu_length;
assign tx_pdu_length = tx_payload_length_i + 2;
assign mic_cnt = 8*tx_pdu_length - pdu_cnt - 1;
assign on_mic_tx = 8*tx_pdu_length - pdu_cnt - 1 <= 31 && aes_ccm_en_i && tx_payload_length_i;
assign tx_fifo_pop_o = pdu_cnt%8 == 7 && sym_end && c_st == `TX_PDU && !on_mic_tx;

reg payload_byte_pop;
always @(posedge clk,negedge rst_n) begin
    if(!rst_n) payload_byte_pop <= 0;
    else if(tx_fifo_pop_o && pdu_cnt > 8) payload_byte_pop <= 1;
    else payload_byte_pop <= 0;
end

assign aes_ccm_start_o = aes_ccm_en_i && pdu_cnt == 8 && cnt == 0 && tx_payload_length_i;
assign unencrypted_byte_o = tx_byte_i;
assign unecnrypted_byte_cnt_o = pdu_cnt/8 - 2;
assign unencrypted_byte_update_o = payload_byte_pop && !on_mic_tx;

assign tx_crc_en_o = cnt == `TX_DATA_VALID_CNT && c_st == `TX_PDU;
assign crc_data_o = unwhitened_data_o;
assign tx_whiten_shift_o = sym_end&&(c_st==`TX_PDU||c_st==`TX_CRC);
assign tx_data_o = whiten_bypass_i ?  unwhitened_data_o : (c_st==`TX_PDU||c_st==`TX_CRC) ? whitened_data_i : unwhitened_data_o;

always @(*) begin
    case (c_st)
        `TX_PREAMBLE:    unwhitened_data_o = preamble_cnt[0] ? ~aa_i[0] : aa_i[0];
        `TX_SYNC_WORD:   unwhitened_data_o = aa_i[aa_cnt];
        `TX_PDU:         unwhitened_data_o = pdu_cnt >= 16 ? aes_ccm_en_i ? on_mic_tx ? tx_mic_i[31 - mic_cnt] : encrypted_byte_i[pdu_cnt%8] : tx_byte_i[pdu_cnt%8] : tx_byte_i[pdu_cnt%8];
        `TX_CRC:         unwhitened_data_o = tx_crc_i[23 - crc_cnt];
        default:         unwhitened_data_o = 0;
    endcase
end

always @(*) begin
    n_st = c_st;
    case (c_st)
        `TX_IDLE:        if(start_i) n_st = `TX_UP;
        `TX_UP:          if((tx_rate_i ? (up_cnt == 2*tx_up_us_i + 1): (up_cnt == tx_up_us_i))  && sym_end) n_st = `TX_PREAMBLE;
        `TX_PREAMBLE:    if((tx_rate_i == 0 && preamble_cnt == 7 || tx_rate_i == 1 && preamble_cnt == 15) && sym_end) n_st = `TX_SYNC_WORD;
        `TX_SYNC_WORD:   if(aa_cnt == 31 && sym_end) n_st = `TX_PDU;
        `TX_PDU:         if(pdu_cnt + 1 == 8*tx_pdu_length && sym_end) n_st = `TX_CRC;
        `TX_CRC:         if(crc_cnt == 23 && sym_end) n_st = `TX_DOWN;
        `TX_DOWN:        if(down_cnt == tx_down_us_i && sym_end) n_st = `TX_IDLE;
    endcase
end



endmodule