/*
 * Copyright (c) 2019-2023 Beijing Hanwei Innovation Technology Ltd. Co. and
 * its subsidiaries and affiliates (collectly called MKSEMI).
 *
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form, except as embedded into an MKSEMI
 *    integrated circuit in a product or a software update for such product,
 *    must reproduce the above copyright notice, this list of conditions and
 *    the following disclaimer in the documentation and/or other materials
 *    provided with the distribution.
 *
 * 3. Neither the name of MKSEMI nor the names of its contributors may be used
 *    to endorse or promote products derived from this software without
 *    specific prior written permission.
 *
 * 4. This software, with or without modification, must only be used with a
 *    MKSEMI integrated circuit.
 *
 * 5. Any software provided in binary form under this license must not be
 *    reverse engineered, decompiled, modified and/or disassembled.
 *
 * THIS SOFTWARE IS PROVIDED BY MKSEMI "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY, NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL MKSEMI OR CONTRIBUTORS BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include "mk_trace.h"
#include "mk_uwb.h"
#include "mk_misc.h"
#include "mk_power.h"
#include "mk_sleep_timer.h"
#include "lib_ranging.h"

#include "board.h"

#if defined(MK_DS_TWR_INIT)

extern int simple_main(void);

/* Ranging period: 1s */
#define RANGING_PERIOD_MS (1000)

/* This is the delay from the end of the poll frame transmission to the enable of the receiver */
#define POLL_TX_TO_RESP_RX_DLY_US 700U

#define RESP_RX_TO_FINAL_TX_DLY_US 550U

/* Receive response timeout */
#define RESP_RX_TIMEOUT_US 500U

/* Field index in frame */
#define MSG_SEQ_NUM_IDX 2
#define FINAL_MSG_POLL_TX_TS_IDX 10
#define FINAL_MSG_RESP_RX_TS_IDX 14
#define FINAL_MSG_FINAL_TX_TS_IDX 18

struct mk_uwb_configure
{
    uint8_t phy_work_mode; /* PHY_TX / PHY_RX / PHT_TX|PHY_RX */
    struct UWB_CONFIG_T phy_cfg;
};

/* Default communication configuration. */
static struct mk_uwb_configure config = {
    .phy_work_mode = (uint8_t)(PHY_TX | PHY_RX),
    .phy_cfg.ch_num = 5,                      /* Channel number.                           */
    .phy_cfg.code_index = 9,                  /* TRX preamble code                         */
    .phy_cfg.mean_prf = MEAN_PRF_64M,         /* Mean prf 64/128/256M                      */
    .phy_cfg.data_bit_rate = DATA_BR_6M8,     /* Data rate 6.8M                            */
    .phy_cfg.sync_sym = PREAM_LEN_128,        /* Preamble duration, length of preamble 128 */
    .phy_cfg.sfd_sym = BPRF_NSFD2_8,          /* Identifier for SFD sequence               */
    .phy_cfg.ranging_bit = 1,                 /* ranging bit set 1                         */
    .phy_cfg.trx_mode = TRX_MODE_15_4Z_BPRF,  /* IEEE802.15.4z - BPRF mode                 */
    .phy_cfg.sts_pkt_cfg = STS_PKT_CFG_0,     /* SP0 Frame                                 */
    .phy_cfg.sts_segnum = STS_SEGNUM_BPRF_1,  /* Number of STS segments in the frame       */
    .phy_cfg.sts_seglen = STS_SEGLEN_BPRF_64, /* Number of symbols in an STS segment       */
    .phy_cfg.rx_ant_id = UWB_RX_ANT_3,        /* UWB RX antenna port                       */
};

/* Buffer to store received frame */
static uint8_t rx_buf[128];

/*     Frames used in the ranging process
 *     Poll message:
 *     - byte 0 - 1: 0x8841 to indicate a data frame using 16-bit addressing.
 *     - byte 2: sequence number, incremented for each new frame.
 *     - byte 3 - 4: PAN Id 0x4B4d
 *     - byte 5 - 6: Destination address
 *     - byte 7 - 8: Source address
 *     - byte 9: Message type (0x02 RANGING_POLL / 0x03 RANGING_RESPONSE / 0x04 RANGING_FINAL)
 *     Response message:
 *     - byte 10: activity code (0x07 to tell the initiator to go on with the ranging exchange)
 *     Final message:
 *     - byte 10 - 13: poll message transmission timestamp.
 *     - byte 14 - 17: response message reception timestamp.
 *     - byte 18 - 21: final message transmission timestamp.
 */
