/*
 * 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_power.h"
#include "mk_gpio.h"
#include "mk_uwb.h"
#include "mk_clock.h"

#include "wsf_msg.h"

#include "ranging_ss_twr.h"
#include "lib_aoa.h"
#include "lib_ranging.h"
#if KF_EN
#include "lib_kf.h"
#endif
// #include "board.h"

#if KF_EN && FILTER_EN
#define KF_SUPPORT_NUM 6
#define KF_TIMEOUT_MS 2000
static struct KF_MAC_ADDR_T kf_mac_addr_cache[KF_SUPPORT_NUM];
static struct KF_CHANNEL_CACHE_T kf_channel_cache[KF_SUPPORT_NUM];
static struct KF_MAT_VALUE_CACHE_T kf_mat_value_cache[KF_SUPPORT_NUM];
#endif

static uint16_t user_poll_data_len = 0;
static uint8_t user_poll_data[MSG_USER_DATA_LEN];

static uint16_t user_response_data_len = 0;
static uint8_t user_response_data[MSG_USER_DATA_LEN];

struct RANGING_USER_PKT_T ranging_user_pkt;

struct RANGING_ENV_T ranging_env;

static struct RANGING_CB_T ranging_cb;

static struct UWB_OP_T op = {
    .session_configure = ranging_configure,
    .session_start = ranging_start,
    .session_stop = ranging_stop,
    .session_local_addr_set = ranging_local_addr_set,
    .session_peer_addr_set = ranging_peer_addr_set,
    .session_responder_addr_add = NULL,
    .session_responder_list_clr = NULL,
    .session_dynamic_update_responder_list = NULL,
    .session_set_ccc_ursk = NULL,
};

static void ranging_timer_callback(void *dev, uint32_t time);

static void uwb_pkt_tx_done_ind(const struct MAC_HW_REPORT_T *tx, enum RANGING_STAGE_T stage);
static void uwb_pkt_rx_done_ind(const struct MAC_HW_REPORT_T *rx, enum RANGING_STAGE_T stage);

static void ranging_tx_process(struct MAC_HW_REPORT_T *tx_report);
static void ranging_rx_process(struct MAC_HW_REPORT_T *rx_report);
void app_session_init(void);

//------------------------------------------------------------------------------
int ranging_init(uint8_t handle_id)
{
    /* store handler ID */
    ranging_cb.handle_id = handle_id;

    /* init rx queue */
    WSF_QUEUE_INIT(&ranging_cb.msg_queue);

    ranging_cb.daemon_timer.handlerId = handle_id;
    ranging_cb.daemon_timer.msg.event = RANGING_DAEMON_TIMER_MSG;

    LOG_INFO(TRACE_MODULE_APP, "Ranging lib version: %s\r\n", MK8000_get_rangelib_version());
    LOG_INFO(TRACE_MODULE_APP, "AoA lib version: %s\r\n", MK8000_get_aoalib_version());

    return 0;
}

int ranging_deinit(void)
{
    return 0;
}

// This function will be called by uwbapi_session_init()
void app_session_init(void)
{
    // register process handler for MAC TX done and RX done
    mac_register_process_handler(ranging_tx_process, ranging_rx_process);

    uwbs_handler_init(&op);
}

void ranging_configure(void)
{
    ranging_env.ranging_period = MS_TO_PHY_TIMER_COUNT(uwb_app_config.session_param.ranging_interval);
    ranging_env.slot_interval = RSTU_TO_PHY_TIMER_COUNT(uwb_app_config.session_param.slot_duration);
    ranging_env.enable = 0;
    ranging_env.tof = 0;
    ranging_env.range_data.measurements_num = 0;
    for (uint8_t i = 0; i < MEASUREMENT_NUM_MAX; i++)
    {
        ranging_env.range_data.measurements[i].status = STATUS_VENDOR_RESERVED;
    }

    uwbs_configure(PHY_TX | PHY_RX, uwb_app_config.session_param.tx_power_level);

#if CSI_EN
    ranging_aux_out_opt_set(CH_LEN_DEFAULT, 3);
#endif

#if AOA_EN
    aoa_aux_info_set(AOA_AUX_ANT_IQ_RSSI_PDOA_AOA_FOM);
    aoa_steering_vector_set((const float *)((uint32_t)((uwb_app_config.ppdu_params.ch_num == 9) ? svec_ch9_ptr : svec_ch5_ptr) | SRAM_BASE));
#else
    aoa_aux_info_set(AOA_AUX_ANT_IQ_RSSI);
#endif

    aoa_param_config();

#if FILTER_EN
    if (uwb_app_config.filter_en)
    {
#if KF_EN
        loc_post_kf_config(uwb_app_config.session_param.ranging_interval, kf_mac_addr_cache, kf_channel_cache, kf_mat_value_cache, KF_SUPPORT_NUM,
                           KF_TIMEOUT_MS);
#else
        loc_post_filter_config(uwb_app_config.session_param.ranging_interval, 1, uwb_app_config.session_param.aoa_result_req);
#endif
    }
#endif

    ranging_user_pkt.session_id = (uint16_t)uwb_app_config.session_id;
}

