/*
 * 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.
 */

#ifndef LIB_FIRA_H_
#define LIB_FIRA_H_
#include "mk_common.h"
#include "mk_uwb.h"
#include "uwb_api.h"
#include "uwb_fira_message.h"
#include "wsf_msg.h"

/**
 * @addtogroup MK8000_FIRA_LIB
 * @{
 */

/// Static STS session key
#define SESSION_KEY_STATIC_STS "StaticTSStaticTS"

/// Dynamic STS session key for test
#define SESSION_KEY_DYNAMIC_STS_TEST "DynamicSTSNoRot0"

// the index to store slot of message in responder_slot_idx[]
#define SLOT_RESPONSE 0
#define SLOT_FINAL 1
#define SLOT_REPORT 2
#define SLOT_RESULT 3

/** STS config definition */
enum FIRA_STS_CFG_T
{
    FIRA_STATIC_STS = 0, /*!< static sts config */
    FIRA_DYNAMIC_STS = 1 /*!< dynamic sts config */
};

/** FiRa message type */
enum FIRA_MSG_TYPE_T
{
    FIRA_MSG_RANG_INIT = 0x0,      /*!< Ranging Initiation Message */
    FIRA_MSG_RANG_RESP = 0x1,      /*!< Ranging Response Message */
    FIRA_MSG_RANG_FINAL = 0x2,     /*!< Ranging Final Message */
    FIRA_MSG_CONTROL = 0x3,        /*!< Control Message */
    FIRA_MSG_MEASURE_RPT = 0x4,    /*!< Measurement Report Message */
    FIRA_MSG_RESULT_RPT = 0x5,     /*!< Ranging Result Report Message */
    FIRA_MSG_CONTROL_UPDATE = 0x6, /*!< Control Update Message*/
    FIRA_MSG_OWR = 0x7,            /*!< OWR Message*/
    FIRA_MSG_DM = 0x8,             /*!< Data Message*/
};

/** General key data structure */
struct FIRA_KEY_T
{
    union
    {
        uint8_t keyByte[16];
        uint32_t keyWord[4];
    } ukey;
};

/** Fira UWB session configuration, used for key derivation */
struct FIRA_SESSSION_CFG_T
{
    uint8_t rgMethod;
    uint8_t firaStsCfg;
    uint8_t multiMode;
    uint8_t chan;
    uint8_t slotDur[2]; /* NOTE: saved as big endian, uint: us (RSTU*416/499.2) */
    uint8_t fcsType;
    uint8_t spCfg;
    uint8_t preCidx;
    uint8_t sfdId;
    uint8_t dataRate;
    uint8_t preLen;
    uint8_t macCfg;
    uint8_t sessionId[4]; /* NOTE: saved as big endian */
};

/** FiRa key list */
struct FIRA_KEY_LIST_T
{
    struct FIRA_KEY_T sessKey;
    struct FIRA_KEY_T dataProKey;
    struct FIRA_KEY_T devAuthIv;
    struct FIRA_KEY_T devAuthKey;
    struct FIRA_KEY_T devPayKey;
    struct FIRA_KEY_T privacyKey;
    uint32_t phyStsIdxInit;
    uint8_t digest[16];     /* store the calculated digest result to avoid calculate again, need to update when session cfg changed */
    uint8_t staticStsIv[6]; /* transferred over UCI */
    uint8_t vendorId[2];    /* transferred over UCI */
};

/** Responder node structure */
struct RESPONDER_NODE_T
{
    uint8_t addr_is_valid; /* 0 is invalid, 1 is valid */
    uint8_t main_ant_id;
    uint16_t short_add;
};

/** FiRa ranging environment variable */
struct RANGING_ENV_T
{
    int64_t tx_poll_time;
    int64_t tx_response_time;
    int64_t tx_final_time;
    int64_t rx_poll_time;
    int64_t rx_response_time[RESPONDER_NUM_MAX];
    int64_t rx_final_time;
    uint8_t slot_bitmap_tx[BIT_MAP_SIZE_MAX];
    uint8_t slot_bitmap_rx[BIT_MAP_SIZE_MAX];

    uint32_t anchor_point;
    uint32_t ranging_period;
    uint32_t round_duration;
    uint32_t round_offset_in_block;
    uint32_t slot_interval;
    uint32_t phy_sts_index;
    uint32_t tof;

    uint32_t uwb_period_prefetch_time;
    uint32_t uwb_evt_prefetch_time;
    uint32_t uwb_rx_open_in_advance;
    uint32_t uwb_rx_window;
    uint32_t uwb_rx_open_in_advance_wakeup;
    uint32_t uwb_rx_window_wakeup;

    uint32_t random_num[10];
    int32_t freq_offset_filter;

    uint8_t lost_cnt;
    uint8_t enable;
    uint8_t slot_idx;
    uint8_t responder_num;
    struct RESPONDER_NODE_T responder_list[RESPONDER_NUM_MAX];
    uint8_t responder_slot_idx[4];
    uint32_t responder_response_flag;
    uint16_t result_flag; // responder_idx || status (1: update, 2: complete)
    uint8_t rcm_flag;
    uint8_t repeat;
    uint16_t slots_per_block;
    uint8_t count;
    uint8_t count_last;

