/**
 *
 * \section COPYRIGHT
 *
 * Copyright 2013-2020 Software Radio Systems Limited
 *
 * By using this file, you agree to the terms and conditions set
 * forth in the LICENSE file which can be found at the top level of
 * the distribution.
 *
 */

#include "sched_test_common.h"
#include "srsenb/hdr/stack/mac/sched.h"
#include "srsenb/hdr/stack/mac/sched_helpers.h"
#include "srsenb/hdr/stack/upper/common_enb.h"
#include "srslte/mac/pdu.h"

#include "sched_common_test_suite.h"
#include "sched_ue_ded_test_suite.h"
#include "srslte/common/test_common.h"

using namespace srsenb;

/***************************
 *     Random Utils
 **************************/

std::default_random_engine rand_gen;

float ::srsenb::randf()
{
  static std::uniform_real_distribution<float> unif_dist(0, 1.0);
  return unif_dist(rand_gen);
}

void ::srsenb::set_randseed(uint64_t seed)
{
  rand_gen = std::default_random_engine(seed);
}

std::default_random_engine& ::srsenb::get_rand_gen()
{
  return rand_gen;
}

/***********************
 *  User State Tester
 ***********************/

void user_state_sched_tester::new_tti(sched* sched_ptr, tti_point tti_rx)
{
  tic++;
  for (auto& u : sim_users) {
    auto        pending_feedback = u.second.get_pending_events(tti_rx, sched_ptr);
    auto&       ctxt             = u.second.get_ctxt();
    const auto& sim_cfg          = ue_sim_cfg_map.at(u.first);

    for (uint32_t enb_cc_idx = 0; enb_cc_idx < pending_feedback->cc_list.size(); ++enb_cc_idx) {
      auto& cc_feedback = pending_feedback->cc_list[enb_cc_idx];
      if (not cc_feedback.configured) {
        continue;
      }

      // ACK DL HARQs
      if (cc_feedback.dl_pid >= 0) {
        auto& h            = ctxt.cc_list[cc_feedback.ue_cc_idx].dl_harqs[cc_feedback.dl_pid];
        cc_feedback.dl_ack = randf() < sim_cfg.prob_dl_ack_mask[h.nof_retxs % sim_cfg.prob_dl_ack_mask.size()];
      }

      // ACK UL HARQs
      if (cc_feedback.ul_pid >= 0) {
        auto& h            = ctxt.cc_list[cc_feedback.ue_cc_idx].ul_harqs[cc_feedback.ul_pid];
        cc_feedback.ul_ack = randf() < sim_cfg.prob_ul_ack_mask[h.nof_retxs % sim_cfg.prob_ul_ack_mask.size()];
      }

      // DL CQI
      if (cc_feedback.dl_cqi >= 0) {
        cc_feedback.dl_cqi = std::uniform_int_distribution<uint32_t>{5, 24}(get_rand_gen());
      }

      // UL CQI
      if (cc_feedback.ul_cqi >= 0) {
        cc_feedback.ul_cqi = std::uniform_int_distribution<uint32_t>{5, 40}(get_rand_gen());
      }
    }
  }
}

int user_state_sched_tester::add_user(uint16_t rnti, uint32_t preamble_idx, const ue_ctxt_test_cfg& cfg_)
{
  CONDERROR(!srslte_prach_tti_opportunity_config_fdd(
                cell_params[cfg_.ue_cfg.supported_cc_list[0].enb_cc_idx].prach_config, tic.to_uint(), -1),
            "New user added in a non-PRACH TTI\n");
  TESTASSERT(sim_users.find_rnti(rnti) == nullptr);
  sim_users.add_user(rnti, generate_rach_ue_cfg(cfg_.ue_cfg), tic, preamble_idx);
  ue_sim_cfg_map[rnti] = cfg_;

  return SRSLTE_SUCCESS;
}

int user_state_sched_tester::user_reconf(uint16_t rnti, const srsenb::sched_interface::ue_cfg_t& ue_cfg)
{
  TESTASSERT(sim_users.find_rnti(rnti) != nullptr);
  sim_users.ue_recfg(rnti, ue_cfg);
  return SRSLTE_SUCCESS;
}

