/*
 * SPDX-License-Identifier: LicenseRef-CSSL-1.0
 */

#include "ran_func_kpm_subs.h"

#include <search.h>

e2_node_level_stats_t cp_node_level_stats(const e2_node_level_stats_t *src)
{
  e2_node_level_stats_t dst = {
    .mac_stats.dl.total_prb_aggregate = src->mac_stats.dl.total_prb_aggregate,
    .mac_stats.dl.used_prb_aggregate = src->mac_stats.dl.used_prb_aggregate,
    .mac_stats.ul.total_prb_aggregate = src->mac_stats.ul.total_prb_aggregate,
    .mac_stats.ul.used_prb_aggregate = src->mac_stats.ul.used_prb_aggregate,
  };

  return dst;
}

/* measurements that need to store values from previous reporting period have a limitation
   when it comes to multiple subscriptions to the same UEs; ric_req_id is unique per subscription */
typedef struct uldlcounter {
  uint32_t dl;
  uint32_t ul;
} uldlcounter_t;

static uldlcounter_t last_pdcp_sdu_total_bytes[MAX_MOBILES_PER_GNB] = {0};

static nr_pdcp_statistics_t get_pdcp_stats_per_drb(const uint32_t rrc_ue_id, const int rb_id)
{
  nr_pdcp_statistics_t pdcp = {0};
  const int srb_flag = 0;

  // Get PDCP stats for specific DRB
  const bool rc = nr_pdcp_get_statistics(rrc_ue_id, srb_flag, rb_id, &pdcp);
  assert(rc == true && "Cannot get PDCP stats\n");

  return pdcp;
}

/* 3GPP TS 28.522 - section 5.1.2.1.1.1
  note: this measurement is calculated as per spec */
static meas_record_lst_t fill_DRB_PdcpSduVolumeDL(__attribute__((unused)) const label_info_lst_t label,
                                                  __attribute__((unused)) uint32_t gran_period_ms,
                                                  cudu_ue_info_pair_t ue_info,
                                                  const size_t ue_idx,
                                                  __attribute__((unused))e2_node_level_stats_t* node_stats)
{
  meas_record_lst_t meas_record = {0};

  // Get PDCP stats per DRB
  const int rb_id = 1;  // at the moment, only 1 DRB is supported
  nr_pdcp_statistics_t pdcp = get_pdcp_stats_per_drb(ue_info.rrc_ue_id, rb_id);

  meas_record.value = INTEGER_MEAS_VALUE;

  // Get DL data volume delivered to PDCP layer
  meas_record.int_val = (pdcp.rxsdu_bytes - last_pdcp_sdu_total_bytes[ue_idx].dl)*8/1000000;   // [Mb]
  last_pdcp_sdu_total_bytes[ue_idx].dl = pdcp.rxsdu_bytes;

  return meas_record;
}

/* 3GPP TS 28.522 - section 5.1.2.1.2.1
  note: this measurement is calculated as per spec */
static meas_record_lst_t fill_DRB_PdcpSduVolumeUL(__attribute__((unused)) const label_info_lst_t label,
                                                  __attribute__((unused)) uint32_t gran_period_ms,
                                                  cudu_ue_info_pair_t ue_info,
                                                  const size_t ue_idx,
                                                  __attribute__((unused))e2_node_level_stats_t* node_stats)
{
  meas_record_lst_t meas_record = {0};

  // Get PDCP stats per DRB
  const int rb_id = 1;  // at the moment, only 1 DRB is supported
  nr_pdcp_statistics_t pdcp = get_pdcp_stats_per_drb(ue_info.rrc_ue_id, rb_id);

  meas_record.value = INTEGER_MEAS_VALUE;

  // Get UL data volume delivered from PDCP layer
  meas_record.int_val = (pdcp.txsdu_bytes - last_pdcp_sdu_total_bytes[ue_idx].ul)*8/1000000;   // [Mb]
  last_pdcp_sdu_total_bytes[ue_idx].ul = pdcp.txsdu_bytes;

  return meas_record;
}

#if defined (NGRAN_GNB_DU)
typedef struct ue_uldlcounter {
  const NR_UE_info_t *ue;
  rnti_t rnti;
  uldlcounter_t value;
  bool dl_valid;
  bool ul_valid;
} ue_uldlcounter_t;

