/*
 * Copyright (c) 2012 Centre Tecnologic de Telecomunicacions de Catalunya (CTTC)
 *
 * SPDX-License-Identifier: GPL-2.0-only
 *
 * Author: Nicola Baldo <nbaldo@cttc.es>
 */

#include "lte-rrc-protocol-ideal.h"

#include "lte-enb-net-device.h"
#include "lte-enb-rrc.h"
#include "lte-ue-net-device.h"
#include "lte-ue-rrc.h"

#include "ns3/fatal-error.h"
#include "ns3/log.h"
#include "ns3/node-list.h"
#include "ns3/node.h"
#include "ns3/nstime.h"
#include "ns3/simulator.h"

namespace ns3
{

NS_LOG_COMPONENT_DEFINE("LteRrcProtocolIdeal");

/**
 * @ingroup lte
 *
 */

/// RRC ideal message delay
static const Time RRC_IDEAL_MSG_DELAY;

NS_OBJECT_ENSURE_REGISTERED(LteUeRrcProtocolIdeal);

LteUeRrcProtocolIdeal::LteUeRrcProtocolIdeal()
    : m_ueRrcSapProvider(nullptr),
      m_enbRrcSapProvider(nullptr)
{
    m_ueRrcSapUser = new MemberLteUeRrcSapUser<LteUeRrcProtocolIdeal>(this);
}

LteUeRrcProtocolIdeal::~LteUeRrcProtocolIdeal()
{
}

void
LteUeRrcProtocolIdeal::DoDispose()
{
    NS_LOG_FUNCTION(this);
    delete m_ueRrcSapUser;
    m_rrc = nullptr;
}

TypeId
LteUeRrcProtocolIdeal::GetTypeId()
{
    static TypeId tid = TypeId("ns3::LteUeRrcProtocolIdeal")
                            .SetParent<Object>()
                            .SetGroupName("Lte")
                            .AddConstructor<LteUeRrcProtocolIdeal>();
    return tid;
}

void
LteUeRrcProtocolIdeal::SetLteUeRrcSapProvider(LteUeRrcSapProvider* p)
{
    m_ueRrcSapProvider = p;
}

LteUeRrcSapUser*
LteUeRrcProtocolIdeal::GetLteUeRrcSapUser()
{
    return m_ueRrcSapUser;
}

void
LteUeRrcProtocolIdeal::SetUeRrc(Ptr<LteUeRrc> rrc)
{
    m_rrc = rrc;
}

void
LteUeRrcProtocolIdeal::DoSetup(LteUeRrcSapUser::SetupParameters params)
{
    NS_LOG_FUNCTION(this);
    // We don't care about SRB0/SRB1 since we use ideal RRC messages.
}

void
LteUeRrcProtocolIdeal::DoSendRrcConnectionRequest(LteRrcSap::RrcConnectionRequest msg)
{
    // initialize the RNTI and get the EnbLteRrcSapProvider for the
    // eNB we are currently attached to
    m_rnti = m_rrc->GetRnti();
    SetEnbRrcSapProvider();

    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteEnbRrcSapProvider::RecvRrcConnectionRequest,
                        m_enbRrcSapProvider,
                        m_rnti,
                        msg);
}

void
LteUeRrcProtocolIdeal::DoSendRrcConnectionSetupCompleted(LteRrcSap::RrcConnectionSetupCompleted msg)
{
    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteEnbRrcSapProvider::RecvRrcConnectionSetupCompleted,
                        m_enbRrcSapProvider,
                        m_rnti,
                        msg);
}

void
LteUeRrcProtocolIdeal::DoSendRrcConnectionReconfigurationCompleted(
    LteRrcSap::RrcConnectionReconfigurationCompleted msg)
{
    // re-initialize the RNTI and get the EnbLteRrcSapProvider for the
    // eNB we are currently attached to
    m_rnti = m_rrc->GetRnti();
    SetEnbRrcSapProvider();

    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteEnbRrcSapProvider::RecvRrcConnectionReconfigurationCompleted,
                        m_enbRrcSapProvider,
                        m_rnti,
                        msg);
}