int user_state_sched_tester::bearer_cfg(uint16_t                                        rnti,
                                        uint32_t                                        lcid,
                                        const srsenb::sched_interface::ue_bearer_cfg_t& bearer_cfg)
{
  sim_users.bearer_cfg(rnti, lcid, bearer_cfg);
  return SRSLTE_SUCCESS;
}

bool user_state_sched_tester::is_drb_cfg(uint16_t rnti) const
{
  const ue_sim& ue = sim_users.at(rnti);
  for (uint32_t i = 2; i < ue.get_ctxt().ue_cfg.ue_bearers.size(); ++i) {
    if (ue.get_ctxt().ue_cfg.ue_bearers[i].direction != sched_interface::ue_bearer_cfg_t::IDLE) {
      return true;
    }
  }
  return false;
}

void user_state_sched_tester::rem_user(uint16_t rnti)
{
  sim_users.rem_user(rnti);
  ue_sim_cfg_map.erase(rnti);
}

int user_state_sched_tester::test_all(const sf_output_res_t& sf_out)
{
  // Perform UE-dedicated sched result tests
  sim_enb_ctxt_t enb_ctxt;
  enb_ctxt.cell_params = &cell_params;
  enb_ctxt.ue_db       = sim_users.get_ues_ctxt();
  TESTASSERT(test_all_ues(enb_ctxt, sf_out) == SRSLTE_SUCCESS);

  // Update Simulated UEs state
  sim_users.update(sf_out);

  return SRSLTE_SUCCESS;
}

/***********************
 *  Sim Stats Storage
 **********************/

void sched_result_stats::process_results(tti_point                                           tti_rx,
                                         const std::vector<sched_interface::dl_sched_res_t>& dl_result,
                                         const std::vector<sched_interface::ul_sched_res_t>& ul_result)
{
  for (uint32_t ccidx = 0; ccidx < dl_result.size(); ++ccidx) {
    for (uint32_t i = 0; i < dl_result[ccidx].nof_data_elems; ++i) {
      user_stats* user = get_user(dl_result[ccidx].data[i].dci.rnti);
      user->tot_dl_sched_data[ccidx] += dl_result[ccidx].data[i].tbs[0];
      user->tot_dl_sched_data[ccidx] += dl_result[ccidx].data[i].tbs[1];
    }
    for (uint32_t i = 0; i < ul_result[ccidx].nof_dci_elems; ++i) {
      user_stats* user = get_user(ul_result[ccidx].pusch[i].dci.rnti);
      user->tot_ul_sched_data[ccidx] += ul_result[ccidx].pusch[i].tbs;
    }
  }
}

sched_result_stats::user_stats* sched_result_stats::get_user(uint16_t rnti)
{
  if (users.count(rnti) != 0) {
    return &users[rnti];
  }
  users[rnti].rnti = rnti;
  users[rnti].tot_dl_sched_data.resize(cell_params.size(), 0);
  users[rnti].tot_ul_sched_data.resize(cell_params.size(), 0);
  return &users[rnti];
}

/***********************
 * Common Sched Tester
 **********************/

const sched::ue_cfg_t* common_sched_tester::get_current_ue_cfg(uint16_t rnti) const
{
  return ue_tester->get_user_cfg(rnti);
}

int common_sched_tester::sim_cfg(sim_sched_args args)
{
  sim_args0 = std::move(args);

  sched::cell_cfg(sim_args0.cell_cfg); // call parent cfg
  sched::set_sched_cfg(&sim_args0.sched_args);

  ue_tester.reset(new user_state_sched_tester{sim_args0.cell_cfg});
  sched_stats.reset(new sched_result_stats{sim_args0.cell_cfg});

  tester_log = sim_args0.sim_log;

  return SRSLTE_SUCCESS;
}

