// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: CERN-OHL-S-2.0

#include "hls_jfjoch.h"
#include "ip_header_checksum.h"

void generate_packet_jungfrau(STREAM_512 &data_out,
                              ap_uint<256> *d_hbm_p0,
                              ap_uint<256> *d_hbm_p1,
                              ap_uint<32> hbm_size_bytes,
                              ap_uint<32> frame,
                              ap_uint<10> mem_cell,
                              ap_uint<8> module,
                              ap_uint<8> data_stream,
                              ap_uint<7> eth_packet,
                              ap_uint<48> src_mac_addr,
                              ap_uint<48> dest_mac_addr,
                              ap_uint<32> src_ipv4_addr,
                              ap_uint<32> dest_ipv4_addr,
                              ap_uint<64> bunchid,
                              ap_uint<32> exptime,
                              ap_uint<64> timestamp,
                              ap_uint<32> debug,
                              ap_uint<1> user = 0) {
#pragma HLS PIPELINE II=130
    ap_uint<720> header = 0;

    header(47, 0) = dest_mac_addr;
    header(95, 48) = src_mac_addr;
    header(111, 96) = 0x0008; // ETHER_IP = IPv4
    header(eth_payload_pos + 3, eth_payload_pos) = 0x5; // header len of 5
    header(eth_payload_pos + 7, eth_payload_pos + 4) = 0x4; // IPv4
    header(eth_payload_pos + 31, eth_payload_pos + 16) = 0x4C20; // total length = 8268
    header(eth_payload_pos + 79, eth_payload_pos + 72) = PROTOCOL_UDP; // UDP
    header(eth_payload_pos + 127, eth_payload_pos + 96) = src_ipv4_addr;
    header(eth_payload_pos + 159, eth_payload_pos + 128) = dest_ipv4_addr;
    header(eth_payload_pos + 95, eth_payload_pos + 80) = computeCheckSum20B(
        header(eth_payload_pos + 159, eth_payload_pos));

    header(ipv4_payload_pos + 31, ipv4_payload_pos + 24) = 2 * module;
    header(ipv4_payload_pos + 23, ipv4_payload_pos + 16) = data_stream;
    header(ipv4_payload_pos + 47, ipv4_payload_pos + 32) = 0x3820; // UDP length = 8248
    header(udp_payload_pos + 63, udp_payload_pos) = frame + 1;
    header(udp_payload_pos + 95, udp_payload_pos + 64) = exptime;
    header(udp_payload_pos + 127, udp_payload_pos + 96) = eth_packet;
    header(udp_payload_pos + 2 * 64 + 63, udp_payload_pos + 2 * 64) = bunchid;
    header(udp_payload_pos + 3 * 64 + 63, udp_payload_pos + 3 * 64) = timestamp;
    header(udp_payload_pos + 4 * 64 + 31, udp_payload_pos + 4 * 64 + 16) = 0;
    header(udp_payload_pos + 5 * 64 + 31, udp_payload_pos + 5 * 64) = debug;
    header(udp_payload_pos + 5 * 64 + 55, udp_payload_pos + 5 * 64 + 48) = SLS_DETECTOR_TYPE_JUNGFRAU;

    packet_512_t packet;
    packet.data = header(511, 0);
    packet.last = 0;
    packet.dest = 0;
    packet.id = 0;
    packet.strb = UINT64_MAX;
    packet.keep = UINT64_MAX;
    packet.user = user;

    data_out << packet;

    ap_uint<32> offset_hbm_0 = 20 * hbm_size_bytes / 32;
    ap_uint<32> offset_hbm_1 = 21 * hbm_size_bytes / 32;

    ap_uint<208> remainder = header(719, 512);
    for (int i = 0; i < 128; i++) {
        ap_uint<512> tmp;
        tmp(255, 0) = d_hbm_p0[offset_hbm_0 + mem_cell * 128 * 128 + eth_packet * 128 + i];
        tmp(511, 256) = d_hbm_p1[offset_hbm_1 + mem_cell * 128 * 128 + eth_packet * 128 + i];

        packet.data(207, 0) = remainder;
        packet.data(511, 208) = tmp(303, 0);
        data_out << packet;
        remainder = tmp(511, 304);
    }
    packet.data(207, 0) = remainder;
    packet.data(511, 208) = 0;
    packet.keep(63, 26) = 0;
    packet.last = 1;
    data_out << packet;
}