void
LteUeRrcProtocolIdeal::DoSendRrcConnectionReestablishmentRequest(
    LteRrcSap::RrcConnectionReestablishmentRequest msg)
{
    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteEnbRrcSapProvider::RecvRrcConnectionReestablishmentRequest,
                        m_enbRrcSapProvider,
                        m_rnti,
                        msg);
}

void
LteUeRrcProtocolIdeal::DoSendRrcConnectionReestablishmentComplete(
    LteRrcSap::RrcConnectionReestablishmentComplete msg)
{
    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteEnbRrcSapProvider::RecvRrcConnectionReestablishmentComplete,
                        m_enbRrcSapProvider,
                        m_rnti,
                        msg);
}

void
LteUeRrcProtocolIdeal::DoSendMeasurementReport(LteRrcSap::MeasurementReport msg)
{
    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteEnbRrcSapProvider::RecvMeasurementReport,
                        m_enbRrcSapProvider,
                        m_rnti,
                        msg);
}

void
LteUeRrcProtocolIdeal::DoSendIdealUeContextRemoveRequest(uint16_t rnti)
{
    NS_LOG_FUNCTION(this);

    uint16_t cellId = m_rrc->GetCellId();
    // re-initialize the RNTI and get the EnbLteRrcSapProvider for the
    // eNB we are currently attached to or attempting random access to
    // a target eNB
    m_rnti = m_rrc->GetRnti();

    NS_LOG_DEBUG("RNTI " << rnti << " sending UE context remove request to cell id " << cellId);
    NS_ABORT_MSG_IF(m_rnti != rnti, "RNTI mismatch");

    SetEnbRrcSapProvider(); // the provider has to be reset since the cell might have changed due to
                            // handover
    // ideally informing eNB
    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteEnbRrcSapProvider::RecvIdealUeContextRemoveRequest,
                        m_enbRrcSapProvider,
                        m_rnti);
}

void
LteUeRrcProtocolIdeal::SetEnbRrcSapProvider()
{
    NS_LOG_FUNCTION(this);

    uint16_t cellId = m_rrc->GetCellId();
    NS_LOG_DEBUG("RNTI " << m_rnti << " connected to cell " << cellId);

    // walk list of all nodes to get the peer eNB
    Ptr<LteEnbNetDevice> enbDev;
    auto listEnd = NodeList::End();
    bool found = false;
    for (auto i = NodeList::Begin(); (i != listEnd) && (!found); ++i)
    {
        Ptr<Node> node = *i;
        int nDevs = node->GetNDevices();
        for (int j = 0; (j < nDevs) && (!found); j++)
        {
            enbDev = node->GetDevice(j)->GetObject<LteEnbNetDevice>();
            if (!enbDev)
            {
                continue;
            }
            else
            {
                if (enbDev->HasCellId(cellId))
                {
                    found = true;
                    break;
                }
            }
        }
    }
    NS_ASSERT_MSG(found, " Unable to find eNB with CellId =" << cellId);
    m_enbRrcSapProvider = enbDev->GetRrc()->GetLteEnbRrcSapProvider();
    Ptr<LteEnbRrcProtocolIdeal> enbRrcProtocolIdeal =
        enbDev->GetRrc()->GetObject<LteEnbRrcProtocolIdeal>();
    enbRrcProtocolIdeal->SetUeRrcSapProvider(m_rnti, m_ueRrcSapProvider);
}

NS_OBJECT_ENSURE_REGISTERED(LteEnbRrcProtocolIdeal);

LteEnbRrcProtocolIdeal::LteEnbRrcProtocolIdeal()
    : m_enbRrcSapProvider(nullptr)
{
    NS_LOG_FUNCTION(this);
    m_enbRrcSapUser = new MemberLteEnbRrcSapUser<LteEnbRrcProtocolIdeal>(this);
}

