/*
 * Copyright (c) 2020 Universita' degli Studi di Napoli Federico II
 *
 * SPDX-License-Identifier: GPL-2.0-only
 *
 * Author: Stefano Avallone <stavallo@unina.it>
 */

#include "wifi-acknowledgment.h"

#include "wifi-utils.h"

#include "ns3/mac48-address.h"

namespace ns3
{

/*
 * WifiAcknowledgment
 */

WifiAcknowledgment::WifiAcknowledgment(Method m)
    : method(m)
{
}

WifiAcknowledgment::~WifiAcknowledgment()
{
}

WifiMacHeader::QosAckPolicy
WifiAcknowledgment::GetQosAckPolicy(Mac48Address receiver, uint8_t tid) const
{
    auto it = m_ackPolicy.find({receiver, tid});
    NS_ASSERT(it != m_ackPolicy.end());
    return it->second;
}

void
WifiAcknowledgment::SetQosAckPolicy(Mac48Address receiver,
                                    uint8_t tid,
                                    WifiMacHeader::QosAckPolicy ackPolicy)
{
    NS_ABORT_MSG_IF(!CheckQosAckPolicy(receiver, tid, ackPolicy), "QoS Ack policy not admitted");
    m_ackPolicy[{receiver, tid}] = ackPolicy;
}

/*
 * WifiNoAck
 */

WifiNoAck::WifiNoAck()
    : WifiAcknowledgment(NONE)
{
    acknowledgmentTime = Seconds(0);
}

std::unique_ptr<WifiAcknowledgment>
WifiNoAck::Copy() const
{
    return std::make_unique<WifiNoAck>(*this);
}

bool
WifiNoAck::CheckQosAckPolicy(Mac48Address receiver,
                             uint8_t tid,
                             WifiMacHeader::QosAckPolicy ackPolicy) const
{
    return ackPolicy == WifiMacHeader::NO_ACK || ackPolicy == WifiMacHeader::BLOCK_ACK;
}

void
WifiNoAck::Print(std::ostream& os) const
{
    os << "NONE";
}

/*
 * WifiNormalAck
 */

WifiNormalAck::WifiNormalAck()
    : WifiAcknowledgment(NORMAL_ACK)
{
}

std::unique_ptr<WifiAcknowledgment>
WifiNormalAck::Copy() const
{
    return std::make_unique<WifiNormalAck>(*this);
}

bool
WifiNormalAck::CheckQosAckPolicy(Mac48Address receiver,
                                 uint8_t tid,
                                 WifiMacHeader::QosAckPolicy ackPolicy) const
{
    return ackPolicy == WifiMacHeader::NORMAL_ACK;
}

void
WifiNormalAck::Print(std::ostream& os) const
{
    os << "NORMAL_ACK";
}

/*
 * WifiBlockAck
 */

WifiBlockAck::WifiBlockAck()
    : WifiAcknowledgment(BLOCK_ACK)
{
}

std::unique_ptr<WifiAcknowledgment>
WifiBlockAck::Copy() const
{
    return std::make_unique<WifiBlockAck>(*this);
}

bool
WifiBlockAck::CheckQosAckPolicy(Mac48Address receiver,
                                uint8_t tid,
                                WifiMacHeader::QosAckPolicy ackPolicy) const
{
    return ackPolicy == WifiMacHeader::NORMAL_ACK;
}

void
WifiBlockAck::Print(std::ostream& os) const
{
    os << "BLOCK_ACK";
}

/*
 * WifiBarBlockAck
 */

WifiBarBlockAck::WifiBarBlockAck()
    : WifiAcknowledgment(BAR_BLOCK_ACK)
{
}

std::unique_ptr<WifiAcknowledgment>
WifiBarBlockAck::Copy() const
{
    return std::make_unique<WifiBarBlockAck>(*this);
}

bool
WifiBarBlockAck::CheckQosAckPolicy(Mac48Address receiver,
                                   uint8_t tid,
                                   WifiMacHeader::QosAckPolicy ackPolicy) const
{
    return ackPolicy == WifiMacHeader::BLOCK_ACK;
}

void
WifiBarBlockAck::Print(std::ostream& os) const
{
    os << "BAR_BLOCK_ACK";
}

/*
 * WifiDlMuBarBaSequence
 */

WifiDlMuBarBaSequence::WifiDlMuBarBaSequence()
    : WifiAcknowledgment(DL_MU_BAR_BA_SEQUENCE)
{
}

std::unique_ptr<WifiAcknowledgment>
WifiDlMuBarBaSequence::Copy() const
{
    return std::make_unique<WifiDlMuBarBaSequence>(*this);
}

bool
WifiDlMuBarBaSequence::CheckQosAckPolicy(Mac48Address receiver,
                                         uint8_t tid,
                                         WifiMacHeader::QosAckPolicy ackPolicy) const
{
    if (ackPolicy == WifiMacHeader::NORMAL_ACK)
    {
        // The given receiver must be the only one to send an immediate reply
        if (stationsReplyingWithNormalAck.size() == 1 &&
            stationsReplyingWithNormalAck.begin()->first == receiver)
        {
            return true;
        }

        if (stationsReplyingWithBlockAck.size() == 1 &&
            stationsReplyingWithBlockAck.begin()->first == receiver)
        {
            return true;
        }

        return false;
    }

    return ackPolicy == WifiMacHeader::BLOCK_ACK;
}

void
WifiDlMuBarBaSequence::Print(std::ostream& os) const
{
    os << "DL_MU_BAR_BA_SEQUENCE [";
    for (const auto& sta : stationsReplyingWithNormalAck)
    {
        os << " (ACK) " << sta.first;
    }
    for (const auto& sta : stationsReplyingWithBlockAck)
    {
        os << " (BA) " << sta.first;
    }
    for (const auto& sta : stationsSendBlockAckReqTo)
    {
        os << " (BAR+BA) " << sta.first;
    }
    os << "]";
}

/*
 * WifiDlMuTfMuBar
 */

WifiDlMuTfMuBar::WifiDlMuTfMuBar()
    : WifiAcknowledgment(DL_MU_TF_MU_BAR),
      ulLength(0)
{
}

std::unique_ptr<WifiAcknowledgment>
WifiDlMuTfMuBar::Copy() const
{
    return std::make_unique<WifiDlMuTfMuBar>(*this);
}

bool
WifiDlMuTfMuBar::CheckQosAckPolicy(Mac48Address receiver,
                                   uint8_t tid,
                                   WifiMacHeader::QosAckPolicy ackPolicy) const
{
    // the only admitted ack policy is Block Ack because stations need to wait for a MU-BAR
    return ackPolicy == WifiMacHeader::BLOCK_ACK;
}

void
WifiDlMuTfMuBar::Print(std::ostream& os) const
{
    os << "DL_MU_TF_MU_BAR [";
    for (const auto& sta : stationsReplyingWithBlockAck)
    {
        os << " (BA) " << sta.first;
    }
    os << "]";
}

/*
 * WifiDlMuAggregateTf
 */

WifiDlMuAggregateTf::WifiDlMuAggregateTf()
    : WifiAcknowledgment(DL_MU_AGGREGATE_TF),
      ulLength(0)
{
}

std::unique_ptr<WifiAcknowledgment>
WifiDlMuAggregateTf::Copy() const
{
    return std::make_unique<WifiDlMuAggregateTf>(*this);
}

bool
WifiDlMuAggregateTf::CheckQosAckPolicy(Mac48Address receiver,
                                       uint8_t tid,
                                       WifiMacHeader::QosAckPolicy ackPolicy) const
{
    // the only admitted ack policy is No explicit acknowledgment or TB PPDU Ack policy
    return ackPolicy == WifiMacHeader::NO_EXPLICIT_ACK;
}

void
WifiDlMuAggregateTf::Print(std::ostream& os) const
{
    os << "DL_MU_AGGREGATE_TF [";
    for (const auto& sta : stationsReplyingWithBlockAck)
    {
        os << " (BA) " << sta.first;
    }
    os << "]";
}

/*
 * WifiUlMuMultiStaBa
 */

WifiUlMuMultiStaBa::WifiUlMuMultiStaBa()
    : WifiAcknowledgment(UL_MU_MULTI_STA_BA),
      baType(BlockAckType::MULTI_STA)
{
}

std::unique_ptr<WifiAcknowledgment>
WifiUlMuMultiStaBa::Copy() const
{
    return std::make_unique<WifiUlMuMultiStaBa>(*this);
}

bool
WifiUlMuMultiStaBa::CheckQosAckPolicy(Mac48Address receiver,
                                      uint8_t tid,
                                      WifiMacHeader::QosAckPolicy ackPolicy) const
{
    // a Basic Trigger Frame has no QoS ack policy
    return true;
}

void
WifiUlMuMultiStaBa::Print(std::ostream& os) const
{
    os << "UL_MU_MULTI_STA_BA [";
    for (const auto& sta : stationsReceivingMultiStaBa)
    {
        os << "(" << sta.first.first << "," << +sta.first.second << ") ";
    }
    os << "]";
}

/*
 * WifiAckAfterTbPpdu
 */

WifiAckAfterTbPpdu::WifiAckAfterTbPpdu()
    : WifiAcknowledgment(ACK_AFTER_TB_PPDU)
{
}

std::unique_ptr<WifiAcknowledgment>
WifiAckAfterTbPpdu::Copy() const
{
    return std::make_unique<WifiAckAfterTbPpdu>(*this);
}

bool
WifiAckAfterTbPpdu::CheckQosAckPolicy(Mac48Address receiver,
                                      uint8_t tid,
                                      WifiMacHeader::QosAckPolicy ackPolicy) const
{
    return ackPolicy == WifiMacHeader::NORMAL_ACK;
}

void
WifiAckAfterTbPpdu::Print(std::ostream& os) const
{
    os << "ACK_AFTER_TB_PPDU";
}

std::ostream&
operator<<(std::ostream& os, const WifiAcknowledgment* acknowledgment)
{
    acknowledgment->Print(os);
    return os;
}

} // namespace ns3