void generate_packet_eiger(STREAM_512 &data_out,
                           ap_uint<256> *d_hbm_p0,
                           ap_uint<256> *d_hbm_p1,
                           ap_uint<32> hbm_size_bytes,
                           ap_uint<32> frame,
                           ap_uint<10> mem_cell,
                           ap_uint<8> module,
                              ap_uint<8> data_stream,
                           ap_uint<9> eth_packet,
                           ap_uint<48> src_mac_addr,
                           ap_uint<48> dest_mac_addr,
                           ap_uint<32> src_ipv4_addr,
                           ap_uint<32> dest_ipv4_addr,
                           ap_uint<32> exptime,
                           ap_uint<64> timestamp,
                           ap_uint<32> debug,
                           ap_uint<8> bit_depth,
                           ap_uint<1> user = 0) {
#pragma HLS PIPELINE II=66
    ap_uint<720> header = 0;

    header(47, 0) = dest_mac_addr;
    header(95, 48) = src_mac_addr;
    header(111, 96) = 0x0008; // ETHER_IP = IPv4
    header(eth_payload_pos + 3, eth_payload_pos) = 0x5; // header len of 5
    header(eth_payload_pos + 7, eth_payload_pos + 4) = 0x4; // IPv4
    header(eth_payload_pos + 31, eth_payload_pos + 16) = 0x4C10; // total length = 4172
    header(eth_payload_pos + 79, eth_payload_pos + 72) = PROTOCOL_UDP; // UDP
    header(eth_payload_pos + 127, eth_payload_pos + 96) = src_ipv4_addr;
    header(eth_payload_pos + 159, eth_payload_pos + 128) = dest_ipv4_addr;
    header(eth_payload_pos + 95, eth_payload_pos + 80) = computeCheckSum20B(
        header(eth_payload_pos + 159, eth_payload_pos));

    header(ipv4_payload_pos + 47, ipv4_payload_pos + 32) = 0x3810; // UDP length = 4152
    header(udp_payload_pos + 63, udp_payload_pos) = frame + 1;
    header(udp_payload_pos + 95, udp_payload_pos + 64) = exptime;
    header(udp_payload_pos + 2 * 64 + 63, udp_payload_pos + 2 * 64) = 0;
    header(udp_payload_pos + 3 * 64 + 63, udp_payload_pos + 3 * 64) = timestamp;

    ap_uint<1> half_module;
    if (bit_depth == 16) {
        header(udp_payload_pos + 127, udp_payload_pos + 96) = (eth_packet / 2) % 64;
        half_module = ((eth_packet >= 128) ? 1 : 0);
    } else if (bit_depth == 32) {
        header(udp_payload_pos + 127, udp_payload_pos + 96) = (eth_packet / 2) % 128;
        half_module = ((eth_packet >= 256) ? 1 : 0);
    } else if (bit_depth == 8) {
        header(udp_payload_pos + 127, udp_payload_pos + 96) = (eth_packet / 2) % 32;
        half_module = ((eth_packet >= 64) ? 1 : 0);
    }
    // Big endian
    header(ipv4_payload_pos + 31, ipv4_payload_pos + 24) = 2 * module + half_module;
    header(ipv4_payload_pos + 23, ipv4_payload_pos + 16) = data_stream;

    header(udp_payload_pos + 4 * 64 + 47, udp_payload_pos + 4 * 64 + 32) = (eth_packet % 2);
    header(udp_payload_pos + 5 * 64 + 31, udp_payload_pos + 5 * 64) = debug;
    header(udp_payload_pos + 5 * 64 + 55, udp_payload_pos + 5 * 64 + 48) = SLS_DETECTOR_TYPE_EIGER;

    packet_512_t packet;
    packet.data = header(511, 0);
    packet.last = 0;
    packet.dest = 0;
    packet.id = 0;
    packet.strb = UINT64_MAX;
    packet.keep = UINT64_MAX;
    packet.user = user;

    data_out << packet;

    ap_uint<32> offset_hbm_0 = 20 * hbm_size_bytes / 32;
    ap_uint<32> offset_hbm_1 = 21 * hbm_size_bytes / 32;

    ap_uint<208> remainder = header(719, 512);
    for (int i = 0; i < 64; i++) {
        ap_uint<512> tmp;
        tmp(255, 0) = d_hbm_p0[offset_hbm_0 + mem_cell * 64 * 256 + eth_packet * 64 + i];
        tmp(511, 256) = d_hbm_p1[offset_hbm_1 + mem_cell * 64 * 256 + eth_packet * 64 + i];

        packet.data(207, 0) = remainder;
        packet.data(511, 208) = tmp(303, 0);
        data_out << packet;
        remainder = tmp(511, 304);
    }
    packet.data(207, 0) = remainder;
    packet.data(511, 208) = 0;
    packet.keep(63, 26) = 0;
    packet.last = 1;
    data_out << packet;
}