    uint8_t is_hopping;
    uint8_t stride_length;
    uint16_t curr_round_index;
    uint16_t next_round_index;
    uint16_t nround_inblock;
#if FIRA_TEST_EN
    uint16_t num_of_measure;
#endif
    struct RANGE_DATA_T range_data;
};

/** UWB packet TX done indication */
struct UWB_PKT_TX_DONE_IND_T
{
    wsfMsgHdr_t hdr;
#if MCTT_TEST_EN
    uint32_t timestamp;
    uint32_t phy_shr_duration;
#endif
    uint8_t ranging_stage;
    uint8_t slot_idx;
    uint16_t status;
    uint16_t tx_len;
    uint8_t tx_data[];
};

/** UWB packet RX done indication */
struct UWB_PKT_RX_DONE_IND_T
{
    wsfMsgHdr_t hdr;
#if MCTT_TEST_EN
    uint32_t timestamp;
    uint32_t phy_header;
#endif
    uint8_t slot_idx;
    uint8_t ranging_stage;
    uint16_t status;
    int8_t rssi;
    int8_t snr;
    uint16_t rx_len;
    uint8_t rx_data[];
};

#ifdef __cplusplus
extern "C" {
#endif

extern struct FIRA_SESSSION_CFG_T session_config;
extern struct FIRA_KEY_LIST_T fira_key;

extern struct RANGING_ENV_T ranging_env;
extern uint8_t fira_uwb_tx_buf[PHY_PAYLOAD_LEN_MAX];

extern uint8_t *tx_msg;
extern uint16_t tx_msg_len;

/**
 * @brief Calculate timestamp difference.
 * @param[in] ts_a          Ranging timestamp a
 * @param[in] ts_b          Ranging timestamp b
 * @return ts_a - ts_b
 */
int64_t ranging_timestamp_diff(int64_t ts_a, int64_t ts_b);

/**
 * @brief Calculate ranging Tround.
 * @param[in] role          Device role @enum DEV_ROLE_T
 * @param[in] dev_idx       Device index in one-to-many case
 * @return Tround value, unit: 15.6ps
 */
int64_t ranging_tround(enum DEV_ROLE_T role, uint8_t dev_idx);

/**
 * @brief Calculate ranging Treply.
 * @param[in] role          Device role @enum DEV_ROLE_T
 * @param[in] dev_idx       Device index in one-to-many case
 * @return Treply value, unit: 15.6ps
 */
int64_t ranging_treply(enum DEV_ROLE_T role, uint8_t dev_idx);

/**
 * @brief Set local device short address.
 * @param[in] short_addr          Device short address
 */
void ranging_local_addr_set(uint16_t short_addr);

/**
 * @brief Set peer device short address.
 * @param[in] short_addr          Device short address
 */
void ranging_peer_addr_set(uint16_t short_addr);

/**
 * @brief Get initiator short address.
 * @return short address of initiator
 */
uint16_t ranging_initiator_addr_get(void);

/**
 * @brief Add responder device into responder device list for one-to-many case.
 * @param[in] addr              Device short address
 * @return 1: successful, 0: fail
 */
uint8_t ranging_responder_addr_add(uint16_t addr);

/**
 * @brief Clear responder list.
 */
void ranging_responder_list_clr(void);

/**
 * @brief Dynamic Update of Multicast List for Controlees
 * @param[in] action            0x00: Update the multicast list by adding requested controlee short address
 *                              0x01: Delete the requested short address from multicast list
 * @param[in] addr              Device short address
 * @return 0x01: successful, 0x07: short address is not found, 0x08: short address is already present
 */
uint8_t ranging_dynamic_update_respoder_list(uint8_t action, uint16_t addr);

/**
 * @brief Get responder short address from responder device list.
 * @param[in] idx               Responder device index
 * @return short address of responder
 */
uint16_t ranging_responder_addr_get(uint8_t idx);

/**
 * @brief Get responder devices number.
 * @return number of responder
 */
uint8_t ranging_responder_num_get(void);

/**
 * @brief Update ranging slot index, based on this informantion to generate STS or encryted packet every slot.
 * @param[in] slot_idx              Slot index
 */
void ranging_update_slot_index(uint32_t slot_idx);

/**
 * @brief Get the current block index.
 * @return Block index value
 */
uint32_t ranging_get_block_index(void);

/**
 * @brief Get the current round index.
 * @return Round index value
 */
uint16_t ranging_get_curr_round_index(void);

/**
 * @brief Get the next round index.
 * @param[in] session_id              Fira Ranging session Id
 * @param[in] block_index             Block index
 * @return Next block index value
 */
uint16_t ranging_get_next_round_index(uint32_t session_id, uint32_t block_index);

/**
 * @brief Callback function for PHY timer, to program MAC TX or RX event in advance.
 * @param[in] dev                   Pointer of timer device
 * @param[in] time                  Last time setting
 */
void session_timer_callback(void *dev, uint32_t time);

/**
 * @brief Session FSM, to process MAC TX or RX done event in MAC interrupt handler.
 * @param[in] ind                   Pointer of MAC HW report
 * @return next ranging stage
 */
enum RANGING_STAGE_T session_fsm(const struct MAC_HW_REPORT_T *ind);

/**
 * @brief Indicate UWB packet TX done.
 * @param[in] tx                Pointer of MAC HW report for TX
 * @param[in] stage             Stage of ranging FSM @enum RANGING_STAGE_T
 * @param[in] slot_idx          Slot index
 */
void uwb_pkt_tx_done_ind(const struct MAC_HW_REPORT_T *tx, enum RANGING_STAGE_T stage, uint8_t slot_idx);

/**
 * @brief Indicate UWB packet RX done.
 * @param[in] rx                Pointer of MAC HW report for RX
 * @param[in] stage             Stage of ranging FSM @enum RANGING_STAGE_T
 * @param[in] slot_idx          Slot index
 * @return RX status
 */
uint16_t uwb_pkt_rx_done_ind(const struct MAC_HW_REPORT_T *rx, enum RANGING_STAGE_T stage, uint8_t slot_idx);

/**
 * @brief Generate FiRa ranging session configuration.
 * @param[out] pFiraSesscfg          Pointer of FiRa session configuration
 * @param[in] pFiraUwbcfg            Pointer of FiRa UWB configuration
 */
void fira_session_cfg_generate(struct FIRA_SESSSION_CFG_T *pFiraSesscfg, struct UWB_APP_CFG_T *pFiraUwbcfg);

/**
 * @brief Generate FiRa ranging keys.
 * @param[out] pFiraKeyList          Pointer of FiRa ranging key list
 * @param[in] pFiraSesscfg           Pointer of FiRa session configuration
 * @param[in] stsCfg                 Static STS or dynamic STS @enum FIRA_STS_CFG_T
 * @param[in] pSessionKey            Pointer of ranging session key
 */
void fira_key_list_generate(struct FIRA_KEY_LIST_T *pFiraKeyList, struct FIRA_SESSSION_CFG_T *pFiraSesscfg, enum FIRA_STS_CFG_T stsCfg, uint8_t *pSessionKey);

/**
 * @brief Update STS key configuration
 * @param[in] firaStsCfg             Static STS or dynamic STS @enum FIRA_STS_CFG_T
 * @param[in] pFiraKeyList           Pointer of FiRa ranging key list
 * @param[in] write_en               Write updated STS configuration to registers directly, for dynamic STS key rotation
 */
void fira_sts_keys_update(enum FIRA_STS_CFG_T firaStsCfg, struct FIRA_KEY_LIST_T *pFiraKeyList, uint8_t write_en);

/**
 * @brief Generate FiRa keys.
 * @note The fira key is generated by session_config and fira_key global variables.
 */
void fira_keys_generate(void);

/**
 * @brief Rotation FiRa key.
 * @param[in] pFiraKeyList           Pointer of FiRa ranging key list
 * @param[in] cryptoStsIdx           PHY STS index, +1 every slot
 */
void fira_key_rotation(struct FIRA_KEY_LIST_T *pFiraKeyList, uint32_t cryptoStsIdx);

/**
 * @brief Encrypt vendor IE data.
 * @param[in] input             Plaintext of vendor IE data
 * @param[out] output           Cyphertext of vendor IE data
 */
void fira_vendor_ie_encrypt(uint8_t *input, uint8_t *output);

/**
 * @brief Decrypt vendor IE data.
 * @param[in] input             Cyphertext of vendor IE data
 * @param[out] output           Plaintext of vendor IE data
 */
void fira_vendor_ie_decrypt(uint8_t *input, uint8_t *output);

/**
 * @brief Decrypt vendor IE data to get phyStsIndex for Dynamic STS case.
 * @param[in] rx_buf            Pointer of RX buffer
 */
uint32_t fira_vendor_oui_process(uint8_t *rx_buf);

/**
 * @brief Encrypt payload data.
 * @param[in] input             Plaintext of payload data
 * @param[out] output           Cyphertext of payload data
 * @param[in] input_len         Input data total length
 * @param[in] associate_len     Authentication data length
 * @param[in] sec_lvl           Security level
 * @param[in] slot_idx          Slot index
 */
uint16_t RAM_FUNC fira_packet_encrypt(uint8_t *input, uint8_t *output, uint16_t input_len, uint16_t associate_len, uint8_t sec_lvl, uint32_t slot_idx);

/**
 * @brief Decrypt payload data.
 * @param[in] input             Cyphertext of payload data
 * @param[out] output           Plaintext of payload data
 * @param[in] input_len         Input data total length
 * @param[in] slot_idx          Slot index
 */
uint16_t fira_packet_decrypt(uint8_t *input, uint8_t *output, uint16_t input_len, uint32_t slot_idx);

#ifdef __cplusplus
}
#endif

/**
 * @}
 */

#endif /* LIB_FIRA_H_ */