LteEnbRrcProtocolIdeal::~LteEnbRrcProtocolIdeal()
{
    NS_LOG_FUNCTION(this);
}

void
LteEnbRrcProtocolIdeal::DoDispose()
{
    NS_LOG_FUNCTION(this);
    delete m_enbRrcSapUser;
}

TypeId
LteEnbRrcProtocolIdeal::GetTypeId()
{
    static TypeId tid = TypeId("ns3::LteEnbRrcProtocolIdeal")
                            .SetParent<Object>()
                            .SetGroupName("Lte")
                            .AddConstructor<LteEnbRrcProtocolIdeal>();
    return tid;
}

void
LteEnbRrcProtocolIdeal::SetLteEnbRrcSapProvider(LteEnbRrcSapProvider* p)
{
    m_enbRrcSapProvider = p;
}

LteEnbRrcSapUser*
LteEnbRrcProtocolIdeal::GetLteEnbRrcSapUser()
{
    return m_enbRrcSapUser;
}

void
LteEnbRrcProtocolIdeal::SetCellId(uint16_t cellId)
{
    m_cellId = cellId;
}

LteUeRrcSapProvider*
LteEnbRrcProtocolIdeal::GetUeRrcSapProvider(uint16_t rnti)
{
    auto it = m_enbRrcSapProviderMap.find(rnti);
    NS_ASSERT_MSG(it != m_enbRrcSapProviderMap.end(), "could not find RNTI = " << rnti);
    return it->second;
}

void
LteEnbRrcProtocolIdeal::SetUeRrcSapProvider(uint16_t rnti, LteUeRrcSapProvider* p)
{
    auto it = m_enbRrcSapProviderMap.find(rnti);
    // assign UE RRC only if the RNTI is found at eNB
    if (it != m_enbRrcSapProviderMap.end())
    {
        it->second = p;
    }
}

void
LteEnbRrcProtocolIdeal::DoSetupUe(uint16_t rnti, LteEnbRrcSapUser::SetupUeParameters params)
{
    NS_LOG_FUNCTION(this << rnti);

    // // walk list of all nodes to get the peer UE RRC SAP Provider
    // Ptr<LteUeRrc> ueRrc;
    // NodeList::Iterator listEnd = NodeList::End ();
    // bool found = false;
    // for (NodeList::Iterator i = NodeList::Begin (); (i != listEnd) && (found == false); i++)
    //   {
    //     Ptr<Node> node = *i;
    //     int nDevs = node->GetNDevices ();
    //     for (int j = 0; j < nDevs; j++)
    //       {
    //         Ptr<LteUeNetDevice> ueDev = node->GetDevice (j)->GetObject <LteUeNetDevice> ();
    //         if (!ueDev)
    //           {
    //             continue;
    //           }
    //         else
    //           {
    //             ueRrc = ueDev->GetRrc ();
    //             if ((ueRrc->GetRnti () == rnti) && (ueRrc->GetCellId () == m_cellId))
    //               {
    //                 found = true;
    //                 break;
    //               }
    //           }
    //       }
    //   }
    // NS_ASSERT_MSG (found , " Unable to find UE with RNTI=" << rnti << " cellId=" << m_cellId);
    // m_enbRrcSapProviderMap[rnti] = ueRrc->GetLteUeRrcSapProvider ();

    // just create empty entry, the UeRrcSapProvider will be set by the
    // ue upon connection request or connection reconfiguration
    // completed
    m_enbRrcSapProviderMap[rnti] = nullptr;
}

void
LteEnbRrcProtocolIdeal::DoRemoveUe(uint16_t rnti)
{
    NS_LOG_FUNCTION(this << rnti);
    m_enbRrcSapProviderMap.erase(rnti);
}