int frame_generator(STREAM_512 &data_out,
                    ap_uint<256> *d_hbm_p0,
                    ap_uint<256> *d_hbm_p1,
                    ap_uint<32> hbm_size_bytes,
                    ap_uint<48> src_mac_addr,
                    ap_uint<32> src_ipv4_addr,
                    volatile ap_uint<1> &in_cancel,
                    const FrameGeneratorConfig &config) {
#pragma HLS INTERFACE mode=s_axilite port=return
#pragma HLS INTERFACE mode=s_axilite port=config
#pragma HLS INTERFACE mode=ap_none register port=src_mac_addr
#pragma HLS INTERFACE mode=ap_none register port=src_ipv4_addr
#pragma HLS INTERFACE mode=ap_none register port=in_cancel

#pragma HLS INTERFACE register ap_none port=hbm_size_bytes
#pragma HLS INTERFACE m_axi port=d_hbm_p0 bundle=d_hbm_p0 depth=512 offset=off \
        max_read_burst_length=16 max_write_burst_length=2 latency=120 num_write_outstanding=2 num_read_outstanding=9
#pragma HLS INTERFACE m_axi port=d_hbm_p1 bundle=d_hbm_p1 depth=512 offset=off \
        max_read_burst_length=16 max_write_burst_length=2 latency=120 num_write_outstanding=2 num_read_outstanding=9

#pragma HLS INTERFACE register both axis port=data_out

    if ((config.modules > MAX_MODULES_FPGA) || (config.modules == 0))
        return 1;

    if ((config.detector_type != SLS_DETECTOR_TYPE_EIGER)
        && (config.detector_type != SLS_DETECTOR_TYPE_JUNGFRAU))
        return 2;

    if (config.detector_type == SLS_DETECTOR_TYPE_EIGER) {
        if ((config.eiger_bit_depth != 8) &&
            (config.eiger_bit_depth != 16) &&
            (config.eiger_bit_depth != 32))
            return 2;
    }

    if (config.frames == 0)
        return 0;

    if (config.modules * (config.images_in_memory + 1) >= hbm_size_bytes * 2 / (RAW_MODULE_SIZE * sizeof(uint16_t)))
        return 3;
    if ((config.detector_type == SLS_DETECTOR_TYPE_EIGER) && (config.eiger_bit_depth == 32)) {
        if (config.modules * (config.images_in_memory + 1) * 2 >= hbm_size_bytes * 2 / (
                RAW_MODULE_SIZE * sizeof(uint16_t)))
            return 3;
    }

    union {
        uint64_t u64;
        double d;
    } pulse_id_conv;
    pulse_id_conv.d = config.pulse_id;

    // send one packet with TUSER = 1, to reset on network cores on the way
    generate_packet_jungfrau(data_out,
                             d_hbm_p0,
                             d_hbm_p1,
                             hbm_size_bytes,
                             0, 0, 0,
                             config.data_stream,
                             0,
                             src_mac_addr,
                             config.dest_mac_addr,
                             src_ipv4_addr,
                             config.dest_ipv4_addr,
                             pulse_id_conv.u64,
                             config.exptime,
                             config.exptime,
                             config.debug, 1);

    ap_uint<7> frame_cell = 0;
    for (uint32_t f = 0; f < config.frames; f++) {
        ap_uint<1> local_cancel = in_cancel;
        if (local_cancel == 1)
            break;
        if (config.detector_type == SLS_DETECTOR_TYPE_JUNGFRAU) {
        generate_jf_frame:
            for (uint32_t m = 0; m < config.modules; m++) {
            generate_jf_packet:
                for (uint32_t p = 0; p < 128; p++) {
                    generate_packet_jungfrau(data_out,
                                             d_hbm_p0,
                                             d_hbm_p1,
                                             hbm_size_bytes,
                                             f, frame_cell * config.modules + m,
                                             m, config.data_stream,
                                             p,
                                             src_mac_addr,
                                             config.dest_mac_addr,
                                             src_ipv4_addr,
                                             config.dest_ipv4_addr,
                                             pulse_id_conv.u64,
                                             config.exptime,
                                             config.exptime * f,
                                             config.debug, 0);
                }
            }
        } else {
        generate_eiger_frame:
            for (uint32_t m = 0; m < config.modules; m++) {
            generate_eiger_packet:
                for (uint32_t p = 0; p < 16 * config.eiger_bit_depth; p++) {
                    generate_packet_eiger(data_out,
                                          d_hbm_p0,
                                          d_hbm_p1,
                                          hbm_size_bytes,
                                          f,
                                          (config.eiger_bit_depth == 32 ? 2 : 1) * (frame_cell * config.modules + m),
                                          m, config.data_stream,
                                          p,
                                          src_mac_addr,
                                          config.dest_mac_addr,
                                          src_ipv4_addr,
                                          config.dest_ipv4_addr,
                                          config.exptime,
                                          config.exptime * f,
                                          config.debug,
                                          config.eiger_bit_depth,
                                          0);
                }
            }
        }
        pulse_id_conv.d += 1.0;

        if (frame_cell == config.images_in_memory)
            frame_cell = 0;
        else
            frame_cell++;
    }
    return 0;
}