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

#include "epc-ue-nas.h"

#include "lte-as-sap.h"

#include "ns3/epc-helper.h"
#include "ns3/fatal-error.h"
#include "ns3/log.h"
#include "ns3/simulator.h"

namespace ns3
{

NS_LOG_COMPONENT_DEFINE("EpcUeNas");

NS_OBJECT_ENSURE_REGISTERED(EpcUeNas);

EpcUeNas::EpcUeNas()
    : m_state(OFF),
      m_csgId(0),
      m_asSapProvider(nullptr),
      m_bidCounter(0)
{
    NS_LOG_FUNCTION(this);
    m_asSapUser = new MemberLteAsSapUser<EpcUeNas>(this);
}

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

void
EpcUeNas::DoDispose()
{
    NS_LOG_FUNCTION(this);
    delete m_asSapUser;
}

TypeId
EpcUeNas::GetTypeId()
{
    static TypeId tid =
        TypeId("ns3::EpcUeNas")
            .SetParent<Object>()
            .SetGroupName("Lte")
            .AddConstructor<EpcUeNas>()
            .AddTraceSource("StateTransition",
                            "fired upon every UE NAS state transition",
                            MakeTraceSourceAccessor(&EpcUeNas::m_stateTransitionCallback),
                            "ns3::EpcUeNas::StateTracedCallback");
    return tid;
}

void
EpcUeNas::SetDevice(Ptr<NetDevice> dev)
{
    NS_LOG_FUNCTION(this << dev);
    m_device = dev;
}

void
EpcUeNas::SetImsi(uint64_t imsi)
{
    NS_LOG_FUNCTION(this << imsi);
    m_imsi = imsi;
}

void
EpcUeNas::SetCsgId(uint32_t csgId)
{
    NS_LOG_FUNCTION(this << csgId);
    m_csgId = csgId;
    m_asSapProvider->SetCsgWhiteList(csgId);
}

uint32_t
EpcUeNas::GetCsgId() const
{
    NS_LOG_FUNCTION(this);
    return m_csgId;
}

void
EpcUeNas::SetAsSapProvider(LteAsSapProvider* s)
{
    NS_LOG_FUNCTION(this << s);
    m_asSapProvider = s;
}

LteAsSapUser*
EpcUeNas::GetAsSapUser()
{
    NS_LOG_FUNCTION(this);
    return m_asSapUser;
}

void
EpcUeNas::SetForwardUpCallback(Callback<void, Ptr<Packet>> cb)
{
    NS_LOG_FUNCTION(this);
    m_forwardUpCallback = cb;
}

void
EpcUeNas::StartCellSelection(uint32_t dlEarfcn)
{
    NS_LOG_FUNCTION(this << dlEarfcn);
    m_asSapProvider->StartCellSelection(dlEarfcn);
}

void
EpcUeNas::Connect()
{
    NS_LOG_FUNCTION(this);

    // tell RRC to go into connected mode
    m_asSapProvider->Connect();
}

void
EpcUeNas::Connect(uint16_t cellId, uint32_t dlEarfcn)
{
    NS_LOG_FUNCTION(this << cellId << dlEarfcn);

    // force the UE RRC to be camped on a specific eNB
    m_asSapProvider->ForceCampedOnEnb(cellId, dlEarfcn);

    // tell RRC to go into connected mode
    m_asSapProvider->Connect();
}

void
EpcUeNas::Disconnect()
{
    NS_LOG_FUNCTION(this);
    SwitchToState(OFF);
    m_asSapProvider->Disconnect();
}

void
EpcUeNas::ActivateEpsBearer(EpsBearer bearer, Ptr<EpcTft> tft)
{
    NS_LOG_FUNCTION(this);
    switch (m_state)
    {
    case ACTIVE:
        NS_FATAL_ERROR("the necessary NAS signaling to activate a bearer after the initial context "
                       "has already been setup is not implemented");
        break;

    default:
        BearerToBeActivated btba;
        btba.bearer = bearer;
        btba.tft = tft;
        m_bearersToBeActivatedList.push_back(btba);
        m_bearersToBeActivatedListForReconnection.push_back(btba);
        break;
    }
}

bool
EpcUeNas::Send(Ptr<Packet> packet, uint16_t protocolNumber)
{
    NS_LOG_FUNCTION(this << packet << protocolNumber);

    switch (m_state)
    {
    case ACTIVE: {
        uint32_t id = m_tftClassifier.Classify(packet, EpcTft::UPLINK, protocolNumber);
        NS_ASSERT((id & 0xFFFFFF00) == 0);
        auto bid = (uint8_t)(id & 0x000000FF);
        if (bid == 0)
        {
            return false;
        }
        else
        {
            m_asSapProvider->SendData(packet, bid);
            return true;
        }
    }
    break;

    default:
        NS_LOG_WARN(this << " NAS OFF, discarding packet");
        return false;
    }
}

void
EpcUeNas::DoNotifyConnectionSuccessful()
{
    NS_LOG_FUNCTION(this);

    SwitchToState(ACTIVE); // will eventually activate dedicated bearers
}

void
EpcUeNas::DoNotifyConnectionFailed()
{
    NS_LOG_FUNCTION(this);

    // immediately retry the connection
    Simulator::ScheduleNow(&LteAsSapProvider::Connect, m_asSapProvider);
}

void
EpcUeNas::DoRecvData(Ptr<Packet> packet)
{
    NS_LOG_FUNCTION(this << packet);
    m_forwardUpCallback(packet);
}

void
EpcUeNas::DoNotifyConnectionReleased()
{
    NS_LOG_FUNCTION(this);
    // remove tfts
    while (m_bidCounter > 0)
    {
        m_tftClassifier.Delete(m_bidCounter);
        m_bidCounter--;
    }
    // restore the bearer list to be activated for the next RRC connection
    m_bearersToBeActivatedList = m_bearersToBeActivatedListForReconnection;

    Disconnect();
}

void
EpcUeNas::DoActivateEpsBearer(EpsBearer bearer, Ptr<EpcTft> tft)
{
    NS_LOG_FUNCTION(this);
    NS_ASSERT_MSG(m_bidCounter < 11, "cannot have more than 11 EPS bearers");
    uint8_t bid = ++m_bidCounter;
    m_tftClassifier.Add(tft, bid);
}

EpcUeNas::State
EpcUeNas::GetState() const
{
    NS_LOG_FUNCTION(this);
    return m_state;
}

void
EpcUeNas::SwitchToState(State newState)
{
    NS_LOG_FUNCTION(this << newState);
    State oldState = m_state;
    m_state = newState;
    NS_LOG_INFO("IMSI " << m_imsi << " NAS " << oldState << " --> " << newState);
    m_stateTransitionCallback(oldState, newState);

    // actions to be done when entering a new state:
    switch (m_state)
    {
    case ACTIVE:
        for (auto it = m_bearersToBeActivatedList.begin(); it != m_bearersToBeActivatedList.end();
             m_bearersToBeActivatedList.erase(it++))
        {
            DoActivateEpsBearer(it->bearer, it->tft);
        }
        break;

    default:
        break;
    }
}

std::ostream&
operator<<(std::ostream& os, EpcUeNas::State state)
{
    switch (state)
    {
    case EpcUeNas::State::OFF:
        return os << "OFF";
    case EpcUeNas::State::ATTACHING:
        return os << "ATTACHING";
    case EpcUeNas::State::IDLE_REGISTERED:
        return os << "IDLE_REGISTERED";
    case EpcUeNas::State::CONNECTING_TO_EPC:
        return os << "CONNECTING_TO_EPC";
    case EpcUeNas::State::ACTIVE:
        return os << "ACTIVE";
    case EpcUeNas::State::NUM_STATES:
        return os << "NUM_STATES";
    };
    return os << "UNKNOWN(" << static_cast<uint32_t>(state) << ")";
}

} // namespace ns3