void
LteEnbRrcProtocolIdeal::DoSendSystemInformation(uint16_t cellId, LteRrcSap::SystemInformation msg)
{
    NS_LOG_FUNCTION(this << cellId);
    // walk list of all nodes to get UEs with this cellId
    Ptr<LteUeRrc> ueRrc;
    for (auto i = NodeList::Begin(); i != NodeList::End(); ++i)
    {
        Ptr<Node> node = *i;
        int nDevs = node->GetNDevices();
        for (int j = 0; j < nDevs; ++j)
        {
            Ptr<LteUeNetDevice> ueDev = node->GetDevice(j)->GetObject<LteUeNetDevice>();
            if (ueDev)
            {
                Ptr<LteUeRrc> ueRrc = ueDev->GetRrc();
                NS_LOG_LOGIC("considering UE IMSI " << ueDev->GetImsi() << " that has cellId "
                                                    << ueRrc->GetCellId());
                if (ueRrc->GetCellId() == cellId)
                {
                    NS_LOG_LOGIC("sending SI to IMSI " << ueDev->GetImsi());

                    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                                        &LteUeRrcSapProvider::RecvSystemInformation,
                                        ueRrc->GetLteUeRrcSapProvider(),
                                        msg);
                }
            }
        }
    }
}

void
LteEnbRrcProtocolIdeal::DoSendRrcConnectionSetup(uint16_t rnti, LteRrcSap::RrcConnectionSetup msg)
{
    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteUeRrcSapProvider::RecvRrcConnectionSetup,
                        GetUeRrcSapProvider(rnti),
                        msg);
}

void
LteEnbRrcProtocolIdeal::DoSendRrcConnectionReconfiguration(
    uint16_t rnti,
    LteRrcSap::RrcConnectionReconfiguration msg)
{
    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteUeRrcSapProvider::RecvRrcConnectionReconfiguration,
                        GetUeRrcSapProvider(rnti),
                        msg);
}

void
LteEnbRrcProtocolIdeal::DoSendRrcConnectionReestablishment(
    uint16_t rnti,
    LteRrcSap::RrcConnectionReestablishment msg)
{
    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteUeRrcSapProvider::RecvRrcConnectionReestablishment,
                        GetUeRrcSapProvider(rnti),
                        msg);
}

void
LteEnbRrcProtocolIdeal::DoSendRrcConnectionReestablishmentReject(
    uint16_t rnti,
    LteRrcSap::RrcConnectionReestablishmentReject msg)
{
    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteUeRrcSapProvider::RecvRrcConnectionReestablishmentReject,
                        GetUeRrcSapProvider(rnti),
                        msg);
}

void
LteEnbRrcProtocolIdeal::DoSendRrcConnectionRelease(uint16_t rnti,
                                                   LteRrcSap::RrcConnectionRelease msg)
{
    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteUeRrcSapProvider::RecvRrcConnectionRelease,
                        GetUeRrcSapProvider(rnti),
                        msg);
}

void
LteEnbRrcProtocolIdeal::DoSendRrcConnectionReject(uint16_t rnti, LteRrcSap::RrcConnectionReject msg)
{
    Simulator::Schedule(RRC_IDEAL_MSG_DELAY,
                        &LteUeRrcSapProvider::RecvRrcConnectionReject,
                        GetUeRrcSapProvider(rnti),
                        msg);
}

/*
 * The purpose of LteEnbRrcProtocolIdeal is to avoid encoding
 * messages. In order to do so, we need to have some form of encoding for
 * inter-node RRC messages like HandoverPreparationInfo and HandoverCommand. Doing so
 * directly is not practical (these messages includes a lot of
 * information elements, so encoding all of them would defeat the
 * purpose of LteEnbRrcProtocolIdeal. The workaround is to store the
 * actual message in a global map, so that then we can just encode the
 * key in a header and send that between eNBs over X2.
 *
 */

static std::map<uint32_t, LteRrcSap::HandoverPreparationInfo>
    g_handoverPreparationInfoMsgMap; ///< handover preparation info message map