static ue_uldlcounter_t last_rlc_pdu_total_bytes[MAX_MOBILES_PER_GNB] = {0};
static ue_uldlcounter_t last_total_prbs[MAX_MOBILES_PER_GNB] = {0};

/* The connected UE list is compacted when a UE leaves, so its iteration index
 * is not a stable identity. Keep baselines by the MAC-assigned UID instead and
 * reset them when that UID is reused for another UE. */
static ue_uldlcounter_t *get_ue_counter(ue_uldlcounter_t counters[MAX_MOBILES_PER_GNB],
                                       const cudu_ue_info_pair_t ue_info)
{
  AssertFatal(ue_info.ue != NULL, "KPM UE measurement requires MAC UE information\n");
  const uid_t uid = ue_info.ue->uid;
  AssertFatal(uid < MAX_MOBILES_PER_GNB, "KPM UE UID %u exceeds maximum %u\n", uid, MAX_MOBILES_PER_GNB);

  ue_uldlcounter_t *counter = &counters[uid];
  if (counter->ue != ue_info.ue || counter->rnti != ue_info.ue->rnti)
    *counter = (ue_uldlcounter_t){.ue = ue_info.ue, .rnti = ue_info.ue->rnti};
  return counter;
}

/* A new/reused UE counter and a reset/wrap have no valid preceding sample.
 * Establish a fresh baseline and report zero for that single period. */
static uint32_t counter_delta(uint32_t current, uint32_t *previous, bool *valid)
{
  const uint32_t delta = *valid && current >= *previous ? current - *previous : 0;
  *previous = current;
  *valid = true;
  return delta;
}

static nr_rlc_statistics_t get_rlc_stats_per_drb(const rnti_t rnti, const int rb_id)
{
  nr_rlc_statistics_t rlc = {0};
  const int srb_flag = 0;

  // Get RLC stats for specific DRB
  const bool rc = nr_rlc_get_statistics(rnti, srb_flag, rb_id, &rlc);
  assert(rc == true && "Cannot get RLC stats\n");

  // Activate average sojourn time at the RLC buffer for specific DRB
  nr_rlc_activate_avg_time_to_tx(rnti, rb_id+3, 1);
  
  return rlc;  
}

/* 3GPP TS 28.522 - section 5.1.3.3.3
  note: by default this measurement is calculated for previous 100ms (openair2/LAYER2/nr_rlc/nr_rlc_entity.c:118, 173, 213); please, update according to your needs */
static meas_record_lst_t fill_DRB_RlcSduDelayDl(__attribute__((unused)) const label_info_lst_t label,
                                                __attribute__((unused)) uint32_t gran_period_ms,
                                                cudu_ue_info_pair_t ue_info,
                                                __attribute__((unused))const size_t ue_idx,
                                                __attribute__((unused))e2_node_level_stats_t* node_stats)
{
  meas_record_lst_t meas_record = {0};
  
  // Get RLC stats per DRB
  const int rb_id = 1;  // at the moment, only 1 DRB is supported
  nr_rlc_statistics_t rlc = get_rlc_stats_per_drb(ue_info.ue->rnti, rb_id);

  meas_record.value = REAL_MEAS_VALUE;

  // Report the windowed-average RLC SDU sojourn (mean per-packet wait over the
  // last 100 ms), not the instantaneous tx_list HOL snapshot. The snapshot
  // (txsdu_wt_us) reads 0 whenever the tx_list is momentarily drained at the
  // single sampling instant -- common with bursty DL even at high throughput --
  // which systematically under-reports delay (era/load-dependent). The averaging
  // accumulator is now activated once (reset guarded to off->on in
  // nr_rlc_activate_avg_time_to_tx), so it accumulates continuously instead of
  // being wiped on every read; it returns 0 only when no SDU was transmitted in
  // the window (genuine idle/starvation).
  meas_record.real_val = (double)rlc.txsdu_avg_time_to_tx;  // [us] avg RLC SDU sojourn

  return meas_record;
}

/* 3GPP TS 28.522 - section 5.1.1.3.1
  note: per spec, average UE throughput in DL (taken into consideration values from all UEs, and averaged)
        here calculated as: UE specific throughput in DL */
