#include "ns3/core-module.h"
#include "ns3/network-module.h"
#include "ns3/nr-helper.h"
#include "ns3/nr-module.h"
#include "ns3/nr-point-to-point-epc-helper.h"
#include "ns3/ipv4-global-routing-helper.h"
#include "ns3/internet-module.h"
#include "ns3/eps-bearer-tag.h"
#include "ns3/antenna-module.h"

using namespace ns3;

static bool g_rxPdcpCallbackCalled = false;
static bool g_rxRxRlcPDUCallbackCalled = false;
static void SendPacket (Ptr<NetDevice> device, Address& addr, uint32_t packetSize)
{
  Ptr<Packet> pkt = Create<Packet> (packetSize);
  Ipv4Header ipv4Header;
  ipv4Header.SetProtocol (UdpL4Protocol::PROT_NUMBER);
  pkt->AddHeader (ipv4Header);
  EpsBearerTag tag (1, 1);
  pkt->AddPacketTag (tag);
  device->Send (pkt, addr, Ipv4L3Protocol::PROT_NUMBER);
}

void
RxPdcpPDU (std::string path, uint16_t rnti, uint8_t lcid, uint32_t bytes, uint64_t pdcpDelay)
{
  std::cout << "\n Packet PDCP delay:" << pdcpDelay << "\n";
  g_rxPdcpCallbackCalled = true;
}

void
RxRlcPDU (std::string path, uint16_t rnti, uint8_t lcid, uint32_t bytes, uint64_t rlcDelay)
{
  std::cout << "\n\n Data received at RLC layer at:" << Simulator::Now () << std::endl;
  std::cout << "\n rnti:" << rnti << std::endl;
  std::cout << "\n lcid:" << (unsigned)lcid << std::endl;
  std::cout << "\n bytes :" << bytes << std::endl;
  std::cout << "\n delay :" << rlcDelay << std::endl;
  g_rxRxRlcPDUCallbackCalled = true;
}

/**
 * Function that connects PDCP and RLC traces to the corresponding trace sources.
 */
void
ConnectPdcpRlcTraces ()
{
  Config::Connect ("/NodeList/*/DeviceList/*/LteUeRrc/DataRadioBearerMap/1/LtePdcp/RxPDU",
                   MakeCallback (&RxPdcpPDU));

  Config::Connect ("/NodeList/*/DeviceList/*/LteUeRrc/DataRadioBearerMap/1/LteRlc/RxPDU",
                   MakeCallback (&RxRlcPDU));
}

/**
 * Function that connects UL PDCP and RLC traces to the corresponding trace sources.
 */
void
ConnectUlPdcpRlcTraces ()
{
  Config::Connect ("/NodeList/*/DeviceList/*/LteEnbRrc/UeMap/*/DataRadioBearerMap/*/LtePdcp/RxPDU",
                   MakeCallback (&RxPdcpPDU));

  Config::Connect ("/NodeList/*/DeviceList/*/LteEnbRrc/UeMap/*/DataRadioBearerMap/*/LteRlc/RxPDU",
                   MakeCallback (&RxRlcPDU));
}