int common_sched_tester::add_user(uint16_t rnti, const ue_ctxt_test_cfg& ue_cfg_)
{
  CONDERROR(ue_cfg(rnti, generate_rach_ue_cfg(ue_cfg_.ue_cfg)) != SRSLTE_SUCCESS,
            "Configuring new user rnti=0x%x to sched\n",
            rnti);
  //        CONDERROR(!srslte_prach_tti_opportunity_config_fdd(
  //            sched_cell_params[CARRIER_IDX].cfg.prach_config, tti_info.tti_params.tti_rx, -1),
  //                  "New user added in a non-PRACH TTI\n");

  dl_sched_rar_info_t rar_info = {};
  rar_info.prach_tti           = tti_rx.to_uint();
  rar_info.temp_crnti          = rnti;
  rar_info.msg3_size           = 7;
  rar_info.preamble_idx        = tti_info.nof_prachs++;
  uint32_t pcell_idx           = ue_cfg_.ue_cfg.supported_cc_list[0].enb_cc_idx;
  dl_rach_info(pcell_idx, rar_info);

  ue_tester->add_user(rnti, rar_info.preamble_idx, ue_cfg_);

  tester_log->info("Adding user rnti=0x%x\n", rnti);
  return SRSLTE_SUCCESS;
}

int common_sched_tester::reconf_user(uint16_t rnti, const sched_interface::ue_cfg_t& ue_cfg_)
{
  CONDERROR(not ue_tester->user_exists(rnti), "User must already exist to be configured\n");
  CONDERROR(ue_cfg(rnti, ue_cfg_) != SRSLTE_SUCCESS, "Configuring new user rnti=0x%x to sched\n", rnti);
  ue_tester->user_reconf(rnti, ue_cfg_);
  return SRSLTE_SUCCESS;
}

void common_sched_tester::rem_user(uint16_t rnti)
{
  tester_log->info("Removing user rnti=0x%x\n", rnti);
  sched::ue_rem(rnti);
  ue_tester->rem_user(rnti);
}

void common_sched_tester::new_test_tti()
{
  if (not tti_rx.is_valid()) {
    tti_rx = srslte::tti_point{sim_args0.start_tti};
  } else {
    tti_rx++;
  }

  tti_info.nof_prachs = 0;
  tti_info.dl_sched_result.clear();
  tti_info.ul_sched_result.clear();
  tti_info.dl_sched_result.resize(sched_cell_params.size());
  tti_info.ul_sched_result.resize(sched_cell_params.size());

  tester_log->step(tti_rx.to_uint());
}

int common_sched_tester::process_results()
{
  // Perform common eNB result tests
  sf_output_res_t sf_out{sched_cell_params, tti_rx, tti_info.ul_sched_result, tti_info.dl_sched_result};
  TESTASSERT(test_all_common(sf_out) == SRSLTE_SUCCESS);

  TESTASSERT(ue_tester->test_all(sf_out) == SRSLTE_SUCCESS);

  sched_stats->process_results(tti_rx, tti_info.dl_sched_result, tti_info.ul_sched_result);

  return SRSLTE_SUCCESS;
}