static uint32_t g_handoverPreparationInfoMsgIdCounter =
    0; ///< handover preparation info message ID counter

/**
 * This header encodes the map key discussed above. We keep this
 * private since it should not be used outside this file.
 *
 */
class IdealHandoverPreparationInfoHeader : public Header
{
  public:
    /**
     * Get the message ID function
     *
     * @returns the message ID
     */
    uint32_t GetMsgId() const;
    /**
     * Set the message ID function
     *
     * @param id the message ID
     */
    void SetMsgId(uint32_t id);
    /**
     * @brief Get the type ID.
     * @return the object TypeId
     */
    static TypeId GetTypeId();
    TypeId GetInstanceTypeId() const override;
    void Print(std::ostream& os) const override;
    uint32_t GetSerializedSize() const override;
    void Serialize(Buffer::Iterator start) const override;
    uint32_t Deserialize(Buffer::Iterator start) override;

  private:
    uint32_t m_msgId; ///< message ID
};

uint32_t
IdealHandoverPreparationInfoHeader::GetMsgId() const
{
    return m_msgId;
}

void
IdealHandoverPreparationInfoHeader::SetMsgId(uint32_t id)
{
    m_msgId = id;
}

TypeId
IdealHandoverPreparationInfoHeader::GetTypeId()
{
    static TypeId tid = TypeId("ns3::IdealHandoverPreparationInfoHeader")
                            .SetParent<Header>()
                            .SetGroupName("Lte")
                            .AddConstructor<IdealHandoverPreparationInfoHeader>();
    return tid;
}

TypeId
IdealHandoverPreparationInfoHeader::GetInstanceTypeId() const
{
    return GetTypeId();
}

void
IdealHandoverPreparationInfoHeader::Print(std::ostream& os) const
{
    os << " msgId=" << m_msgId;
}

uint32_t
IdealHandoverPreparationInfoHeader::GetSerializedSize() const
{
    return 4;
}

void
IdealHandoverPreparationInfoHeader::Serialize(Buffer::Iterator start) const
{
    start.WriteU32(m_msgId);
}

uint32_t
IdealHandoverPreparationInfoHeader::Deserialize(Buffer::Iterator start)
{
    m_msgId = start.ReadU32();
    return GetSerializedSize();
}

Ptr<Packet>
LteEnbRrcProtocolIdeal::DoEncodeHandoverPreparationInformation(
    LteRrcSap::HandoverPreparationInfo msg)
{
    uint32_t msgId = ++g_handoverPreparationInfoMsgIdCounter;
    NS_ASSERT_MSG(g_handoverPreparationInfoMsgMap.find(msgId) ==
                      g_handoverPreparationInfoMsgMap.end(),
                  "msgId " << msgId << " already in use");
    NS_LOG_INFO(" encoding msgId = " << msgId);
    g_handoverPreparationInfoMsgMap.insert(
        std::pair<uint32_t, LteRrcSap::HandoverPreparationInfo>(msgId, msg));
    IdealHandoverPreparationInfoHeader h;
    h.SetMsgId(msgId);
    Ptr<Packet> p = Create<Packet>();
    p->AddHeader(h);
    return p;
}

LteRrcSap::HandoverPreparationInfo
LteEnbRrcProtocolIdeal::DoDecodeHandoverPreparationInformation(Ptr<Packet> p)
{
    IdealHandoverPreparationInfoHeader h;
    p->RemoveHeader(h);
    uint32_t msgId = h.GetMsgId();
    NS_LOG_INFO(" decoding msgId = " << msgId);
    auto it = g_handoverPreparationInfoMsgMap.find(msgId);
    NS_ASSERT_MSG(it != g_handoverPreparationInfoMsgMap.end(), "msgId " << msgId << " not found");
    LteRrcSap::HandoverPreparationInfo msg = it->second;
    g_handoverPreparationInfoMsgMap.erase(it);
    return msg;
}