static uint8_t tx_poll_msg[] = {0x41, 0x88, 0, 0x4D, 0x4B, 0x53, 0x45, 0x4D, 0x49, 0x02};
static uint8_t rx_resp_msg[] = {0x41, 0x88, 0, 0x4D, 0x4B, 0x4D, 0x49, 0x53, 0x45, 0x03, 0x07};
static uint8_t tx_final_msg[] = {0x41, 0x88, 0, 0x4D, 0x4B, 0x53, 0x45, 0x4D, 0x49, 0x04, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};

/* Count value of phy counter when transmitting and receiving frames */
static uint32_t poll_tx_en_start_u32;
static uint32_t resp_rx_en_start_u32;
static uint32_t final_tx_en_start_u32;

/* 41 bits timestamps of frames transmission/reception. */
static int64_t poll_tx_ts_i64;
static int64_t resp_rx_ts_i64;
static int64_t final_tx_ts_i64;

/* Frame sequence number, incremented after each transmission. */
static uint8_t frame_seq_nb = 0;

static volatile uint16_t rx_state;
static volatile uint16_t rx_length;

enum SIMPLE_FSM_T
{
    SIMPLE_IDLE = 0,
    SIMPLE_POLL = 1,
    SIMPLE_RESPONSE = 2,
    SIMPLE_FINAL = 3,
};

static enum SIMPLE_FSM_T state = SIMPLE_IDLE;

/**
 * @brief Correct TX timestamp of the ranging frame.
 *
 * @param[in] timestamp    PHY timer count of TX
 * @return TX timestamp (unit: 15.65ps)
 */
static int64_t ranging_tx_time_correct(uint32_t timestamp)
{
    int64_t tx_timestamp = ranging_tx_time(timestamp);

    // correct antenna delay (TX using the same antenna as RX)
    tx_timestamp += ranging_ant_delays_get(config.phy_cfg.rx_ant_id) / 2;

    return tx_timestamp;
}

/**
 * @brief Correct RX timestamp of the ranging frame.
 *
 * @param[in] ind    MAC RX report
 * @return RX timestamp (unit: 15.65ps)
 */
static int64_t ranging_rx_time_correct(const struct MAC_HW_REPORT_T *ind)
{
    int64_t rx_timestamp = ranging_rx_time(ind);

    // correct antenna delay
    rx_timestamp -= ranging_ant_delays_get(config.phy_cfg.rx_ant_id) / 2;

    return rx_timestamp;
}

/* RX done process handler. */
static void rx_int_callback(struct MAC_HW_REPORT_T *rx_report)
{
    // Power off radio
    power_off_radio();

    rx_state = rx_report->err_code;

    /** UWB RX success */
    if (rx_state == UWB_RX_OK)
    {
        /* Received data does not contain FCS */
        rx_length = rx_report->pkt_len;
        memcpy(rx_buf, rx_report->pkt_data, rx_length);
        /* Calculate rx timestamp */
        resp_rx_ts_i64 = ranging_rx_time_correct(rx_report);
    }
    else
    {
        /* UWB_PLD_ERR     payload error             */
        /* UWB_PHR_ERR     PHR error                 */
        /* UWB_SFD_ERR     Sfd error                 */
        /* UWB_BD_ERR      Preamble detection error  */
        /* UWB_TO_ERR      Receive timeout           */
        /* UWB_STS_ERR     STS error                 */
        rx_length = 0;
        memset(rx_buf, 0, sizeof(rx_buf));
    }
}

/* TX done process handler. */
static void tx_int_callback(struct MAC_HW_REPORT_T *tx_report)
{
    // Power off radio
    power_off_radio();

    /** UWB TX success */
    if (tx_report->err_code == UWB_TX_OK)
    {
    }
}