int common_sched_tester::process_tti_events(const tti_ev& tti_ev)
{
  for (const tti_ev::user_cfg_ev& ue_ev : tti_ev.user_updates) {
    // There is a new configuration
    if (ue_ev.ue_sim_cfg != nullptr) {
      if (not ue_tester->user_exists(ue_ev.rnti)) {
        // new user
        TESTASSERT(add_user(ue_ev.rnti, *ue_ev.ue_sim_cfg) == SRSLTE_SUCCESS);
      } else {
        // reconfiguration
        TESTASSERT(reconf_user(ue_ev.rnti, ue_ev.ue_sim_cfg->ue_cfg) == SRSLTE_SUCCESS);
      }
    }

    // There is a user to remove
    if (ue_ev.rem_user) {
      rem_user(ue_ev.rnti);
    }

    // configure bearers
    if (ue_ev.bearer_cfg != nullptr) {
      CONDERROR(not ue_tester->user_exists(ue_ev.rnti), "User rnti=0x%x does not exist\n", ue_ev.rnti);
      // TODO: Instantiate more bearers
      bearer_ue_cfg(ue_ev.rnti, 0, ue_ev.bearer_cfg.get());
    }

    const ue_sim* user = ue_tester->get_user_ctxt(ue_ev.rnti);

    if (user != nullptr) {
      const auto& ue_sim_ctxt = user->get_ctxt();
      if (not ue_sim_ctxt.msg4_tti_rx.is_valid() and ue_sim_ctxt.msg3_tti_rx.is_valid() and
          to_tx_ul(ue_sim_ctxt.msg3_tti_rx) <= tti_rx) {
        // Msg3 has been received but Msg4 has not been yet transmitted
        // Setup default UE config
        reconf_user(ue_sim_ctxt.rnti, generate_setup_ue_cfg(sim_args0.default_ue_sim_cfg.ue_cfg));

        // Schedule RRC Setup and ConRes CE
        uint32_t pending_dl_new_data = ue_db[ue_ev.rnti].get_pending_dl_rlc_data();
        if (pending_dl_new_data == 0) {
          uint32_t lcid = RB_ID_SRB0; // Use SRB0 to schedule Msg4
          dl_rlc_buffer_state(ue_ev.rnti, lcid, 50, 0);
          dl_mac_buffer_state(ue_ev.rnti, (uint32_t)srslte::dl_sch_lcid::CON_RES_ID);
        } else {
          // Let SRB0 Msg4 get fully transmitted
        }
      }
    }

    // push UL SRs and DL packets
    if (ue_ev.buffer_ev != nullptr) {
      CONDERROR(user == nullptr, "TESTER ERROR: Trying to schedule data for user that does not exist\n");
      const auto& ue_sim_ctxt = user->get_ctxt();
      if (ue_ev.buffer_ev->dl_data > 0 and ue_sim_ctxt.msg4_tti_rx.is_valid()) {
        // If Msg4 has already been tx and there DL data to transmit
        uint32_t lcid                = RB_ID_DRB1;
        uint32_t pending_dl_new_data = ue_db[ue_ev.rnti].get_pending_dl_rlc_data();
        if (ue_tester->is_drb_cfg(ue_ev.rnti) or pending_dl_new_data == 0) {
          // If RRCSetup finished
          if (not ue_tester->is_drb_cfg(ue_ev.rnti)) {
            reconf_user(ue_sim_ctxt.rnti, sim_args0.default_ue_sim_cfg.ue_cfg);
            // setup lcid==drb1 bearer
            sched::ue_bearer_cfg_t cfg = {};
            cfg.direction              = ue_bearer_cfg_t::BOTH;
            cfg.group                  = 1;
            ue_tester->bearer_cfg(ue_ev.rnti, lcid, cfg);
            bearer_ue_cfg(ue_ev.rnti, lcid, &cfg);
          }
          // DRB is set. Update DL buffer
          uint32_t tot_dl_data = pending_dl_new_data + ue_ev.buffer_ev->dl_data; // TODO: derive pending based on rx
          dl_rlc_buffer_state(ue_ev.rnti, lcid, tot_dl_data, 0);                 // TODO: Check retx_queue
        } else {
          // Let SRB0 get emptied
        }
      }

      if (ue_ev.buffer_ev->sr_data > 0 and ue_tester->is_drb_cfg(ue_ev.rnti)) {
        uint32_t tot_ul_data =
            ue_db[ue_ev.rnti].get_pending_ul_new_data(to_tx_ul(tti_rx), -1) + ue_ev.buffer_ev->sr_data;
        uint32_t lcg = 1;
        ul_bsr(ue_ev.rnti, lcg, tot_ul_data);
      }
    }
  }
  return SRSLTE_SUCCESS;
}

int common_sched_tester::run_tti(const tti_ev& tti_events)
{
  new_test_tti();
  tester_log->info("---- tti=%u | nof_ues=%zd ----\n", tti_rx.to_uint(), ue_db.size());

  ue_tester->new_tti(this, tti_rx);
  process_tti_events(tti_events);
  before_sched();

  // Call scheduler for all carriers
  tti_info.dl_sched_result.resize(sched_cell_params.size());
  for (uint32_t i = 0; i < sched_cell_params.size(); ++i) {
    dl_sched(to_tx_dl(tti_rx).to_uint(), i, tti_info.dl_sched_result[i]);
  }
  tti_info.ul_sched_result.resize(sched_cell_params.size());
  for (uint32_t i = 0; i < sched_cell_params.size(); ++i) {
    ul_sched(to_tx_ul(tti_rx).to_uint(), i, tti_info.ul_sched_result[i]);
  }

  process_results();
  tti_count++;
  return SRSLTE_SUCCESS;
}

int common_sched_tester::test_next_ttis(const std::vector<tti_ev>& tti_events)
{
  while (tti_count < tti_events.size()) {
    TESTASSERT(run_tti(tti_events[tti_count]) == SRSLTE_SUCCESS);
  }
  return SRSLTE_SUCCESS;
}