void ranging_start(void)
{
    ranging_env.enable = 1;
    ranging_env.anchor_point = phy_timer_count_get();
    enum DEV_ROLE_T dev_role = uwb_app_config.session_param.device_role;
    if (dev_role == DEV_ROLE_INITIATOR)
    {
        uwb_app_config.ranging_stage = RANGING_POLL;
        phy_timer_target_set(ranging_env.anchor_point + ranging_env.slot_interval - UWB_PERIOD_PREFETCH_TIME, ranging_timer_callback);
        power_mode_request(POWER_UNIT_APP, POWER_MODE_POWER_DOWN);
    }
    else
    {
        uwb_app_config.ranging_stage = RANGING_SYNC;
        phy_timer_target_set(ranging_env.anchor_point + ranging_env.slot_interval - UWB_PERIOD_PREFETCH_TIME, ranging_timer_callback);
        power_mode_request(POWER_UNIT_APP, POWER_MODE_SLEEP);
    }

    ranging_env.count = 0;
    ranging_env.count_last = 0;
    ranging_env.lost_cnt = 0;

    WsfTimerStartMs(&ranging_cb.daemon_timer, uwb_app_config.session_param.ranging_interval * 10, WSF_TIMER_PERIODIC);
    LOG_INFO(TRACE_MODULE_APP, "Ranging start, role %d\r\n", dev_role);
}

void ranging_stop(void)
{
    ranging_env.enable = 0;
    WsfTimerStop(&ranging_cb.daemon_timer);
}

void ranging_restart(void)
{
    ranging_stop();
    mac_restart();
    ranging_start();
    LOG_INFO(TRACE_MODULE_APP, "Ranging restart\r\n");
}

int8_t ranging_tx_power_get(void)
{
    return uwb_tx_power_get(uwb_app_config.ppdu_params.ch_num, uwb_app_config.session_param.tx_power_level);
}

void ranging_local_addr_set(uint16_t short_addr)
{
    uwbs_local_short_addr_set(short_addr);
}

void ranging_peer_addr_set(uint16_t short_addr)
{
    uwbs_peer_short_addr_set(short_addr);
}

uint8_t ranging_poll_msg_set(uint16_t data_len, const uint8_t *data)
{
    user_poll_data_len = sizeof(user_poll_data);

    if (user_poll_data_len > data_len)
    {
        user_poll_data_len = data_len;
    }

    for (uint8_t ii = 0; ii < user_poll_data_len; ii++)
    {
        user_poll_data[ii] = data[ii];
    }

    // The return value is 0, the poll message is set successfully
    // The return value is 1 is set failed
    return (user_poll_data_len == data_len) ? 0 : 1;
}

uint8_t ranging_response_msg_set(uint16_t data_len, const uint8_t *data)
{
    user_response_data_len = sizeof(user_response_data);

    if (user_response_data_len > data_len)
    {
        user_response_data_len = data_len;
    }

    for (uint8_t ii = 0; ii < user_response_data_len; ii++)
    {
        user_response_data[ii] = data[ii];
    }

    // The return value is 0, the response message is set successfully
    // The return value is 1 is set failed
    return (user_response_data_len == data_len) ? 0 : 1;
}

// ts_a - ts_b
int64_t ranging_timestamp_diff(int64_t ts_a, int64_t ts_b)
{
    if (ts_a < ts_b)
    {
        return (0x20000000000 - ts_b + ts_a);
    }
    else
    {
        return (ts_a - ts_b);
    }
}