static meas_record_lst_t fill_DRB_UEThpDl(__attribute__((unused)) const label_info_lst_t label,
                                          uint32_t gran_period_ms,
                                          cudu_ue_info_pair_t ue_info,
                                          __attribute__((unused)) const size_t ue_idx,
                                          __attribute__((unused))e2_node_level_stats_t* node_stats)
{
  meas_record_lst_t meas_record = {0};
  
  const int rb_id = 1;  // at the moment, only 1 DRB is supported
  nr_rlc_statistics_t rlc = get_rlc_stats_per_drb(ue_info.ue->rnti, rb_id);
  meas_record.value = REAL_MEAS_VALUE;

  // Calculate DL Thp
  ue_uldlcounter_t *counter = get_ue_counter(last_rlc_pdu_total_bytes, ue_info);
  const uint32_t bytes = counter_delta(rlc.txpdu_bytes, &counter->value.dl, &counter->dl_valid);
  meas_record.real_val = gran_period_ms == 0 ? 0 : (double)bytes * 8 / gran_period_ms;  // [kbps]

  return meas_record;
}

/* 3GPP TS 28.522 - section 5.1.1.3.3
  note: per spec, average UE throughput in UL (taken into consideration values from all UEs, and averaged)
        here calculated as: UE specific throughput in UL */
static meas_record_lst_t fill_DRB_UEThpUl(__attribute__((unused)) const label_info_lst_t label,
                                          uint32_t gran_period_ms,
                                          cudu_ue_info_pair_t ue_info,
                                          __attribute__((unused)) const size_t ue_idx,
                                          __attribute__((unused))e2_node_level_stats_t* node_stats)
{
  meas_record_lst_t meas_record = {0};
  
  // Get RLC stats per DRB
  const int rb_id = 1;  // at the moment, only 1 DRB is supported
  nr_rlc_statistics_t rlc = get_rlc_stats_per_drb(ue_info.ue->rnti, rb_id);

  meas_record.value = REAL_MEAS_VALUE;

  // Calculate UL Thp
  ue_uldlcounter_t *counter = get_ue_counter(last_rlc_pdu_total_bytes, ue_info);
  const uint32_t bytes = counter_delta(rlc.rxpdu_bytes, &counter->value.ul, &counter->ul_valid);
  meas_record.real_val = gran_period_ms == 0 ? 0 : (double)bytes * 8 / gran_period_ms;  // [kbps]
  
  return meas_record;
}

/* 3GPP TS 28.522 - section 5.1.1.2.1
  note: this measurement is calculated as per spec */
static meas_record_lst_t fill_RRU_PrbTotDl(__attribute__((unused)) const label_info_lst_t label,
                                           __attribute__((unused)) uint32_t gran_period_ms,
                                           cudu_ue_info_pair_t ue_info,
                                           __attribute__((unused)) const size_t ue_idx,
                                           e2_node_level_stats_t* node_stats)
{
  meas_record_lst_t meas_record = {0};
  
  meas_record.value = INTEGER_MEAS_VALUE;

  // Get the number of DL PRBs
  ue_uldlcounter_t *counter = get_ue_counter(last_total_prbs, ue_info);
  const uint32_t ue_delta = counter_delta(ue_info.ue->mac_stats.dl.total_rbs, &counter->value.dl, &counter->dl_valid);
  const uint64_t current = node_stats[1].mac_stats.dl.total_prb_aggregate;
  const uint64_t previous = node_stats[0].mac_stats.dl.total_prb_aggregate;
  const uint64_t total_delta = current >= previous ? current - previous : 0;
  const uint64_t percentage = total_delta == 0 ? 0 : (uint64_t)ue_delta * 100 / total_delta;
  meas_record.int_val = percentage > 100 ? 100 : percentage;  // [%]

  return meas_record;
}

/* 3GPP TS 28.522 - section 5.1.1.2.2
  note: this measurement is calculated as per spec */