int simple_main(void)
{
    // The following peripherals will be initialized in the uwb_open function
    // phy/mac/aes/lsp/phy timers initialized
    uwb_open();

    // Set calibration parameters
    uwb_calibration_params_set(config.phy_cfg.ch_num);

    // set advanced parameters
    struct PHY_ADV_CONFIG_T adv_config = {
        .thres_fap_detect = 40,
        .nth_scale_factor = 4,
        .ranging_performance_mode = 0,
        .skip_weakest_port_en = 0,
    };
    phy_adv_params_configure(&adv_config);

    // uwb configure
    uwb_configure(config.phy_work_mode, board_param.tx_power_fcc[CALIB_CH(config.phy_cfg.ch_num)], &config.phy_cfg);

    ranging_lib_init();
    ranging_frame_type_set(config.phy_cfg.sts_pkt_cfg);

    // Register rx interrupt callback function
    mac_register_process_handler(tx_int_callback, rx_int_callback);

    // Initialize low power mode
    power_init();

#if LOW_POWER_EN
    // Enable sleep timer
    sleep_timer_open(true, SLEEP_TIMER_MODE_ONESHOT, NULL);
#endif

    state = SIMPLE_POLL;
    poll_tx_en_start_u32 = phy_timer_count_get();

    while (1)
    {
        switch (state)
        {
            case SIMPLE_IDLE:
            case SIMPLE_POLL:
            {
                // Target time equals 0 means send immediately
                // Program the MAC to transmit UWB packet
                // Send Poll
                tx_poll_msg[MSG_SEQ_NUM_IDX] = frame_seq_nb;

                poll_tx_en_start_u32 += MS_TO_PHY_TIMER_COUNT(RANGING_PERIOD_MS);
                uwb_tx(tx_poll_msg, sizeof(tx_poll_msg), 1, poll_tx_en_start_u32);

                /* Calculate tx timestamp */
                poll_tx_ts_i64 = ranging_tx_time_correct(poll_tx_en_start_u32 + phy_shr_duration());

                // Check the MAC busy state
                while (mac_is_busy())
                {
                }

                state = SIMPLE_RESPONSE;
                /* Increment frame sequence number after transmission of the poll message */
                frame_seq_nb++;
            }
            break;

            case SIMPLE_RESPONSE:
            {
                // Receive Respond
                resp_rx_en_start_u32 = poll_tx_en_start_u32 + US_TO_PHY_TIMER_COUNT(POLL_TX_TO_RESP_RX_DLY_US);
                uwb_rx(1, resp_rx_en_start_u32, RESP_RX_TIMEOUT_US);

                // Check the MAC busy state
                while (mac_is_busy())
                {
                }

                if (rx_state == UWB_RX_OK)
                {
                    state = SIMPLE_FINAL;
                }
                else
                {
                    state = SIMPLE_POLL;
                }
            }
            break;

            case SIMPLE_FINAL:
            {
                rx_buf[MSG_SEQ_NUM_IDX] = rx_resp_msg[MSG_SEQ_NUM_IDX];
                if (0 == memcmp(rx_resp_msg, rx_buf, rx_length))
                {
                    final_tx_en_start_u32 = resp_rx_en_start_u32 + US_TO_PHY_TIMER_COUNT(RESP_RX_TO_FINAL_TX_DLY_US);
                    final_tx_ts_i64 = ranging_tx_time_correct(final_tx_en_start_u32 + phy_shr_duration());

                    // Write all timestamps in the final message
                    int32_t ts1, ts2, ts3;
                    ts1 = (int32_t)poll_tx_ts_i64;
                    for (uint8_t i = 0; i < 4; i++)
                    {
                        tx_final_msg[FINAL_MSG_POLL_TX_TS_IDX + i] = (uint8_t)ts1;
                        ts1 >>= 8;
                    }

                    ts2 = (int32_t)resp_rx_ts_i64;
                    for (uint8_t i = 0; i < 4; i++)
                    {
                        tx_final_msg[FINAL_MSG_RESP_RX_TS_IDX + i] = (uint8_t)ts2;
                        ts2 >>= 8;
                    }

                    ts3 = (int32_t)final_tx_ts_i64;
                    for (uint8_t i = 0; i < 4; i++)
                    {
                        tx_final_msg[FINAL_MSG_FINAL_TX_TS_IDX + i] = (uint8_t)ts3;
                        ts3 >>= 8;
                    }

                    // Send Final
                    tx_final_msg[MSG_SEQ_NUM_IDX] = frame_seq_nb;
                    uwb_tx(tx_final_msg, sizeof(tx_final_msg), 1, final_tx_en_start_u32);

                    // Check the MAC busy state
                    while (mac_is_busy())
                    {
                    }

                    /* Increment frame sequence number after transmission of the final message */
                    frame_seq_nb++;
                }
                state = SIMPLE_POLL;
            }
            break;
        }

        if (state == SIMPLE_POLL)
        {
            LOG_INFO(TRACE_MODULE_APP, "SEQ NUM: %u\r\n", (uint8_t)(frame_seq_nb - 1));
#if LOW_POWER_EN
            trace_flush();
            uint32_t lock = int_lock();
            sleep_timer_start(__MS_TO_32K_CNT(RANGING_PERIOD_MS - 4));
            power_enter_power_down_mode(0);
            int_unlock(lock);
#else
            /* Execute a delay between ranging exchanges. */
            sys_timer_delay_ms(RANGING_PERIOD_MS - 2);
#endif
        }
    }
}
#endif