// unit: 15.6ps
int64_t ranging_tround(void)
{
    int64_t tx_timestamp = ranging_env.tx_poll_time;
    int64_t rx_timestamp = ranging_env.rx_response_time;

    int64_t tround = ranging_timestamp_diff(rx_timestamp, tx_timestamp);

    // correct antenna delay
    tround -= ranging_ant_delays_get(uwb_app_config.ppdu_params.rx_ant_id);

    return tround;
}

// unit: 15.6ps
int64_t ranging_treply(void)
{
    int64_t tx_timestamp = ranging_env.tx_response_time;
    int64_t rx_timestamp = ranging_env.rx_poll_time;

    int64_t treply = ranging_timestamp_diff(tx_timestamp, rx_timestamp);

    // correct antenna delay
    treply += ranging_ant_delays_get(uwb_app_config.ppdu_params.rx_ant_id);

    return treply;
}

#if defined(__ARMCC_VERSION) && (__ARMCC_VERSION >= 6010050)
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wswitch-enum"
#pragma clang diagnostic ignored "-Wunreachable-code-break"
#pragma clang diagnostic ignored "-Wcast-align"
#endif

static void ranging_timer_callback(void *dev, uint32_t time)
{
    // board_led_on(BOARD_LED_2);
    if (ranging_env.enable)
    {
        enum RANGING_STAGE_T stage = uwb_app_config.ranging_stage;
        if (uwb_app_config.session_param.device_role == DEV_ROLE_INITIATOR)
        {
            switch (stage)
            {
                case RANGING_IDLE:
                    ASSERT(0, "RANGING_IDLE\r\n");
                    break;

                case RANGING_POLL:
                {
                    ranging_env.anchor_point += ranging_env.ranging_period;
                    ranging_user_pkt.seq_num++;

                    // send poll
                    ranging_user_pkt.msg_type = RANGING_POLL;
                    memcpy(ranging_user_pkt.msg.poll_msg, user_poll_data, user_poll_data_len);

                    power_on_radio(1, 0);
                    mac_tx(EVT_MODE_MAC_ASAP_PHY_FIX, ranging_env.anchor_point, 0, (uint8_t *)&ranging_user_pkt, user_poll_data_len + MSG_HEADER_LEN);

                    // Store tx timestamp
                    ranging_env.tx_poll_time = ranging_tx_time(ranging_env.anchor_point + phy_shr_duration());
                }
                break;

                case RANGING_RESPONSE:
                {
                    power_on_radio(0, 1);
                    // receive response
                    mac_rx(EVT_MODE_MAC_ASAP_PHY_FIX, ranging_env.anchor_point + ranging_env.slot_interval - UWB_RX_OPEN_IN_ADVANCE, UWB_RX_WINDOW);
                }
                break;

                default:
                    ASSERT(0, "UNKNOWN STAGE\r\n");
                    break;
            }
        }
        else
        {
            switch (stage)
            {
                case RANGING_IDLE:
                    ASSERT(0, "RANGING_IDLE\r\n");
                    break;

                case RANGING_SYNC:
                {
                    power_on_radio(0, 1);
                    // receive sync
                    mac_rx(EVT_MODE_MAC_PHY_ASAP, 0, ranging_env.ranging_period);
                }
                break;

                case RANGING_POLL:
                {
                    ranging_env.anchor_point += ranging_env.ranging_period;

                    power_on_radio(0, 1);
                    // receive poll
                    mac_rx(EVT_MODE_MAC_ASAP_PHY_FIX, ranging_env.anchor_point - UWB_RX_OPEN_IN_ADVANCE, UWB_RX_WINDOW);
                }
                break;

                case RANGING_RESPONSE:
                {
                    ranging_user_pkt.msg_type = RANGING_RESPONSE;

                    // Store tx timestamp
                    ranging_env.tx_response_time = ranging_tx_time(ranging_env.anchor_point + ranging_env.slot_interval + phy_shr_duration());

                    // update reply time in response packet
                    int64_t treply = ranging_treply();
                    for (int i = 0; i < 5; i++)
                    {
                        ranging_user_pkt.msg.response_msg[MSG_TREPLY_IDX - MSG_HEADER_LEN + i] = (treply >> (i * 8)) & 0xff;
                    }
                    memcpy(&ranging_user_pkt.msg.response_msg[MSG_TIMESTAMP_LEN], user_response_data, user_response_data_len);

                    power_on_radio(1, 0);
                    // send response
                    mac_tx(EVT_MODE_MAC_ASAP_PHY_FIX, ranging_env.anchor_point + ranging_env.slot_interval, 0, (uint8_t *)&ranging_user_pkt,
                           user_response_data_len + MSG_HEADER_LEN + MSG_TIMESTAMP_LEN);
                }
                break;

                default:
                    ASSERT(0, "UNKNOWN STAGE\r\n");
                    break;
            }
        }

        mac_start();
    }
    else
    {
        power_mode_clear(POWER_UNIT_APP);
        LOG_INFO(TRACE_MODULE_APP, "Ranging stop\r\n");
    }
    // board_led_off(BOARD_LED_2);
}

