/* 
 * 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_decode_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_decode::sptr
    srs3_decode::make(unsigned int size, bool crc, bool cc, bool rand, bool rs, bool v1compat)
    {
      return gnuradio::get_initial_sptr
        (new srs3_decode_impl(size, crc, cc, rand, rs, v1compat));
    }

    void
    srs3_decode_impl::decode(pmt::pmt_t pdu)
    {
        uint8_t frame_data[SPACEFRAME_LENGTH_MAX];
        size_t frame_length = 0;
        int errs, cc_errs = 0, rs_errs = 0;

        /* Copy frame data to output */
        pmt::pmt_t vector = pmt::cdr(pdu);
        frame_length = pmt::blob_length(vector);
        memcpy(frame_data, pmt::blob_data(vector), frame_length);

        /* Viterbi decode */
        if (d_cc) {
            size_t conv_length = (frame_length / SPACEFRAME_CC_RATE) - SPACEFRAME_CC_TAIL;
            /*
             * The sender always starts encoding a frame from state 0, but the
             * first 4 input bytes are the ASM. Since we sync to the last bits
             * of the _encoded_ ASM, we need to start decoding from the state
             * the encoder is in after being primed with the ASM. This
             * corresponds to the last 6 bits of the ASM: 011101 = 1D
             */
            int starting_state = d_v1compat ? 0 : 0x1d;
            init_viterbi27(vp, starting_state);
            update_viterbi27_blk(vp, frame_data, (conv_length * 8) + (SPACEFRAME_CC_CONSTR - 1));
            cc_errs = chainback_viterbi27(vp, frame_data, conv_length * 8, 0);
            frame_length = conv_length;
        }

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

        /* Check RS parity */
        if (d_rs) {
            /*
             * We follow the CCSDS 131.0-B-3 interleaving specification, but
             * allow frame sizes that are not a multiple of the interleaving
             * depth I. In that case, the last 1 to I-1 blocks have an
             * additional virtual fill byte.
             *
             * See Figure 4-2 in CCSDS 131.0-B-3 or Figure 9 in
             * NASA/JPL/DSN 810-005-208, Rev. B for explanation of the
             * interleaving.
             */
            unsigned int rs_size, rs_blocks, rs_block_size_cur, rs_pad;
            uint8_t rs_data[SPACEFRAME_RS_DATA_SIZE + SPACEFRAME_RS_PAR_SIZE];

            /* Calculate number of RS data bytes */
            rs_size = d_size + (d_crc ? sizeof(uint32_t) : 0);

            /* Calculate number of RS blocks required to fit data */
            rs_blocks = (rs_size + 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 = (rs_size + rs_blocks - 1) / rs_blocks;

                /* Calculate remaining _data_ bytes */
                unsigned int rs_remain = rs_size;

                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;

                    /* Copy data and parity bytes to check array */
                    memcpy(rs_data, rs_data_v1, rs_block_size_cur + SPACEFRAME_RS_PAR_SIZE);

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

                    /* Check parity */
                    errs = decode_rs_ccsds(rs_data, NULL, 0, rs_pad);
                    if (errs < 0)
                        return;
                    rs_errs += errs;

                    /* Copy corrected data back to frame */
                    memcpy(rs_data_v1, rs_data, rs_block_size_cur + SPACEFRAME_RS_PAR_SIZE);

                    /* Remove parity bytes and advance to next block boundary */
                    memmove(rs_data_v1 + rs_block_size_cur,
                            rs_data_v1 + rs_block_size_cur + SPACEFRAME_RS_PAR_SIZE,
                            rs_remain + (rs_blocks - i - 1) * SPACEFRAME_RS_PAR_SIZE);
                    rs_data_v1 += rs_block_size_cur;
                }
            } else {
                for (int i = 0; i < rs_blocks; i++) {
                    /* Keep track of current block size */
                    rs_block_size_cur = 0;

                    /* Copy data bytes to check array */
                    for (int j = 0; j < SPACEFRAME_RS_DATA_SIZE && i + j * rs_blocks < rs_size; j++) {
                        rs_data[j] = frame_data[i + j * rs_blocks];
                        rs_block_size_cur++;
                    }

                    /* Copy parity bytes to check array */
                    for (int j = 0; j < SPACEFRAME_RS_PAR_SIZE; j++)
                        rs_data[j + rs_block_size_cur] = frame_data[rs_size + i + j * rs_blocks];

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

                    /* Check parity */
                    errs = decode_rs_ccsds(rs_data, NULL, 0, rs_pad);
                    if (errs < 0)
                        return;
                    rs_errs += errs;

                    /* Copy corrected data back to frame */
                    for (int j = 0; j < rs_block_size_cur; j++)
                        frame_data[i + j * rs_blocks] = rs_data[j];
                }
            }

            frame_length -= rs_blocks * SPACEFRAME_RS_PAR_SIZE;
        }

        /* Check CRC */
        if (d_crc) {
            uint32_t crc32, crc32_calc;
            frame_length -= sizeof(crc32);
            memcpy(&crc32, &frame_data[frame_length], sizeof(crc32));
            crc32 = be32toh(crc32);
            crc32_calc = crc32c(frame_data, frame_length);
            if (crc32 != crc32_calc)
                return;
        }

        /* Check if frame contains all-zero data (OID/idle frame) */
        bool oid = true;
        for (int i = 0; i < frame_length; i++) {
            if (frame_data[i] != 0x00) {
                oid = false;
                break;
            }
        }

        /* Drop OID frames */
        if (oid)
            return;

        /* Add error counters as PMT meta */
        pmt::pmt_t meta = pmt::is_dict(pmt::car(pdu)) ? pmt::car(pdu) : pmt::make_dict();
        pmt::pmt_t key, val;

        key = pmt::intern("cc_errs");
        val = pmt::from_long(cc_errs);
        meta = pmt::dict_add(meta, key, val);

        key = pmt::intern("rs_errs");
        val = pmt::from_long(rs_errs);
        meta = pmt::dict_add(meta, key, val);

        /* Create PMT blob */
        pmt::pmt_t vector_out = pmt::make_blob(frame_data, frame_length);
        pmt::pmt_t msg = pmt::cons(meta, vector_out);

        /* Route frame to correct output */
        message_port_pub(pmt::mp("output"), msg);
    }

    void srs3_decode_impl::set_size(unsigned int bytes)
    {
        d_size = bytes;
    }

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

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

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

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

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

    /*
     * The private constructor
     */
    srs3_decode_impl::srs3_decode_impl(unsigned int size, bool crc, bool cc, bool rand, bool rs, bool v1compat)
      : gr::block("srs3_decode",
              gr::io_signature::make(0, 0, 0),
              gr::io_signature::make(0, 0, 0)),
        d_size(size),
        d_crc(crc),
        d_cc(cc),
        d_rand(rand),
        d_rs(rs),
        d_v1compat(v1compat)
    {
        /* Create Viterbi decoder */
        vp = create_viterbi27(SPACEFRAME_LENGTH_MAX * 8);
        if (!vp) {
            throw std::runtime_error("failed to allocate Viterbi decoder");
        }

        message_port_register_out(pmt::mp("output"));

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

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