static meas_record_lst_t fill_RRU_PrbTotUl(__attribute__((unused)) const label_info_lst_t label,
                                           __attribute__((unused)) uint32_t gran_period_ms,
                                           cudu_ue_info_pair_t ue_info,
                                           __attribute__((unused)) const size_t ue_idx,
                                           e2_node_level_stats_t* node_stats)
{
  meas_record_lst_t meas_record = {0};

  meas_record.value = INTEGER_MEAS_VALUE;

  // Get the number of UL PRBs
  ue_uldlcounter_t *counter = get_ue_counter(last_total_prbs, ue_info);
  const uint32_t ue_delta = counter_delta(ue_info.ue->mac_stats.ul.total_rbs, &counter->value.ul, &counter->ul_valid);
  const uint64_t current = node_stats[1].mac_stats.ul.total_prb_aggregate;
  const uint64_t previous = node_stats[0].mac_stats.ul.total_prb_aggregate;
  const uint64_t total_delta = current >= previous ? current - previous : 0;
  const uint64_t percentage = total_delta == 0 ? 0 : (uint64_t)ue_delta * 100 / total_delta;
  meas_record.int_val = percentage > 100 ? 100 : percentage;  // [%]

  return meas_record;
}

/* 3GPP TS 28.552 - section 5.1.1.12.1*/
static meas_record_lst_t fill_CARR_PDSCHMCSDist(const label_info_lst_t label,
                                                __attribute__((unused)) uint32_t gran_period_ms,
                                                __attribute__((unused)) cudu_ue_info_pair_t ue_info,
                                                __attribute__((unused)) const size_t ue_idx,
                                                __attribute__((unused))e2_node_level_stats_t* node_stats)
{
  meas_record_lst_t meas_record = {.value = INTEGER_MEAS_VALUE};

  uint32_t bin_x = *label.distBinX, bin_y = *label.distBinY, bin_z = *label.distBinZ;

  NR_du_stats_t* du_stats = &RC.nrmac[0]->du_stats;

  if (bin_x <= 8 && bin_y <= 3 && bin_z <= 31) {
    meas_record.int_val = du_stats->pdsch_mcs_dist[bin_x - 1][bin_y - 1][bin_z];
  } else {
    printf("[E2 AGENT] Unknown binX %d, binY %d, binZ %d for \"CARR.PDSCHMCSDist\" measurement\n", bin_x, bin_y, bin_z);
    meas_record.int_val = 0;
  }

  return meas_record;
}
#endif

static kv_measure_t lst_measure[] = {
  {.key = "DRB.PdcpSduVolumeDL", .value = fill_DRB_PdcpSduVolumeDL }, 
  {.key = "DRB.PdcpSduVolumeUL", .value = fill_DRB_PdcpSduVolumeUL },
#if defined (NGRAN_GNB_DU)
  {.key = "DRB.RlcSduDelayDl", .value =  fill_DRB_RlcSduDelayDl }, 
  {.key = "DRB.UEThpDl", .value =  fill_DRB_UEThpDl }, 
  {.key = "DRB.UEThpUl", .value =  fill_DRB_UEThpUl }, 
  {.key = "RRU.PrbTotDl", .value =  fill_RRU_PrbTotDl }, 
  {.key = "RRU.PrbTotUl", .value =  fill_RRU_PrbTotUl }, 
  {.key = "CARR.PDSCHMCSDist", .value = fill_CARR_PDSCHMCSDist },
#endif
}; 

void init_kpm_subs_data(void)
{
  const size_t ht_len = sizeof(lst_measure) / sizeof(lst_measure[0]);
  hcreate(ht_len);

  ENTRY kv_pair;

  for (size_t i = 0; i < ht_len; i++) {
    kv_pair.key = lst_measure[i].key;
    kv_pair.data = &lst_measure[i];
    hsearch(kv_pair, ENTER);
  }
}

meas_record_lst_t get_kpm_meas_value(char* kpm_meas_name, const label_info_lst_t label, uint32_t gran_period_ms, cudu_ue_info_pair_t ue_info, const size_t ue_idx, e2_node_level_stats_t* node_stats)
{
  assert(kpm_meas_name != NULL);

  ENTRY search_entry = {.key = kpm_meas_name};
  ENTRY *found_entry = hsearch(search_entry, FIND);
  assert(found_entry != NULL && "Unsupported KPM measurement name");

  kv_measure_t *kv_found = (kv_measure_t *)found_entry->data;
  meas_record_lst_t meas_record = kv_found->value(label, gran_period_ms, ue_info, ue_idx, node_stats);

  return meas_record;
}