enum RANGING_STAGE_T ranging_fsm(const struct MAC_HW_REPORT_T *ind)
{
    struct RANGING_USER_PKT_T *ranging_recv_pkt = (struct RANGING_USER_PKT_T *)&ind->pkt_data[0];
    enum RANGING_STAGE_T stage = uwb_app_config.ranging_stage;
    // LOG_INFO(TRACE_MODULE_APP, "ranging_fsm %d\r\n", stage);
    // board_led_on(BOARD_LED_2);

    power_off_radio();
    ranging_env.count++;

    if (uwb_app_config.session_param.device_role == DEV_ROLE_INITIATOR)
    {
        switch (stage)
        {
            case RANGING_IDLE:
                break;

            case RANGING_POLL:
            {
                // process tx-poll packet
                uwb_pkt_tx_done_ind(ind, stage);

                uwb_app_config.ranging_stage = RANGING_RESPONSE;
            }
            break;

            case RANGING_RESPONSE:
            {
                // process rx-response packet
                uwb_pkt_rx_done_ind(ind, stage);

                if ((ind->err_code == UWB_RX_OK) && (ranging_recv_pkt->session_id == ranging_user_pkt.session_id) &&
                    (ranging_recv_pkt->msg_type == RANGING_RESPONSE))
                {
                    ranging_env.rx_response_time = ranging_rx_time(ind);
                }

                uwb_app_config.ranging_stage = (RANGING_POLL);
            }
            break;

            default:
                break;
        }
    }
    else
    {
        switch (stage)
        {
            case RANGING_IDLE:
                break;
            case RANGING_SYNC:
            {
                if ((ind->err_code == UWB_RX_OK) && (ranging_recv_pkt->session_id == ranging_user_pkt.session_id) &&
                    (ranging_recv_pkt->msg_type == RANGING_POLL))
                {
                    // update anchor point
                    ranging_env.anchor_point = ind->timestamp - phy_shr_duration();
                    ranging_env.rx_poll_time = ranging_rx_time(ind);
                    ranging_env.lost_cnt = 0;

                    ranging_user_pkt.seq_num = ranging_recv_pkt->seq_num;
                    uwb_app_config.ranging_stage = (RANGING_RESPONSE);

                    // process rx-poll packet
                    uwb_pkt_rx_done_ind(ind, RANGING_POLL);
                }
                else
                {
                    // process rx-poll packet
                    uwb_pkt_rx_done_ind(ind, RANGING_POLL);

                    if (ranging_env.enable)
                    {
                        power_on_radio(0, 1);
                        // receive sync again
                        mac_rx(EVT_MODE_MAC_PHY_ASAP, 0, ranging_env.ranging_period);
                        mac_start();
                        return stage;
                    }
                }
            }
            break;

            case RANGING_POLL:
            {
                if ((ind->err_code == UWB_RX_OK) && (ranging_recv_pkt->session_id == ranging_user_pkt.session_id) &&
                    (ranging_recv_pkt->msg_type == RANGING_POLL))
                {
                    // update anchor point
                    ranging_env.anchor_point = ind->timestamp - phy_shr_duration();
                    ranging_env.rx_poll_time = ranging_rx_time(ind);
                    ranging_env.lost_cnt = 0;

                    ranging_user_pkt.seq_num = ranging_recv_pkt->seq_num;

                    uwb_app_config.ranging_stage = (RANGING_RESPONSE);
                }
                else
                {
                    ranging_env.lost_cnt++;
                    if (ranging_env.lost_cnt > 3)
                    {
                        // receive sync again
                        uwb_app_config.ranging_stage = RANGING_SYNC;
                    }
                }

                // process rx-poll packet
                uwb_pkt_rx_done_ind(ind, RANGING_POLL);
            }
            break;

            case RANGING_RESPONSE:
            {
                uwb_app_config.ranging_stage = RANGING_POLL;

                // process tx-response packet
                uwb_pkt_tx_done_ind(ind, stage);
            }
            break;

            case RANGING_MRM:
                break;
            default:
                break;
        }
    }