static std::map<uint32_t, LteRrcSap::RrcConnectionReconfiguration>
    g_handoverCommandMsgMap;                       ///< handover command message map
static uint32_t g_handoverCommandMsgIdCounter = 0; ///< handover command message ID counter

/**
 * This header encodes the map key discussed above. We keep this
 * private since it should not be used outside this file.
 *
 */
class IdealHandoverCommandHeader : public Header
{
  public:
    /**
     * Get the message ID function
     *
     * @returns the message ID
     */
    uint32_t GetMsgId() const;
    /**
     * Set the message ID function
     *
     * @param id the message ID
     */
    void SetMsgId(uint32_t id);
    /**
     * @brief Get the type ID.
     * @return the object TypeId
     */
    static TypeId GetTypeId();
    TypeId GetInstanceTypeId() const override;
    void Print(std::ostream& os) const override;
    uint32_t GetSerializedSize() const override;
    void Serialize(Buffer::Iterator start) const override;
    uint32_t Deserialize(Buffer::Iterator start) override;

  private:
    uint32_t m_msgId; ///< message ID
};

uint32_t
IdealHandoverCommandHeader::GetMsgId() const
{
    return m_msgId;
}

void
IdealHandoverCommandHeader::SetMsgId(uint32_t id)
{
    m_msgId = id;
}

TypeId
IdealHandoverCommandHeader::GetTypeId()
{
    static TypeId tid = TypeId("ns3::IdealHandoverCommandHeader")
                            .SetParent<Header>()
                            .SetGroupName("Lte")
                            .AddConstructor<IdealHandoverCommandHeader>();
    return tid;
}

TypeId
IdealHandoverCommandHeader::GetInstanceTypeId() const
{
    return GetTypeId();
}

void
IdealHandoverCommandHeader::Print(std::ostream& os) const
{
    os << " msgId=" << m_msgId;
}

uint32_t
IdealHandoverCommandHeader::GetSerializedSize() const
{
    return 4;
}

void
IdealHandoverCommandHeader::Serialize(Buffer::Iterator start) const
{
    start.WriteU32(m_msgId);
}

uint32_t
IdealHandoverCommandHeader::Deserialize(Buffer::Iterator start)
{
    m_msgId = start.ReadU32();
    return GetSerializedSize();
}

Ptr<Packet>
LteEnbRrcProtocolIdeal::DoEncodeHandoverCommand(LteRrcSap::RrcConnectionReconfiguration msg)
{
    uint32_t msgId = ++g_handoverCommandMsgIdCounter;
    NS_ASSERT_MSG(g_handoverCommandMsgMap.find(msgId) == g_handoverCommandMsgMap.end(),
                  "msgId " << msgId << " already in use");
    NS_LOG_INFO(" encoding msgId = " << msgId);
    g_handoverCommandMsgMap.insert(
        std::pair<uint32_t, LteRrcSap::RrcConnectionReconfiguration>(msgId, msg));
    IdealHandoverCommandHeader h;
    h.SetMsgId(msgId);
    Ptr<Packet> p = Create<Packet>();
    p->AddHeader(h);
    return p;
}

LteRrcSap::RrcConnectionReconfiguration
LteEnbRrcProtocolIdeal::DoDecodeHandoverCommand(Ptr<Packet> p)
{
    IdealHandoverCommandHeader h;
    p->RemoveHeader(h);
    uint32_t msgId = h.GetMsgId();
    NS_LOG_INFO(" decoding msgId = " << msgId);
    auto it = g_handoverCommandMsgMap.find(msgId);
    NS_ASSERT_MSG(it != g_handoverCommandMsgMap.end(), "msgId " << msgId << " not found");
    LteRrcSap::RrcConnectionReconfiguration msg = it->second;
    g_handoverCommandMsgMap.erase(it);
    return msg;
}

} // namespace ns3
