/* 
 * Copyright (c) 2021 Satlab A/S <satlab@satlab.com>
 * 
 * This is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3, or (at your option)
 * any later version.
 * 
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this software; see the file COPYING.  If not, write to
 * the Free Software Foundation, Inc., 51 Franklin Street,
 * Boston, MA 02110-1301, USA.
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <endian.h>
#include <gnuradio/io_signature.h>
#include "srs3_encode_impl.h"

extern "C" {
  #include "fec/crc32c.h"
  #include "fec/randomizer.h"
  #include "fec/rs.h"
  #include "fec/viterbi.h"
  #include "spaceframe.h"
}

namespace gr {
  namespace satlab {

    srs3_encode::sptr
    srs3_encode::make(unsigned int preamble, bool crc, bool cc, bool rand, bool rs, bool v1compat)
    {
      return gnuradio::get_initial_sptr
        (new srs3_encode_impl(preamble, crc, cc, rand, rs, v1compat));
    }

    void
    srs3_encode_impl::encode(pmt::pmt_t pdu)
    {
        uint32_t syncword = htobe32(SPACEFRAME_ASM);
        uint32_t crc32;

        /* Add one byte for postamble */
        uint8_t frame_data[SPACEFRAME_PREAMBLE_LENGHT_MAX + SPACEFRAME_LENGTH_MAX + 1];
        size_t frame_length;

        pmt::pmt_t vector = pmt::cdr(pdu);
        frame_length = pmt::blob_length(vector);

        /* Copy frame data to output */
        memcpy(frame_data, pmt::blob_data(vector), frame_length);

        /* Add CRC */
        if (d_crc) {
                crc32 = crc32c(frame_data, frame_length);
                crc32 = htobe32(crc32);
                memcpy(&frame_data[frame_length], &crc32, sizeof(crc32));
                frame_length += sizeof(crc32);
        }

        /* Add RS partity */
        if (d_rs) {
            /* See srs3_decode_impl.cc for explanation on interleaving */

            unsigned int rs_blocks, rs_block_size_cur, rs_pad;
            uint8_t rs_data[SPACEFRAME_RS_DATA_SIZE];
            uint8_t rs_parity[SPACEFRAME_RS_PAR_SIZE];

            /* Calculate number of RS blocks required to fit data */
            rs_blocks = (frame_length + SPACEFRAME_RS_DATA_SIZE - 1) / SPACEFRAME_RS_DATA_SIZE;

            if (d_v1compat) {
                uint8_t *rs_data_v1 = frame_data;

                /* Divide data in (almost) equal size blocks to distribute parity */
                unsigned int rs_block_size = (frame_length + rs_blocks - 1) / rs_blocks;

                /* Start with first block */
                unsigned int rs_remain = frame_length;

                for (int i = 0; i < rs_blocks; i++) {
                    /* Calculate current block size */
                    rs_block_size_cur = std::min(rs_remain, rs_block_size);

                    /* Consume current block size bytes */
                    rs_remain -= rs_block_size_cur;

                    /* Make room for parity bytes */
                    memmove(rs_data_v1 + rs_block_size_cur + SPACEFRAME_RS_PAR_SIZE, rs_data_v1 + rs_block_size_cur, rs_remain);

                    /* Copy data to input buffer */
                    memcpy(rs_data, rs_data_v1, rs_block_size_cur);

                    /* Calculate padding of current block */
                    rs_pad = SPACEFRAME_RS_DATA_SIZE - rs_block_size_cur;

                    /* Calculate parity */
                    encode_rs_ccsds(rs_data, rs_parity, rs_pad);

                    /* Copy parity to output frame buffer */
                    memcpy(rs_data_v1 + rs_block_size_cur, rs_parity, SPACEFRAME_RS_PAR_SIZE);
                    rs_data_v1 += rs_block_size_cur + SPACEFRAME_RS_PAR_SIZE;
                }
            } else {
                for (int i = 0; i < rs_blocks; i++) {
                    /* Keep track of current block size */
                    rs_block_size_cur = 0;

                    /* Copy data to interleave buffer */
                    for (int j = 0; j < SPACEFRAME_RS_DATA_SIZE && i + j * rs_blocks < frame_length; j++) {
                        rs_data[j] = frame_data[i + j * rs_blocks];
                        rs_block_size_cur++;
                    }

                    /* Calculate padding of current block */
                    rs_pad = SPACEFRAME_RS_DATA_SIZE - rs_block_size_cur;

                    /* Calculate parity */
                    encode_rs_ccsds(rs_data, rs_parity, rs_pad);

                    /* Copy parity to output frame buffer */
                    for (int j = 0; j < SPACEFRAME_RS_PAR_SIZE; j++)
                        frame_data[frame_length + i + j * rs_blocks] = rs_parity[j];
                }
            }

            /* Increase frame size */
            frame_length += rs_blocks * SPACEFRAME_RS_PAR_SIZE;
        }

        /* Randomize */
        if (d_rand) {
            randomize_xor(frame_data, frame_length);
        }

        /* Add ASM if not in compatibility mode */
        if (!d_v1compat) {
            memmove(&frame_data[sizeof(syncword)], frame_data, frame_length);
            memcpy(frame_data, &syncword, sizeof(syncword));
            frame_length = frame_length + sizeof(syncword);
        }

        /* Add convolutional code */
        if (d_cc) {
            /* Move data to end of buffer to avoid overwrite by inline encode_viterbi */
            uint8_t *enc_data = &frame_data[frame_length + 1];
            memcpy(enc_data, frame_data, frame_length);

            encode_viterbi(frame_data, enc_data, frame_length * 8);
            frame_length = (frame_length + SPACEFRAME_CC_TAIL) * SPACEFRAME_CC_RATE;
        }

        /* Add ASM if in compatibility mode */
        if (d_v1compat) {
            memmove(&frame_data[sizeof(syncword)], frame_data, frame_length);
            memcpy(frame_data, &syncword, sizeof(syncword));
            frame_length = frame_length + sizeof(syncword);
        }

        /* Add preamble */
        if (d_preamble > 0)  {
            if (d_preamble > SPACEFRAME_PREAMBLE_LENGHT_MAX)
                d_preamble = SPACEFRAME_PREAMBLE_LENGHT_MAX;
            memmove(&frame_data[d_preamble], frame_data, frame_length);
            memset(&frame_data[0], 0x55, d_preamble);
            frame_length += d_preamble;
        }

        pmt::pmt_t vector_out = pmt::make_blob(frame_data, frame_length);
        pmt::pmt_t msg = pmt::cons(pmt::car(pdu), vector_out);

        message_port_pub(pmt::mp("output"), msg);
    }

    void srs3_encode_impl::set_preamble(unsigned int bytes)
    {
        d_preamble = bytes;
    }

    void srs3_encode_impl::set_crc(bool crc)
    {
        d_crc = crc;
    }

    void srs3_encode_impl::set_cc(bool cc)
    {
        d_cc = cc;
    }

    void srs3_encode_impl::set_rand(bool rand)
    {
        d_rand = rand;
    }

    void srs3_encode_impl::set_rs(bool rs)
    {
        d_rs = rs;
    }

    void srs3_encode_impl::set_v1compat(bool v1compat)
    {
        d_v1compat = v1compat;
    }

    /*
     * The private constructor
     */
    srs3_encode_impl::srs3_encode_impl(unsigned int preamble, bool crc, bool cc, bool rand, bool rs, bool v1compat)
      : gr::block("srs3_encode",
              gr::io_signature::make(0, 0, 0),
              gr::io_signature::make(0, 0, 0)),
        d_preamble(preamble),
        d_crc(crc),
        d_cc(cc),
        d_rand(rand),
        d_rs(rs),
        d_v1compat(v1compat)
    {
        message_port_register_out(pmt::mp("output"));

        message_port_register_in(pmt::mp("input"));
        set_msg_handler(pmt::mp("input"), [this](pmt::pmt_t msg) {
            this->encode(msg);
        });
    }

    /*
     * Our virtual destructor.
     */
    srs3_encode_impl::~srs3_encode_impl()
    {
    }
  } /* namespace satlab */
} /* namespace gr */