    stage = uwb_app_config.ranging_stage;

    if (ranging_env.enable)
    {
        if (RANGING_POLL == stage)
        {
            phy_timer_target_set(ranging_env.anchor_point + ranging_env.ranging_period - UWB_PERIOD_PREFETCH_TIME, ranging_timer_callback);
            power_mode_request(POWER_UNIT_APP, POWER_MODE_POWER_DOWN);
        }
        else
        {
            phy_timer_target_set(ranging_env.anchor_point + ranging_env.slot_interval - UWB_PERIOD_PREFETCH_TIME, ranging_timer_callback);
            power_mode_request(POWER_UNIT_APP, POWER_MODE_SLEEP);
        }
    }
    else
    {
        power_mode_clear(POWER_UNIT_APP);
        LOG_INFO(TRACE_MODULE_APP, "Ranging stop\r\n");
    }

    // board_led_off(BOARD_LED_2);

    return stage;
}
#if defined(__ARMCC_VERSION) && (__ARMCC_VERSION >= 6010050)
#pragma clang diagnostic pop
#endif

static void uwb_pkt_tx_done_ind(const struct MAC_HW_REPORT_T *tx, enum RANGING_STAGE_T stage)
{
    struct UWB_PKT_TX_DONE_IND_T *ind;
    if ((ind = WsfMsgAlloc(sizeof(struct UWB_PKT_TX_DONE_IND_T) + tx->pkt_len)) != NULL)
    {
        ind->hdr.event = UWB_PKT_TX_DONE_MSG;
        ind->ranging_stage = (uint8_t)stage;
        ind->status = tx->err_code;

        ind->tx_len = tx->pkt_len;

        if ((ind->tx_len) && (tx->pkt_data != NULL))
        {
            memcpy(ind->tx_data, tx->pkt_data, tx->pkt_len);
        }

        // Send the message
        WsfMsgSend(ranging_cb.handle_id, ind);
    }
    else
    {
        LOG_WARNING(TRACE_MODULE_UWB, "memory is not enough for UWB_PKT_TX_DONE_IND_T\r\n");
    }
}

#if defined(__ARMCC_VERSION) && (__ARMCC_VERSION >= 6010050)
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wcast-qual"
#endif
static void uwb_pkt_rx_done_ind(const struct MAC_HW_REPORT_T *rx, enum RANGING_STAGE_T stage)
{
    // send an indication to application
    struct UWB_PKT_RX_DONE_IND_T *ind = WsfMsgAlloc(sizeof(struct UWB_PKT_RX_DONE_IND_T) + rx->pkt_len);

    if (ind != NULL)
    {
        ind->hdr.event = UWB_PKT_RX_DONE_MSG;
        ind->ranging_stage = (uint8_t)stage;
        ind->status = rx->err_code;
        ind->rssi = rx->rssi;
        ind->snr = rx->snr;

        ind->rx_len = rx->pkt_len;

        if ((rx->pkt_len) && (rx->pkt_data != NULL))
        {
            memcpy(ind->rx_data, rx->pkt_data, rx->pkt_len);
        }

        // Send the message
        WsfMsgSend(ranging_cb.handle_id, ind);
    }
    else
    {
        LOG_WARNING(TRACE_MODULE_UWB, "memory is not enough for UWB_PKT_RX_DONE_IND_T %d\r\n", rx->pkt_len);
    }
}
#if defined(__ARMCC_VERSION) && (__ARMCC_VERSION >= 6010050)
#pragma clang diagnostic pop
#endif

static void ranging_tx_process(struct MAC_HW_REPORT_T *tx_report)
{
    ranging_fsm(tx_report);
}

static void ranging_rx_process(struct MAC_HW_REPORT_T *rx_report)
{
    ranging_fsm(rx_report);
}