int
main (int argc, char *argv[])
{
  uint16_t numerologyBwp1 = 0;
  uint32_t udpPacketSize = 1000;
  double centralFrequencyBand1 = 28e9;
  double bandwidthBand1 = 400e6;
  uint16_t gNbNum = 1;
  uint16_t ueNumPergNb = 2;
  //bool enableUl = true;

  Time sendPacketTime = Seconds (0.4);

  int64_t randomStream = 1;
  //Create the scenario
GridScenarioHelper gridScenario;
gridScenario.SetRows(1);                        // Single row since we are arranging elements horizontally.
gridScenario.SetColumns(2);                     // Two columns to place gNB and UE side by side.
gridScenario.SetHorizontalBsDistance(10.0);     // Distance between gNB and UE set to 10 meters.
gridScenario.SetBsHeight(1.5);                  // Height of the gNB.
gridScenario.SetUtHeight(1.5);                  // Height of the UE, same as gNB for alignment.

gridScenario.SetSectorization(GridScenarioHelper::SINGLE);
gridScenario.SetBsNumber(gNbNum);               // Assuming gNbNum is set to 1.
gridScenario.SetUtNumber(ueNumPergNb * gNbNum); // Assuming ueNumPergNb is set to match the number of UEs you want per gNb.
gridScenario.SetScenarioHeight(3);
gridScenario.SetScenarioLength(3);
randomStream += gridScenario.AssignStreams(randomStream);
gridScenario.CreateScenario();



  Ptr<NrPointToPointEpcHelper> epcHelper = CreateObject<NrPointToPointEpcHelper> ();
  Ptr<IdealBeamformingHelper> idealBeamformingHelper = CreateObject<IdealBeamformingHelper>();
  Ptr<NrHelper> nrHelper = CreateObject<NrHelper> ();

  nrHelper->SetBeamformingHelper (idealBeamformingHelper);
  nrHelper->SetEpcHelper (epcHelper);

  // Create one operational band containing one CC with one bandwidth part
  BandwidthPartInfoPtrVector allBwps;
  CcBwpCreator ccBwpCreator;
  const uint8_t numCcPerBand = 1;

  // Create the configuration for the CcBwpHelper
  CcBwpCreator::SimpleOperationBandConf bandConf1 (centralFrequencyBand1, bandwidthBand1,
                                                   numCcPerBand, BandwidthPartInfo::UMi_StreetCanyon_LoS);

  // By using the configuration created, it is time to make the operation band
  OperationBandInfo band1 = ccBwpCreator.CreateOperationBandContiguousCc (bandConf1);

  Config::SetDefault ("ns3::ThreeGppChannelModel::UpdatePeriod",TimeValue (MilliSeconds (0)));
  nrHelper->SetSchedulerAttribute ("FixedMcsDl", BooleanValue (true));
  nrHelper->SetSchedulerAttribute ("StartingMcsDl", UintegerValue (28));
  nrHelper->SetChannelConditionModelAttribute ("UpdatePeriod", TimeValue (MilliSeconds (0)));
  nrHelper->SetPathlossAttribute ("ShadowingEnabled", BooleanValue (false));

  nrHelper->InitializeOperationBand (&band1);
  allBwps = CcBwpCreator::GetAllBwps ({band1});

  // Beamforming method
  idealBeamformingHelper->SetAttribute ("BeamformingMethod", TypeIdValue (DirectPathBeamforming::GetTypeId ()));

  // Antennas for all the UEs
  nrHelper->SetUeAntennaAttribute ("NumRows", UintegerValue (2));
  nrHelper->SetUeAntennaAttribute ("NumColumns", UintegerValue (4));
  nrHelper->SetUeAntennaAttribute ("AntennaElement", PointerValue (CreateObject<IsotropicAntennaModel> ()));

  // Antennas for all the gNbs
  nrHelper->SetGnbAntennaAttribute ("NumRows", UintegerValue (4));
  nrHelper->SetGnbAntennaAttribute ("NumColumns", UintegerValue (8));
  nrHelper->SetGnbAntennaAttribute ("AntennaElement", PointerValue (CreateObject<IsotropicAntennaModel> ()));

  //Install and get the pointers to the NetDevices
  NetDeviceContainer enbNetDev = nrHelper->InstallGnbDevice (gridScenario.GetBaseStations (), allBwps);
  NetDeviceContainer ueNetDev = nrHelper->InstallUeDevice (gridScenario.GetUserTerminals (), allBwps);

  randomStream += nrHelper->AssignStreams (enbNetDev, randomStream);
  randomStream += nrHelper->AssignStreams (ueNetDev, randomStream);

  // Set the attribute of the netdevice (enbNetDev.Get (0)) and bandwidth part (0)
  nrHelper->GetGnbPhy (enbNetDev.Get (0), 0)->SetAttribute ("Numerology", UintegerValue (numerologyBwp1));

  for (auto it = enbNetDev.Begin (); it != enbNetDev.End (); ++it)
    {
      DynamicCast<NrGnbNetDevice> (*it)->UpdateConfig ();
    }

  for (auto it = ueNetDev.Begin (); it != ueNetDev.End (); ++it)
    {
      DynamicCast<NrUeNetDevice> (*it)->UpdateConfig ();
    }

 InternetStackHelper internet;
  internet.Install (gridScenario.GetUserTerminals ());
  Ipv4InterfaceContainer ueIpIface = epcHelper->AssignUeIpv4Address (NetDeviceContainer (ueNetDev));

  // Attach UEs
  nrHelper->AttachToClosestEnb (ueNetDev, enbNetDev);

  // Uplink Transmission: UE1 -> gNB 
  Simulator::Schedule (sendPacketTime, &SendPacket, ueNetDev.Get (0), enbNetDev.Get (0)->GetAddress (), udpPacketSize);

  // Downlink Transmission: gNB -> UE2
  Simulator::Schedule (sendPacketTime, &SendPacket, enbNetDev.Get (0), ueNetDev.Get (1)->GetAddress (), udpPacketSize); 

  // Separate traces for uplink and downlink
  std::cout << "\n Sending data in uplink (UE1 -> gNB)." << std::endl;
  Simulator::Schedule (Seconds (0.2), &ConnectUlPdcpRlcTraces); 

  std::cout << "\n Sending data in downlink (gNB -> UE2)." << std::endl;
  Simulator::Schedule (Seconds (0.2), &ConnectPdcpRlcTraces); // Modify function name if needed

  nrHelper->EnableTraces ();

  Simulator::Stop (Seconds (1));
  Simulator::Run ();
  Simulator::Destroy ();

  if (g_rxPdcpCallbackCalled && g_rxRxRlcPDUCallbackCalled)
    {
      return EXIT_SUCCESS;
    }
  else
    {
      return EXIT_FAILURE;
    }

}